mod9fa084ee42

openpilot/selfdrive/modeld/models/driving_supercombo.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:65a08adc31d5c456219687d99b7bf5e44d61dae2d49ea67850e76105c7248cce
-size 60918562
+oid sha256:fbbb4c07c0de022af7b28c60ca34794b9326825375da37af6ae463c2fee4961e
+size 83124285

openpilot/sunnypilot/modeld_v2/compile_warp.py

@@ -17,16 +17,18 @@ def main():
   parser.add_argument("--frame", required=True, help="width,height,stride,y_height,uv_height,buffer_size")
   parser.add_argument("--warp-to", required=True, help="model input WxH")
   parser.add_argument("--output", required=True)
+  parser.add_argument("--light", action="store_true")
   args = parser.parse_args()
 
   from examples.openpilot.compile_warp import make_frame_prepare, parse_frame, parse_size
   from tinygrad import Device, Tensor, TinyJit
   frame = parse_frame(args.frame)
   warp = make_frame_prepare(frame, *parse_size(args.warp_to))
-  run = TinyJit(lambda input_frame, M_inv: Tensor.stack(*(warp(input_frame[i], M_inv[i]) for i in range(2))), prune=True)
+  n = 1 if args.light else 2
+  run = TinyJit(lambda input_frame, M_inv: Tensor.stack(*(warp(input_frame[i], M_inv[i]) for i in range(n))), prune=True)
   rng = np.random.default_rng(42)
   inputs = {"input_frame": Tensor(rng.integers(0, 256, (2, frame.size), dtype=np.uint8), device=Device.DEFAULT).realize(),
-            "M_inv": Tensor((rng.standard_normal((2, 3, 3)) * 8).astype(np.float32), device=Device.DEFAULT).realize()}
+            "M_inv": Tensor((rng.standard_normal((n, 3, 3)) * 8).astype(np.float32), device=Device.DEFAULT).realize()}
   with record_kernel_linear() as recorded:
     run(**inputs)
     run(**inputs)
@@ -34,7 +36,7 @@ def main():
   assert len(recorded) == 1, f"expected one captured kernel list, got {len(recorded)}"
   part = serialize_jit(run, recorded[0], check)
   PlainJit(*part)
-  specs = {"input_frame": ((2, frame.size), "|u1", Device.DEFAULT), "M_inv": ((2, 3, 3), "<f4", Device.DEFAULT)}
+  specs = {"input_frame": ((2, frame.size), "|u1", Device.DEFAULT), "M_inv": ((n, 3, 3), "<f4", Device.DEFAULT)}
   write_bundle(args.output, {"run": part}, {"input_specs": specs, "metadata": {}})
   print(f"Saved warp to {args.output} ({len(part[0]['calls'])} kernels)")
 

openpilot/sunnypilot/modeld_v2/frame_fusion.py

