IQ.Pilot Release Commit @ 0798119

selfdrive/ui/onroad/environment_renderer.py

@@ -0,0 +1,141 @@
+import numpy as np
+import pyray as rl
+from cereal import custom
+
+from openpilot.selfdrive.ui.ui_state import ui_state
+from openpilot.system.ui.lib.application import gui_app
+from openpilot.system.ui.widgets import Widget
+
+MODE_OFF = 0
+MODE_OVERLAY = 1
+MODE_REPLACE = 2
+
+_ENV_LABEL = custom.IQEnvironment.Object.Label
+_OBJECT_COLORS = {
+  _ENV_LABEL.car: (40, 210, 200),
+  _ENV_LABEL.truck: (40, 210, 200),
+  _ENV_LABEL.bus: (40, 210, 200),
+  _ENV_LABEL.motorcycle: (90, 220, 255),
+  _ENV_LABEL.bicycle: (90, 220, 255),
+  _ENV_LABEL.person: (255, 210, 90),
+  _ENV_LABEL.stopSign: (255, 60, 45),
+  _ENV_LABEL.trafficLight: (255, 190, 0),
+}
+
+_BOX_EDGES = (
+  (0, 1), (1, 3), (3, 2), (2, 0),
+  (4, 5), (5, 7), (7, 6), (6, 4),
+  (0, 4), (1, 5), (2, 6), (3, 7),
+)
+
+GRID_HALF_WIDTH = 12.0
+GRID_MAX_DISTANCE = 90.0
+GRID_STEP = 6.0
+
+
+class EnvironmentRenderer(Widget):
+  def __init__(self):
+    Widget.__init__(self)
+    self._car_space_transform = np.zeros((3, 3), dtype=np.float32)
+    self._mode = MODE_OFF
+    self._counter = 0
+
+  def set_transform(self, transform: np.ndarray):
+    self._car_space_transform = transform.astype(np.float32)
+
+  def _project(self, pt: np.ndarray):
+    p = self._car_space_transform @ pt
+    if abs(p[2]) < 1e-6:
+      return None
+    return p[0] / p[2], p[1] / p[2]
+
+  def _in_rect(self, x: float, y: float) -> bool:
+    r = self._rect
+    return r.x - 400 <= x <= r.x + r.width + 400 and r.y - 400 <= y <= r.y + r.height + 400
+
+  def _render(self, rect: rl.Rectangle):
+    sm = ui_state.sm
+    if self._counter % 30 == 0:
+      self._mode = int(ui_state.params.get("EnvironmentView", return_default=True) or 0) if ui_state.active_bundle else 0
+    self._counter += 1
+
+    if self._mode == MODE_OFF:
+      return
+    if sm.recv_frame["liveCalibration"] < ui_state.started_frame:
+      return
+
+    if self._mode == MODE_REPLACE:
+      self._draw_backdrop(rect)
+      self._draw_ground_grid()
+      if sm.valid["modelV2"]:
+        self._draw_model_scene(sm["modelV2"])
+
+    if sm.alive["iqEnvironment"] and sm.valid["iqEnvironment"]:
+      self._draw_objects(sm["iqEnvironment"])
+
+  def _draw_backdrop(self, rect: rl.Rectangle):
+    rl.draw_rectangle_gradient_v(int(rect.x), int(rect.y), int(rect.width), int(rect.height),
+                                 rl.Color(14, 17, 22, 255), rl.Color(6, 8, 11, 255))
+
+  def _draw_ground_grid(self):
+    col = rl.Color(60, 70, 82, 90)
+    dist = GRID_STEP
+    while dist <= GRID_MAX_DISTANCE:
+      a = self._project(np.array([dist, -GRID_HALF_WIDTH, 0.0]))
+      b = self._project(np.array([dist, GRID_HALF_WIDTH, 0.0]))
+      if a and b and self._in_rect(*a) and self._in_rect(*b):
+        rl.draw_line_ex(rl.Vector2(*a), rl.Vector2(*b), 1.5, col)
+      dist += GRID_STEP
+    for off in np.arange(-GRID_HALF_WIDTH, GRID_HALF_WIDTH + 0.1, 3.0):
+      a = self._project(np.array([GRID_STEP, float(off), 0.0]))
+      b = self._project(np.array([GRID_MAX_DISTANCE, float(off), 0.0]))
+      if a and b and self._in_rect(*a) and self._in_rect(*b):
+        rl.draw_line_ex(rl.Vector2(*a), rl.Vector2(*b), 1.5, col)
+
+  def _draw_polyline(self, xs, ys, zs, color, thick):
+    pts = []
+    for x, y, z in zip(xs, ys, zs, strict=False):
+      if x < 0:
+        continue
+      s = self._project(np.array([x, y, z], dtype=np.float32))
+      if s and self._in_rect(*s):
+        pts.append(rl.Vector2(*s))
+    for i in range(len(pts) - 1):
+      rl.draw_line_ex(pts[i], pts[i + 1], thick, color)
+
+  def _draw_model_scene(self, model):
+    for i, lane in enumerate(model.laneLines):
+      a = int(np.clip(model.laneLineProbs[i], 0.0, 0.9) * 255)
+      self._draw_polyline(lane.x, lane.y, lane.z, rl.Color(235, 235, 235, a), 3.0)
+    for edge in model.roadEdges:
+      self._draw_polyline(edge.x, edge.y, edge.z, rl.Color(230, 70, 70, 180), 3.0)
+    pos = model.position
+    self._draw_polyline(pos.x, pos.y, pos.z, rl.Color(40, 210, 200, 220), 6.0)
+
+  def _draw_objects(self, env):
+    for obj in env.objects:
+      self._draw_box(obj)
+
+  def _draw_box(self, obj):
+    hx, hy = obj.length / 2.0, obj.width / 2.0
+    base = np.array([
+      [obj.x - hx, obj.y - hy, obj.z], [obj.x - hx, obj.y + hy, obj.z],
+      [obj.x + hx, obj.y - hy, obj.z], [obj.x + hx, obj.y + hy, obj.z],
+      [obj.x - hx, obj.y - hy, obj.z + obj.height], [obj.x - hx, obj.y + hy, obj.z + obj.height],
+      [obj.x + hx, obj.y - hy, obj.z + obj.height], [obj.x + hx, obj.y + hy, obj.z + obj.height],
+    ], dtype=np.float32)
+
+    screen = []
+    for corner in base:
+      s = self._project(corner)
+      if s is None or not self._in_rect(*s):
+        return
+      screen.append(s)
+
+    r, g, b = _OBJECT_COLORS.get(obj.label, (40, 210, 200))
+    a = int(np.clip(obj.prob, 0.3, 1.0) * 210)
+    floor = [rl.Vector2(*screen[i]) for i in (0, 1, 3, 2)]
+    rl.draw_triangle(floor[0], floor[1], floor[2], rl.Color(r, g, b, a // 5))
+    rl.draw_triangle(floor[0], floor[2], floor[3], rl.Color(r, g, b, a // 5))
+    for i, j in _BOX_EDGES:
+      rl.draw_line_ex(rl.Vector2(*screen[i]), rl.Vector2(*screen[j]), 2.0, rl.Color(r, g, b, a))