IQ.Pilot Release Commit @ d2ce8a8

iqpilot/selfdrive/iqmodeld/temporal_state.py

@@ -0,0 +1,131 @@
+"""
+Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
+"""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+
+DEFAULT_FRAME_SKIP = 4
+
+MODEL_INPUT_SPEC: dict[str, tuple[tuple[int, ...], str]] = {
+  "img": ((1, 12, 128, 256), "uint8"),
+  "big_img": ((1, 12, 128, 256), "uint8"),
+  "desire_pulse": ((1, 25, 8), "float32"),
+  "traffic_convention": ((1, 2), "float32"),
+  "features_buffer": ((1, 24, 512), "float32"),
+  "action_t": ((1, 2), "float32"),
+}
+
+
+def spec_from_meta(meta: dict) -> dict[str, tuple[tuple[int, ...], str]] | None:
+  shapes = meta.get("input_shapes")
+  if not shapes:
+    return None
+  return {name: (tuple(shape), "uint8" if name in ("img", "big_img") else "float32")
+          for name, shape in shapes.items()}
+
+
+class TemporalInputState:
+  def __init__(self, frame_skip: int, spec: dict[str, tuple[tuple[int, ...], str]] = MODEL_INPUT_SPEC):
+    self.frame_skip = frame_skip
+    img = spec["img"][0]
+    fb = spec["features_buffer"][0]
+    dp = spec["desire_pulse"][0]
+
+    self.n_frames = img[1] // 6
+    img_q_shape = (frame_skip * (self.n_frames - 1) + 1, 6, img[2], img[3])
+    self._img_shape = img
+    self._fb_shape = fb
+    self._dp_shape = dp
+    feat_dim = math.prod(fb[2:])
+
+    self.img_q = np.zeros(img_q_shape, dtype=np.uint8)
+    self.big_img_q = np.zeros(img_q_shape, dtype=np.uint8)
+    self.feat_q = np.zeros((frame_skip * fb[1], fb[0], feat_dim), dtype=np.float32)
+    self.desire_q = np.zeros((frame_skip * dp[1], dp[0], dp[2]), dtype=np.float32)
+    self.prev_desire = np.zeros(dp[2], dtype=np.float32)
+    self.prev_feat = np.zeros((fb[0], feat_dim), dtype=np.float32)
+
+  @staticmethod
+  def _shift_append(q: np.ndarray, new_val: np.ndarray) -> None:
+    q[:-1] = q[1:]
+    q[-1] = new_val
+
+  def push_and_materialize(self, warped: np.ndarray, desire_pulse: np.ndarray,
+                           traffic_convention: np.ndarray, action_t: np.ndarray,
+                           ) -> dict[str, np.ndarray]:
+    fs = self.frame_skip
+
+    cur = desire_pulse.astype(np.float32).copy()
+    cur[0] = 0
+    pulse = np.where(cur - self.prev_desire > 0.99, cur, 0).astype(np.float32)
+    self.prev_desire[:] = cur
+
+    self._shift_append(self.img_q, warped[0])
+    self._shift_append(self.big_img_q, warped[1])
+    self._shift_append(self.desire_q, pulse.reshape(self._dp_shape[0], self._dp_shape[2]))
+    self._shift_append(self.feat_q, self.prev_feat)
+
+    dp = self._dp_shape
+    return {
+      "img": np.ascontiguousarray(self.img_q[::fs]).reshape(self._img_shape),
+      "big_img": np.ascontiguousarray(self.big_img_q[::fs]).reshape(self._img_shape),
+      "features_buffer": np.ascontiguousarray(self.feat_q[::fs]).reshape(self._fb_shape),
+      "desire_pulse": self.desire_q.reshape(dp[1], fs, dp[0], dp[2]).max(axis=1).reshape(dp),
+      "traffic_convention": traffic_convention.astype(np.float32).reshape(1, -1),
+      "action_t": action_t.astype(np.float32).reshape(1, -1),
+    }
+
+  def note_hidden_state(self, model_output: np.ndarray, hidden_slice: slice) -> None:
+    self.prev_feat[:] = model_output[hidden_slice].reshape(self.prev_feat.shape)
+
+
+class SplitTemporalState:
+
+  def __init__(self, frame_skip: int, img_shape: tuple[int, ...],
+               feature_shape: tuple[int, ...], desire_shape: tuple[int, ...]):
+    self.frame_skip = frame_skip
+    self._img_shape = tuple(img_shape)
+    self._fb_shape = tuple(feature_shape)
+    self._dp_shape = tuple(desire_shape)
+
+    n_frames = img_shape[1] // 6
+    img_q_shape = (frame_skip * (n_frames - 1) + 1, 6, img_shape[2], img_shape[3])
+    self.img_q = np.zeros(img_q_shape, dtype=np.uint8)
+    self.big_img_q = np.zeros(img_q_shape, dtype=np.uint8)
+    self.feat_q = np.zeros((frame_skip * (feature_shape[1] - 1) + 1, feature_shape[0], feature_shape[2]),
+                           dtype=np.float32)
+    self.desire_q = np.zeros((frame_skip * desire_shape[1], desire_shape[0], desire_shape[2]), dtype=np.float32)
+    self.prev_desire = np.zeros(desire_shape[2], dtype=np.float32)
+
+  def materialize_vision(self, warped: np.ndarray, desire: np.ndarray) -> dict[str, np.ndarray]:
+    fs = self.frame_skip
+    cur = desire.astype(np.float32).copy()
+    cur[0] = 0
+    pulse = np.where(cur - self.prev_desire > 0.99, cur, 0).astype(np.float32)
+    self.prev_desire[:] = cur
+
+    TemporalInputState._shift_append(self.img_q, warped[0])
+    TemporalInputState._shift_append(self.big_img_q, warped[1])
+    TemporalInputState._shift_append(self.desire_q, pulse.reshape(self._dp_shape[0], self._dp_shape[2]))
+    return {
+      "img": np.ascontiguousarray(self.img_q[::fs]).reshape(self._img_shape),
+      "big_img": np.ascontiguousarray(self.big_img_q[::fs]).reshape(self._img_shape),
+    }
+
+  def materialize_policy(self, vision_feature: np.ndarray, traffic_convention: np.ndarray,
+                         action_t: np.ndarray | None = None) -> dict[str, np.ndarray]:
+    fs = self.frame_skip
+    TemporalInputState._shift_append(self.feat_q, vision_feature.reshape(self._fb_shape[0], self._fb_shape[2]))
+    dp = self._dp_shape
+    out = {
+      "features_buffer": np.ascontiguousarray(self.feat_q[::fs]).reshape(self._fb_shape),
+      "desire_pulse": self.desire_q.reshape(dp[1], fs, dp[0], dp[2]).max(axis=1).reshape(dp),
+      "traffic_convention": traffic_convention.astype(np.float32).reshape(1, -1),
+    }
+    if action_t is not None:
+      out["action_t"] = action_t.astype(np.float32).reshape(1, -1)
+    return out