@@ -20,6 +20,11 @@ def pack_layout(shapes: dict[str, tuple], frame_size: int) -> tuple[dict[str, in
   return offsets | {"frames": offset}, offset + 2 * frame_size
 
 
+def host_shapes(input_shapes: dict[str, tuple], states) -> dict[str, tuple]:
+  shapes = {name: shape for name, shape in input_shapes.items() if name not in states and name not in ("new_img", "light_img")}
+  return shapes | ({"light_tfm": (1, 3, 3)} if "light_img" in input_shapes else {})
+
+
 class _Blobs:  # the model's and the warp's blobs as one address space, for a part whose ranges span both files
   def __init__(self, *blobs):
     self.blobs, self.starts = blobs, list(itertools.accumulate([0] + [len(blob) for blob in blobs[:-1]]))
@@ -30,18 +35,23 @@ class _Blobs:  # the model's and the warp's blobs as one address space, for a pa
 
 
 def fused_with_warp(cam_w: int, cam_h: int, chestnut: bool):  # load_model's extend: a native model runs upload, warp and model as one submit
-  return partial(frame_part, f"{BASEDIR}/openpilot/sunnypilot/modeld_v2/models/{'big_' if chestnut else ''}driving_warp_{cam_w}x{cam_h}_tinygrad.pkl")
+  models = f"{BASEDIR}/openpilot/sunnypilot/modeld_v2/models"
+  return partial(frame_part, f"{models}/{'big_' if chestnut else ''}driving_warp_{cam_w}x{cam_h}_tinygrad.pkl",
+                 light_warp_path=f"{models}/driving_lightwarp_{cam_w}x{cam_h}_tinygrad.pkl")
 
 
-def frame_part(warp_path: str, parts: dict, extras: dict) -> dict:
+def frame_part(warp_path: str, parts: dict, extras: dict, light_warp_path: str | None = None) -> dict:
   if "input_specs" not in extras:  # a legacy model warps in its own parts
     return {}
   (model, model_blob), (warp, warp_blob) = parts["run"], read_parts(warp_path)[0]["run"]
   states = sorted(name for name in extras["input_specs"] if f"next_{name}" in extras["metadata"]["output_shapes"])
-  host = {name: shape for name, (shape, _, _) in extras["input_specs"].items() if name not in states and name != "new_img"}
+  host = host_shapes({name: shape for name, (shape, _, _) in extras["input_specs"].items()}, states)
+  light = "light_img" in extras["input_specs"]
+  light_warp, light_blob = read_parts(light_warp_path)[0]["run"] if light else ({"buffers": [], "calls": [], "inputs": []}, memoryview(b""))
   frame_size = next(s["numel"] for s in warp["inputs"] if s["name"] == "input_frame") // 2
   offsets, packed = pack_layout(host, frame_size)
-  shift, upload, u8 = len(model_blob), len(model["buffers"]) + len(warp["buffers"]), repr(_dtype("dtypes.uint8"))
+  shift, light_shift = len(model_blob), len(model_blob) + len(warp_blob)
+  upload, u8 = len(model["buffers"]) + len(warp["buffers"]) + len(light_warp["buffers"]), repr(_dtype("dtypes.uint8"))
 
   def view(name: str, nbytes: int, dtype: str) -> dict:
     return {"buffer": upload, "offset": offsets[name], "nbytes": nbytes, "dtype": dtype}
@@ -49,12 +59,15 @@ def frame_part(warp_path: str, parts: dict, extras: dict) -> dict:
   names = sorted(["host", *states])
   slots = {s["slot"]: names.index(s["name"]) for s in model["inputs"] if s["name"] in states} | \
           {slot: len(names) + i for i, slot in enumerate(outputs)}
-  returned = warp["returns"]["tensors"][0]
+  def warped(part: dict, base: int) -> dict:
+    returned = part["returns"]["tensors"][0]
+    return {k: v for k, v in returned.items() if k != "shape"} | {"buffer": returned["buffer"] + base}
   model_inputs = {s["slot"]: {"input": slots[s["slot"]]} if s["slot"] in slots else
-                  {k: v for k, v in returned.items() if k != "shape"} | {"buffer": returned["buffer"] + len(model["buffers"])} if s["name"] == "new_img" else
+                  warped(warp, len(model["buffers"])) if s["name"] == "new_img" else
+                  warped(light_warp, len(model["buffers"]) + len(warp["buffers"])) if s["name"] == "light_img" else
                   view(s["name"], s["numel"] * _dtype(s["dtype"]).itemsize, s["dtype"]) for s in model["inputs"]}
-  warp_inputs = {s["slot"]: view("frames" if s["name"] == "input_frame" else "tfm", s["numel"] * _dtype(s["dtype"]).itemsize, s["dtype"])
-                 for s in warp["inputs"]}
+  def warp_inputs(part: dict, tfm: str) -> dict:
+    return {s["slot"]: view("frames" if s["name"] == "input_frame" else tfm, s["numel"] * _dtype(s["dtype"]).itemsize, s["dtype"]) for s in part["inputs"]}
 
   def remap(calls: list, inputs: dict, buffers: int, lib: int) -> list:
     return [call | ({"lib": [call["lib"][0] + lib, call["lib"][1] + lib]} if "lib" in call else {}) |
@@ -62,8 +75,11 @@ def frame_part(warp_path: str, parts: dict, extras: dict) -> dict:
   header = {"input_names": names, "returns": {"kind": "NoneType", "tensors": []}, "check": None,
             "inputs": [{"slot": names.index("host"), "name": "host", "dtype": u8, "numel": packed, "device": "NPY"}] +
                       [s | {"slot": slots[s["slot"]]} for s in model["inputs"] if s["slot"] in slots],
-            "buffers": model["buffers"] + [b | ({"init": [b["init"][0] + shift, b["init"][1] + shift]} if "init" in b else {}) for b in warp["buffers"]] +
+            "buffers": model["buffers"] + [b | ({"init": [b["init"][0] + d, b["init"][1] + d]} if "init" in b else {})
+                                           for part, d in ((warp, shift), (light_warp, light_shift)) for b in part["buffers"]] +
                        [{"nbytes": packed, "device": model["inputs"][0]["device"], "dtype": u8}],
             "calls": [{"kind": "copy", "args": [{"buffer": upload, "offset": 0, "nbytes": packed, "dtype": u8}, {"input": names.index("host")}]}] +
-                     remap(warp["calls"], warp_inputs, len(model["buffers"]), shift) + remap(model["calls"], model_inputs, 0, 0)}
-  return {"run": (header, _Blobs(model_blob, warp_blob))}  # replaces the model's run: one set of buffers, no separate warp
+                     remap(warp["calls"], warp_inputs(warp, "tfm"), len(model["buffers"]), shift) +
+                     remap(light_warp["calls"], warp_inputs(light_warp, "light_tfm"), len(model["buffers"]) + len(warp["buffers"]), light_shift) +
+                     remap(model["calls"], model_inputs, 0, 0)}
+  return {"run": (header, _Blobs(model_blob, warp_blob, light_blob) if light else _Blobs(model_blob, warp_blob))}

openpilot/sunnypilot/modeld_v2/light_frame.py

@@ -0,0 +1,22 @@
+"""
+Copyright (c) 2021-, sunnypilot, and a number of other contributors.
+
+This file is part of sunnypilot and is licensed under the MIT License.
+See the LICENSE.md file in the root directory for more details.
+"""
+import numpy as np
+
+from openpilot.common.transformations.camera import get_view_frame_from_calib_frame, view_frame_from_device_frame
+from openpilot.common.transformations.orientation import rot_from_euler
+
+LIGHT_INPUT_SIZE = (512, 256)
+light_intrinsics = np.array([
+  [1000.0, 0.0, 256.0],
+  [0.0, 1000.0, 200.0],
+  [0.0, 0.0, 1.0]])
+calib_from_light = np.linalg.inv((light_intrinsics @ get_view_frame_from_calib_frame(0, 0, 0, 0))[:, :3])
+
+
+def get_light_warp_matrix(device_from_calib_euler: np.ndarray, intrinsics: np.ndarray) -> np.ndarray:
+  camera_from_calib = intrinsics @ view_frame_from_device_frame @ rot_from_euler(device_from_calib_euler)
+  return (camera_from_calib @ calib_from_light).astype(np.float32)

openpilot/sunnypilot/modeld_v2/model_adapters.py

@@ -11,7 +11,7 @@ import pickle
 import numpy as np
 
 from openpilot.common.basedir import BASEDIR
-from openpilot.sunnypilot.modeld_v2.frame_fusion import pack_layout
+from openpilot.sunnypilot.modeld_v2.frame_fusion import host_shapes, pack_layout
 from openpilot.sunnypilot.modeld_v2.helpers import load_model
 from openpilot.sunnypilot.modeld_v2.compile_modeld import (POLICY_INPUTS, derive_frame_skip,
                                                            make_split_input_queues, make_supercombo_input_queues)
@@ -172,7 +172,7 @@ class NativeTinygradAdapter(BaseModelAdapter):
 
     self.input_queues = {name: Tensor(np.zeros(shape, dtype=dtype), device=self.model_device).realize()
                          for name, (shape, dtype) in self.input_shapes.items() if name in self.state_pairs}
-    shapes = {name: shape for name, (shape, _) in self.input_shapes.items() if name not in self.state_pairs and name != 'new_img'}
+    shapes = host_shapes({name: shape for name, (shape, _) in self.input_shapes.items()}, self.state_pairs)
     offsets, packed = pack_layout(shapes, self.warp_frame_size)
     self.packed_input = np.zeros(packed, dtype=np.uint8)
     self.host = Tensor(self.packed_input, device='NPY')

openpilot/sunnypilot/modeld_v2/modeld.py

@@ -44,6 +44,7 @@ from openpilot.sunnypilot.modeld_v2.blendful import Blendful
 from openpilot.sunnypilot.modeld_v2.constants import ModelConstants, Plan
 from openpilot.sunnypilot.modeld_v2.meta_helper import load_meta_constants
 from openpilot.sunnypilot.modeld_v2.camera_offset_helper import CameraOffsetHelper
+from openpilot.sunnypilot.modeld_v2.light_frame import get_light_warp_matrix
 from openpilot.sunnypilot.livedelay.helpers import get_lat_delay
 from openpilot.sunnypilot.modeld_v2.modeld_base import ModelStateBase
 from openpilot.sunnypilot.modeld_v2.helpers import load_model
@@ -171,7 +172,7 @@ class ModelState(ModelStateBase):
       np.where(inputs[desire_key] - self.prev_desire > .99, inputs[desire_key], 0)
     self.prev_desire[:] = inputs[desire_key]
 
-    for key in ('traffic_convention', 'lateral_control_params', 'action_t'):
+    for key in ('traffic_convention', 'lateral_control_params', 'action_t', 'speed_limit'):
       if key in self.numpy_inputs and key in inputs:
         self.numpy_inputs[key][:] = inputs[key]
 
@@ -179,6 +180,8 @@ class ModelState(ModelStateBase):
       for i, key in enumerate(self._vision_input_names):
         if key in transforms:
           self.numpy_inputs['tfm'][i] = transforms[key].reshape(3, 3)
+      if 'light_tfm' in self.numpy_inputs and 'light' in transforms:
+        self.numpy_inputs['light_tfm'][0] = transforms['light'].reshape(3, 3)
     else:
       self.numpy_inputs['tfm'][:, :] = transforms[self._road_key].reshape(3, 3)
       self.numpy_inputs['big_tfm'][:, :] = transforms[self._wide_key].reshape(3, 3)
@@ -338,7 +341,7 @@ def main(demo=False):
   pub_socks = ["modelV2", "drivingModelData", "cameraOdometry", "modelDataV2SP"] + (["chestnutGpuState"] if CHESTNUT else [])
   pm = PubMaster(pub_socks)
   sm = SubMaster(["deviceState", "carState", "narrowRoadCameraState", "extrinsicsCalibration", "driverMonitoringState",
-                  "carControl", "controlsState", "lateralDelay"])
+                  "carControl", "controlsState", "lateralDelay", "longitudinalPlanSP"])
 
   publish_state = PublishState()
   chestnut_state = ChestnutGpuState(pm, model.chestnut) if CHESTNUT else None
