mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-08 13:33:42 +08:00
firestar5683
22a8575199
Keep official openpilot ancestry and dependency licenses. Store resources directly in Git, remove LFS setup requirements, and leave the downloadable big model outside the source tree.
149 lines
6.7 KiB
Python
149 lines
6.7 KiB
Python
"""Compile perspective warps for NV12 frames."""
|
|
import argparse, pickle
|
|
from typing import NamedTuple
|
|
import numpy as np
|
|
from tinygrad import Tensor, Device, Context, TinyJit
|
|
from tinygrad.engine.realize import lower_and_compile
|
|
from examples.openpilot.helpers import allocate_inputs, benchmark, make_retargetable
|
|
|
|
|
|
class NV12Frame(NamedTuple):
|
|
width: int
|
|
height: int
|
|
stride: int
|
|
y_height: int
|
|
uv_height: int
|
|
size: int
|
|
|
|
|
|
def parse_frame(value): return NV12Frame(*map(int, value.split(',')))
|
|
|
|
def parse_size(value): return tuple(map(int, value.lower().split('x')))
|
|
|
|
|
|
def warp_perspective_tinygrad(src_flat, M_inv, dst_shape, src_shape, stride_pad, border_fill_val=None):
|
|
w_dst, h_dst = dst_shape
|
|
h_src, w_src = src_shape
|
|
|
|
x = Tensor.arange(w_dst).reshape(1, w_dst).expand(h_dst, w_dst).reshape(-1)
|
|
y = Tensor.arange(h_dst).reshape(h_dst, 1).expand(h_dst, w_dst).reshape(-1)
|
|
|
|
# inline 3x3 matmul as elementwise to avoid reduce op (enables fusion with gather)
|
|
src_x = M_inv[0, 0] * x + M_inv[0, 1] * y + M_inv[0, 2]
|
|
src_y = M_inv[1, 0] * x + M_inv[1, 1] * y + M_inv[1, 2]
|
|
src_w = M_inv[2, 0] * x + M_inv[2, 1] * y + M_inv[2, 2]
|
|
|
|
src_x = src_x / src_w
|
|
src_y = src_y / src_w
|
|
|
|
x_round = Tensor.round(src_x)
|
|
y_round = Tensor.round(src_y)
|
|
x_nn_clipped = x_round.clip(0, w_src - 1).cast('int')
|
|
y_nn_clipped = y_round.clip(0, h_src - 1).cast('int')
|
|
idx = y_nn_clipped * (w_src + stride_pad) + x_nn_clipped
|
|
sampled = src_flat[idx]
|
|
|
|
if border_fill_val is None:
|
|
return sampled
|
|
|
|
in_bounds = ((x_round >= 0) & (x_round <= w_src - 1) &
|
|
(y_round >= 0) & (y_round <= h_src - 1)).cast(sampled.dtype)
|
|
return sampled * in_bounds + Tensor(border_fill_val, dtype=sampled.dtype) * (1 - in_bounds)
|
|
|
|
|
|
def frames_to_tensor(frames):
|
|
H = (frames.shape[0] * 2) // 3
|
|
W = frames.shape[1]
|
|
in_img1 = Tensor.cat(frames[0:H:2, 0::2],
|
|
frames[1:H:2, 0::2],
|
|
frames[0:H:2, 1::2],
|
|
frames[1:H:2, 1::2],
|
|
frames[H:H+H//4].reshape((H//2, W//2)),
|
|
frames[H+H//4:H+H//2].reshape((H//2, W//2)), dim=0).reshape((6, H//2, W//2))
|
|
return in_img1
|
|
|
|
|
|
def make_frame_prepare(nv12: NV12Frame, model_w, model_h):
|
|
cam_w, cam_h, stride, y_height, uv_height, _ = nv12
|
|
uv_offset = stride * y_height
|
|
stride_pad = stride - cam_w
|
|
|
|
def frame_prepare_tinygrad(input_frame, M_inv):
|
|
M_inv = M_inv.to(Device.DEFAULT).realize()
|
|
# UV_SCALE @ M_inv @ UV_SCALE_INV simplifies to elementwise scaling
|
|
M_inv_uv = M_inv * Tensor([[1.0, 1.0, 0.5], [1.0, 1.0, 0.5], [2.0, 2.0, 1.0]], device=Device.DEFAULT)
|
|
# deinterleave NV12 UV plane (UVUV... -> separate U, V)
|
|
uv = input_frame[uv_offset:uv_offset + uv_height * stride].reshape(uv_height, stride)
|
|
with Context(SPLIT_REDUCEOP=0):
|
|
y = warp_perspective_tinygrad(input_frame[:cam_h*stride],
|
|
M_inv, (model_w, model_h),
|
|
(cam_h, cam_w), stride_pad).realize()
|
|
u = warp_perspective_tinygrad(uv[:cam_h//2, :cam_w:2].flatten(),
|
|
M_inv_uv, (model_w//2, model_h//2),
|
|
(cam_h//2, cam_w//2), 0).realize()
|
|
v = warp_perspective_tinygrad(uv[:cam_h//2, 1:cam_w:2].flatten(),
|
|
M_inv_uv, (model_w//2, model_h//2),
|
|
(cam_h//2, cam_w//2), 0).realize()
|
|
yuv = y.cat(u).cat(v).reshape((model_h * 3 // 2, model_w))
|
|
tensor = frames_to_tensor(yuv)
|
|
return tensor
|
|
return frame_prepare_tinygrad
|
|
|
|
|
|
def make_luma_warp(nv12:NV12Frame, width, height, border_fill=None):
|
|
def warp(input_frame, M_inv):
|
|
M_inv = M_inv.to(Device.DEFAULT).realize()
|
|
return warp_perspective_tinygrad(input_frame[:nv12.height*nv12.stride], M_inv,
|
|
(width, height), (nv12.height, nv12.width), nv12.stride-nv12.width,
|
|
border_fill_val=border_fill).reshape(-1, height*width)
|
|
return warp
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument('--frame', type=parse_frame, required=True, help='width,height,stride,y_height,uv_height,buffer_size')
|
|
parser.add_argument('--warp-to', type=parse_size, required=True)
|
|
parser.add_argument('--layout', choices=['luma', 'yuv420'], default='luma')
|
|
parser.add_argument('--border-fill', type=int, help='luma outside the frame; omit to clamp coordinates')
|
|
parser.add_argument('--frames', type=int, default=1, help='number of frames to warp together')
|
|
parser.add_argument('--transform-device', help='device holding the transforms; defaults to DEV')
|
|
parser.add_argument('--output', required=True)
|
|
parser.add_argument('--benchmark-runs', type=int, default=20)
|
|
parser.add_argument('--retargetable', action='store_true', help='output retargetable jit (requires recompilation when loading)')
|
|
args = parser.parse_args()
|
|
|
|
warp = make_luma_warp(args.frame, *args.warp_to, args.border_fill) if args.layout == 'luma' else make_frame_prepare(args.frame, *args.warp_to)
|
|
prefix = () if args.frames == 1 else (args.frames,)
|
|
specs = {'input_frame': (prefix + (args.frame.size,), np.dtype(np.uint8).str, Device.DEFAULT),
|
|
'M_inv': (prefix + (3, 3), np.dtype(np.float32).str, args.transform_device or Device.DEFAULT)}
|
|
|
|
def make_inputs(seed):
|
|
rng = np.random.default_rng(seed)
|
|
def initialize(views):
|
|
views['input_frame'][:] = rng.integers(0, 256, views['input_frame'].shape, dtype=np.uint8)
|
|
views['M_inv'][:] = rng.standard_normal(views['M_inv'].shape)*8
|
|
return allocate_inputs(specs, initialize)
|
|
|
|
@TinyJit(prune=True)
|
|
def run(input_frame, M_inv):
|
|
if args.frames == 1: return warp(input_frame, M_inv)
|
|
return Tensor.stack(*(warp(input_frame[i], M_inv[i]) for i in range(args.frames)))
|
|
|
|
expected = benchmark(run, **(inputs:=make_inputs(42)))
|
|
# capture jit
|
|
for _ in range(2): np.testing.assert_array_equal(benchmark(run, **inputs), expected)
|
|
# test jit output actually changes with different inputs
|
|
with np.testing.assert_raises(AssertionError): np.testing.assert_array_equal(benchmark(run, **make_inputs(43)), expected)
|
|
# benchmarks
|
|
for i in range(args.benchmark_runs):
|
|
np.testing.assert_array_equal(benchmark(run, cb=lambda t: print(f" [{i}/{args.benchmark_runs}] {t*1e3:.2f} ms"), **inputs), expected)
|
|
|
|
if args.retargetable: make_retargetable(run)
|
|
|
|
artifact = {'metadata': {}, 'run': run, 'input_specs': specs}
|
|
with open(args.output, 'wb') as f: pickle.dump(artifact, f)
|
|
|
|
# test pickled jit
|
|
with open(args.output, 'rb') as f: loaded = pickle.load(f)
|
|
if args.retargetable: loaded['run'].captured._linear = lower_and_compile(loaded['run'].captured._linear)
|
|
np.testing.assert_array_equal(benchmark(loaded['run'], **make_inputs(42)), expected)
|