github-actions[bot]
de77c89cf6
version: sunnypilot v2026.002.000 (staging) date: 2026-07-21T07:44:02 master commit: 2c334ede443d7391d27575af8c854a095ba702a8
172 lines
7.2 KiB
Python
172 lines
7.2 KiB
Python
import pickle
|
|
import time
|
|
import numpy as np
|
|
from pathlib import Path
|
|
from tinygrad.tensor import Tensor
|
|
from tinygrad.engine.jit import TinyJit
|
|
from tinygrad.device import Device
|
|
|
|
from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
|
|
from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE
|
|
from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye
|
|
from openpilot.selfdrive.modeld.compile_modeld import NV12Frame, make_frame_prepare as _make_frame_prepare
|
|
|
|
CAMERA_CONFIGS = [
|
|
(_ar_ox_fisheye.width, _ar_ox_fisheye.height),
|
|
(_os_fisheye.width, _os_fisheye.height),
|
|
]
|
|
|
|
|
|
def make_frame_prepare(cam_w, cam_h, model_w, model_h):
|
|
nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
|
|
return _make_frame_prepare(nv12, model_w, model_h)
|
|
|
|
|
|
def warp_pkl_path(w, h):
|
|
from openpilot.selfdrive.modeld.helpers import MODELS_DIR
|
|
return MODELS_DIR / f'warp_{w}x{h}_tinygrad.pkl'
|
|
|
|
|
|
def make_update_img_input(frame_prepare, model_w, model_h):
|
|
def update_img_input_tinygrad(tensor, frame, M_inv):
|
|
M_inv = M_inv.to(Device.DEFAULT)
|
|
new_img = frame_prepare(frame, M_inv)
|
|
tensor.assign(tensor[6:].cat(new_img, dim=0).contiguous())
|
|
return Tensor.cat(tensor[:6], tensor[-6:], dim=0).contiguous().reshape(1, 12, model_h//2, model_w//2)
|
|
return update_img_input_tinygrad
|
|
|
|
|
|
def make_update_both_imgs(frame_prepare, model_w, model_h):
|
|
update_img = make_update_img_input(frame_prepare, model_w, model_h)
|
|
def update_both_imgs_tinygrad(calib_img_buffer, new_img, M_inv,
|
|
calib_big_img_buffer, new_big_img, M_inv_big):
|
|
calib_img_pair = update_img(calib_img_buffer, new_img, M_inv)
|
|
calib_big_img_pair = update_img(calib_big_img_buffer, new_big_img, M_inv_big)
|
|
return calib_img_pair, calib_big_img_pair
|
|
return update_both_imgs_tinygrad
|
|
|
|
MODELS_DIR = Path(__file__).parent / 'models'
|
|
MODEL_W, MODEL_H = MEDMODEL_INPUT_SIZE
|
|
UPSTREAM_BUFFER_LENGTH = 5
|
|
|
|
|
|
def v2_warp_pkl_path(cam_w, cam_h, buffer_length):
|
|
return MODELS_DIR / f'warp_{cam_w}x{cam_h}_b{buffer_length}_tinygrad.pkl'
|
|
|
|
|
|
def compile_v2_warp(cam_w, cam_h, buffer_length):
|
|
_, _, _, yuv_size = get_nv12_info(cam_w, cam_h)
|
|
img_buffer_shape = (buffer_length * 6, MODEL_H // 2, MODEL_W // 2)
|
|
|
|
print(f"Compiling v2 warp for {cam_w}x{cam_h} buffer_length={buffer_length}...")
|
|
|
|
frame_prepare = make_frame_prepare(cam_w, cam_h, MODEL_W, MODEL_H)
|
|
update_both_imgs = make_update_both_imgs(frame_prepare, MODEL_W, MODEL_H)
|
|
update_img_jit = TinyJit(update_both_imgs, prune=True)
|
|
|
|
full_buffer = Tensor.zeros(img_buffer_shape, dtype='uint8').contiguous().realize()
|
|
big_full_buffer = Tensor.zeros(img_buffer_shape, dtype='uint8').contiguous().realize()
|
|
new_frame_np = np.random.randint(0, 256, yuv_size, dtype=np.uint8)
|
|
new_big_frame_np = np.random.randint(0, 256, yuv_size, dtype=np.uint8)
|
|
for i in range(10):
|
|
img_inputs = [full_buffer,
|
|
Tensor.from_blob(new_frame_np.ctypes.data, (yuv_size,), dtype='uint8').realize(),
|
|
Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')]
|
|
big_img_inputs = [big_full_buffer,
|
|
Tensor.from_blob(new_big_frame_np.ctypes.data, (yuv_size,), dtype='uint8').realize(),
|
|
Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')]
|
|
inputs = img_inputs + big_img_inputs
|
|
Device.default.synchronize()
|
|
|
|
st = time.perf_counter()
|
|
_ = update_img_jit(*inputs)
|
|
mt = time.perf_counter()
|
|
Device.default.synchronize()
|
|
et = time.perf_counter()
|
|
print(f" [{i+1}/10] enqueue {(mt-st)*1e3:6.2f} ms -- total {(et-st)*1e3:6.2f} ms")
|
|
|
|
pkl_path = v2_warp_pkl_path(cam_w, cam_h, buffer_length)
|
|
with open(pkl_path, "wb") as f:
|
|
pickle.dump(update_img_jit, f)
|
|
print(f" Saved to {pkl_path}")
|
|
|
|
jit = pickle.load(open(pkl_path, "rb"))
|
|
verify_frame = np.random.randint(0, 256, yuv_size, dtype=np.uint8)
|
|
verify_big_frame = np.random.randint(0, 256, yuv_size, dtype=np.uint8)
|
|
fresh_inputs = [
|
|
Tensor.zeros(img_buffer_shape, dtype='uint8').contiguous().realize(),
|
|
Tensor.from_blob(verify_frame.ctypes.data, (yuv_size,), dtype='uint8').realize(),
|
|
Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY'),
|
|
Tensor.zeros(img_buffer_shape, dtype='uint8').contiguous().realize(),
|
|
Tensor.from_blob(verify_big_frame.ctypes.data, (yuv_size,), dtype='uint8').realize(),
|
|
Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY'),
|
|
]
|
|
jit(*fresh_inputs)
|
|
|
|
|
|
class Warp:
|
|
def __init__(self, buffer_length=2):
|
|
self.buffer_length = buffer_length
|
|
self.img_buffer_shape = (buffer_length * 6, MODEL_H // 2, MODEL_W // 2)
|
|
|
|
self.jit_cache = {}
|
|
self.full_buffers = {k: Tensor.zeros(self.img_buffer_shape, dtype='uint8').contiguous().realize() for k in ['img', 'big_img']}
|
|
self._blob_cache: dict[int, Tensor] = {}
|
|
self._nv12_cache: dict[tuple[int, int], int] = {}
|
|
self.transforms_np = {k: np.zeros((3, 3), dtype=np.float32) for k in ['img', 'big_img']}
|
|
self.transforms = {k: Tensor(v, device='NPY').realize() for k, v in self.transforms_np.items()}
|
|
|
|
def process(self, bufs, transforms):
|
|
if not bufs:
|
|
return {}
|
|
road = next(n for n in bufs if 'big' not in n)
|
|
wide = next(n for n in bufs if 'big' in n)
|
|
cam_w, cam_h = bufs[road].width, bufs[road].height
|
|
key = (cam_w, cam_h)
|
|
|
|
if key not in self.jit_cache:
|
|
v2_pkl = v2_warp_pkl_path(cam_w, cam_h, self.buffer_length)
|
|
if v2_pkl.exists():
|
|
with open(v2_pkl, 'rb') as f:
|
|
self.jit_cache[key] = pickle.load(f)
|
|
elif self.buffer_length == UPSTREAM_BUFFER_LENGTH:
|
|
upstream_pkl = warp_pkl_path(cam_w, cam_h)
|
|
if upstream_pkl.exists():
|
|
with open(upstream_pkl, 'rb') as f:
|
|
self.jit_cache[key] = pickle.load(f)
|
|
if key not in self.jit_cache:
|
|
frame_prepare = make_frame_prepare(cam_w, cam_h, MODEL_W, MODEL_H)
|
|
update_both_imgs = make_update_both_imgs(frame_prepare, MODEL_W, MODEL_H)
|
|
self.jit_cache[key] = TinyJit(update_both_imgs, prune=True)
|
|
|
|
if key not in self._nv12_cache:
|
|
self._nv12_cache[key] = get_nv12_info(cam_w, cam_h)[3]
|
|
yuv_size = self._nv12_cache[key]
|
|
|
|
road_ptr = bufs[road].data.ctypes.data
|
|
wide_ptr = bufs[wide].data.ctypes.data
|
|
if road_ptr not in self._blob_cache:
|
|
self._blob_cache[road_ptr] = Tensor.from_blob(road_ptr, (yuv_size,), dtype='uint8')
|
|
if wide_ptr not in self._blob_cache:
|
|
self._blob_cache[wide_ptr] = Tensor.from_blob(wide_ptr, (yuv_size,), dtype='uint8')
|
|
road_blob = self._blob_cache[road_ptr]
|
|
wide_blob = self._blob_cache[wide_ptr] if wide_ptr != road_ptr else Tensor.from_blob(wide_ptr, (yuv_size,), dtype='uint8')
|
|
np.copyto(self.transforms_np['img'], transforms[road].reshape(3, 3))
|
|
np.copyto(self.transforms_np['big_img'], transforms[wide].reshape(3, 3))
|
|
|
|
Device.default.synchronize()
|
|
res = self.jit_cache[key](
|
|
self.full_buffers['img'], road_blob, self.transforms['img'],
|
|
self.full_buffers['big_img'], wide_blob, self.transforms['big_img'],
|
|
)
|
|
out_road = res[0].realize()
|
|
out_wide = res[1].realize()
|
|
|
|
return {road: out_road, wide: out_wide}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
for cam_w, cam_h in CAMERA_CONFIGS:
|
|
for bl in [2, 5]:
|
|
compile_v2_warp(cam_w, cam_h, bl)
|