@@ -351,6 +354,7 @@ def main(demo=False):
 
   model_transform_main = np.zeros((3, 3), dtype=np.float32)
   model_transform_extra = np.zeros((3, 3), dtype=np.float32)
+  model_transform_light = np.zeros((3, 3), dtype=np.float32)
   live_calib_seen = False
   buf_main, buf_extra = None, None
   meta_main = FrameMeta()
@@ -425,6 +429,7 @@ def main(demo=False):
       model_transform_main = get_warp_matrix(device_from_calib_euler, main_intrinsics, False).astype(np.float32)
       model_transform_extra = get_warp_matrix(device_from_calib_euler, dc.wide_road.intrinsics, True).astype(np.float32)
       model_transform_main, model_transform_extra = camera_offset_helper.update(model_transform_main, model_transform_extra, sm, main_wide_camera)
+      model_transform_light = get_light_warp_matrix(device_from_calib_euler, main_intrinsics)
       live_calib_seen = True
 
     traffic_convention = np.zeros(2)
@@ -446,6 +451,7 @@ def main(demo=False):
 
     bufs = {name: buf_extra if 'big' in name else buf_main for name in model.vision_input_names}
     transforms = {name: model_transform_extra if 'big' in name else model_transform_main for name in model.vision_input_names}
+    transforms['light'] = model_transform_light
 
     frame_delay = DT_MDL # compensate for time passed since the frame was captured: current_time - timestamp_eof is 50ms on average
     action_delay = DT_MDL / 2 # middle of the interval between model output (current state) and next frame (expected state)
@@ -463,6 +469,10 @@ def main(demo=False):
     if 'action_t' in model.numpy_inputs:
       inputs['action_t'] = np.array([lat_action_t, long_action_t], dtype=np.float32)
 
+    if 'speed_limit' in model.numpy_inputs:
+      resolver = sm['longitudinalPlanSP'].speedLimit.resolver
+      inputs['speed_limit'] = np.array([resolver.speedLimitFinal if resolver.speedLimitValid else 0.0], dtype=np.float32)
+
     mt1 = time.perf_counter()
     try:
       send_chestnut = (chestnut_state is not None and

openpilot/sunnypilot/modeld_v2/models/driving_lightwarp_1344x760_tinygrad.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c08cffa27671d6773e316fdf4dcab80ffce2c4cc6fa28775615bc6726a63d7ba
+size 3473263

openpilot/sunnypilot/modeld_v2/models/driving_lightwarp_1928x1208_tinygrad.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6cddff2e3320057735fd4812d1ef9cc7d6d9fe405e3b6e7c4df85ee654a0e886
+size 7699823