Batch small UI lane and road edge projections

selfdrive/ui/mici/onroad/model_renderer.py

@@ -247,20 +247,13 @@ class ModelRenderer(Widget):
     max_distance = np.clip(path_x_array[-1], MIN_DRAW_DISTANCE, MAX_DRAW_DISTANCE)
     max_idx = self._get_path_length_idx(self._lane_lines[0].raw_points[:, 0], max_distance)
 
-    # Update lane lines using raw points
-    line_width_factor = 0.12
-    for i, lane_line in enumerate(self._lane_lines):
-      if i in (1, 2):
-        line_width_factor = 0.16
-      line_width = lane_line_width if lane_line_width is not None else line_width_factor
-      lane_line.projected_points = self._map_line_to_polygon(
-        lane_line.raw_points, line_width * self._lane_line_probs[i], 0.0, max_idx
-      )
-
-    # Update road edges using raw points
-    edge_width = road_edge_width if road_edge_width is not None else line_width_factor
-    for road_edge in self._road_edges:
-      road_edge.projected_points = self._map_line_to_polygon(road_edge.raw_points, edge_width, 0.0, max_idx)
+    line_widths = [0.12, 0.16, 0.16, 0.16] if lane_line_width is None else [lane_line_width] * 4
+    widths = [width * probability for width, probability in zip(line_widths, self._lane_line_probs, strict=True)]
+    widths.extend([0.16 if road_edge_width is None else road_edge_width] * len(self._road_edges))
+    lines = [*self._lane_lines, *self._road_edges]
+    polygons = self._map_lines_to_polygons([line.raw_points for line in lines], widths, max_idx)
+    for line, polygon in zip(lines, polygons, strict=True):
+      line.projected_points = polygon
 
     # Update path using raw points
     if lead and lead.status:
@@ -586,7 +579,39 @@ class ModelRenderer(Widget):
       return np.empty((0, 2), dtype=np.float32)
 
     offsets = np.array([[0, -y_off, z_off], [0, y_off, z_off]], dtype=np.float32)
-    points_3d = points[None, :, :] + offsets[:, None, :]
+    proj, valid = self._project_points(points, offsets[:, None, :])
+
+    screen = proj[:2, :, valid]
+    if not allow_invert:
+      keep = screen[1, 0] == np.minimum.accumulate(screen[1, 0])
+      screen = screen[:, :, keep]
+
+    return np.concatenate((screen[:, 0].T, screen[:, 1, ::-1].T)).astype(np.float32, copy=False)
+
+  def _map_lines_to_polygons(self, lines: list[np.ndarray], y_offsets: list[float], max_idx: int) -> list[np.ndarray]:
+    points = [line[:max_idx + 1] for line in lines]
+    counts = [len(line) for line in points]
+    total = sum(counts)
+    if total == 0:
+      return [np.empty((0, 2), dtype=np.float32) for _ in lines]
+
+    offsets = np.zeros((2, total, 3), dtype=np.float32)
+    offsets[1, :, 1] = np.repeat(y_offsets, counts)
+    offsets[0, :, 1] = -offsets[1, :, 1]
+    proj, valid = self._project_points(np.concatenate(points), offsets)
+
+    polygons = []
+    start = 0
+    for count in counts:
+      end = start + count
+      screen = proj[:2, :, start:end][:, :, valid[start:end]]
+      polygons.append(np.concatenate((screen[:, 0].T, screen[:, 1, ::-1].T)).astype(np.float32, copy=False))
+      start = end
+    return polygons
+
+  def _project_points(self, points: np.ndarray, offsets: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    n = len(points)
+    points_3d = points[None, :, :] + offsets
     proj = (self._car_space_transform @ points_3d.reshape(2 * n, 3).T).reshape(3, 2, n)
 
     valid = (points[:, 0] >= 0) & (np.abs(proj[2, 0]) >= 1e-6) & (np.abs(proj[2, 1]) >= 1e-6)
@@ -598,13 +623,7 @@ class ModelRenderer(Widget):
     for side in (0, 1):
       valid &= ((proj[0, side] >= x_min) & (proj[0, side] <= x_max) &
                 (proj[1, side] >= y_min) & (proj[1, side] <= y_max))
-
-    screen = proj[:2, :, valid]
-    if not allow_invert:
-      keep = screen[1, 0] == np.minimum.accumulate(screen[1, 0])
-      screen = screen[:, :, keep]
-
-    return np.concatenate((screen[:, 0].T, screen[:, 1, ::-1].T)).astype(np.float32, copy=False)
+    return proj, valid
 
   @staticmethod
   def _hsla_to_color(h, s, l, a):

selfdrive/ui/tests/test_mici_model_renderer.py

@@ -101,3 +101,37 @@ def test_changed_transform_invalidates_geometry_and_copies_input(renderer):
   np.testing.assert_array_equal(renderer._car_space_transform, transform)
   transform[0, 0] += 1
   assert renderer._car_space_transform[0, 0] != transform[0, 0]
+
+
+@pytest.mark.parametrize('lengths', [(33,) * 6, (0, 1, 7, 33, 15, 100), (0,) * 6, ()])
+@pytest.mark.parametrize('max_idx', [-1, 0, 15, 100])
+@pytest.mark.parametrize('dtype', [np.float32, np.float64])
+def test_batched_projection_preserves_each_lane(renderer, lengths, max_idx, dtype):
+  rng = np.random.default_rng(67)
+  lines = [np.column_stack((np.linspace(-5, 110, length), rng.normal(0, 6, length), rng.normal(0, 1, length))).astype(dtype)
+           for length in lengths]
+  widths = [0.08 + i * 0.06 for i in range(len(lines))]
+
+  actual = renderer._map_lines_to_polygons(lines, widths, max_idx)
+
+  assert len(actual) == len(lines)
+  for line, width, polygon in zip(lines, widths, actual, strict=True):
+    assert polygon.dtype == np.float32
+    assert polygon.flags.c_contiguous
+    expected = reference_polygon(renderer, line, width, 0, max_idx, True)
+    np.testing.assert_allclose(polygon, expected, rtol=1e-5, atol=1e-3)
+
+
+def test_batched_projection_clips_lines_independently(renderer):
+  renderer._clip_region = SimpleNamespace(x=0, y=0, width=100, height=100)
+  renderer._car_space_transform = np.eye(3, dtype=np.float32)
+  lines = [np.array([[0, 0, 0], [20, 50, 1], [30, 99, 1]], dtype=np.float32),
+           np.empty((0, 3), dtype=np.float32),
+           np.array([[-1, 10, 1], [40, 80, 1], [50, 10, 1]], dtype=np.float32)]
+
+  with np.errstate(divide='raise', invalid='raise'):
+    polygons = renderer._map_lines_to_polygons(lines, [2, 1, 15], 100)
+
+  np.testing.assert_array_equal(polygons[0], [[20, 48], [20, 52]])
+  assert polygons[1].shape == (0, 2)
+  np.testing.assert_array_equal(polygons[2], [[40, 65], [40, 95]])