optimize _map_line_to_polygon

Author: deanlee

selfdrive/ui/onroad/model_renderer.py

@@ -354,60 +354,61 @@ def _map_line_to_polygon(self, line: np.ndarray, y_off: float, z_off: float, max
     if points.shape[0] == 0:
       return np.empty((0, 2), dtype=np.float32)
 
-    N = points.shape[0]
-    # Generate left and right 3D points in one array using broadcasting
-    offsets = np.array([[0, -y_off, z_off], [0, y_off, z_off]], dtype=np.float32)
-    points_3d = points[None, :, :] + offsets[:, None, :]  # Shape: 2xNx3
-    points_3d = points_3d.reshape(2 * N, 3)  # Shape: (2*N)x3
+    n = points.shape[0]
+    lr_points = np.empty((n, 2, 3), dtype=np.float32)
+    lr_points[:, 0, :] = points
+    lr_points[:, 1, :] = points
+    lr_points[:, 0, 1] -= y_off
+    lr_points[:, 1, 1] += y_off
+    lr_points[:, :, 2] += z_off
 
     # Transform all points to projected space in one operation
-    proj = self._car_space_transform @ points_3d.T  # Shape: 3x(2*N)
-    proj = proj.reshape(3, 2, N)
-    left_proj = proj[:, 0, :]
-    right_proj = proj[:, 1, :]
+    proj = (self._car_space_transform @ lr_points.reshape(-1, 3).T).T.reshape(-1, 2, 3)
 
     # Filter points where z is sufficiently large
-    valid_proj = (np.abs(left_proj[2]) >= 1e-6) & (np.abs(right_proj[2]) >= 1e-6)
-    if not np.any(valid_proj):
+    z_vals = np.abs(proj[:, :, 2])
+    valid_pairs = (z_vals >= 1e-6).all(axis=1)
+    if not valid_pairs.any():
       return np.empty((0, 2), dtype=np.float32)
 
-    # Compute screen coordinates
-    left_screen = left_proj[:2, valid_proj] / left_proj[2, valid_proj][None, :]
-    right_screen = right_proj[:2, valid_proj] / right_proj[2, valid_proj][None, :]
+    proj = proj[valid_pairs]
+
+    screen = proj[:, :, :2] / proj[:, :, 2:3]
+    left_screen = screen[:, 0, :]
+    right_screen = screen[:, 1, :]
 
     # Define clip region bounds
     clip = self._clip_region
     x_min, x_max = clip.x, clip.x + clip.width
     y_min, y_max = clip.y, clip.y + clip.height
 
     # Filter points within clip region
-    left_in_clip = (
-      (left_screen[0] >= x_min) & (left_screen[0] <= x_max) &
-      (left_screen[1] >= y_min) & (left_screen[1] <= y_max)
-    )
-    right_in_clip = (
-      (right_screen[0] >= x_min) & (right_screen[0] <= x_max) &
-      (right_screen[1] >= y_min) & (right_screen[1] <= y_max)
+    in_clip = (
+      (left_screen[:, 0] >= x_min) & (left_screen[:, 0] <= x_max) &
+      (left_screen[:, 1] >= y_min) & (left_screen[:, 1] <= y_max) &
+      (right_screen[:, 0] >= x_min) & (right_screen[:, 0] <= x_max) &
+      (right_screen[:, 1] >= y_min) & (right_screen[:, 1] <= y_max)
     )
-    both_in_clip = left_in_clip & right_in_clip
-
-    if not np.any(both_in_clip):
+    if not in_clip.any():
       return np.empty((0, 2), dtype=np.float32)
 
-    # Select valid and clipped points
-    left_screen = left_screen[:, both_in_clip]
-    right_screen = right_screen[:, both_in_clip]
+    left_screen = left_screen[in_clip]
+    right_screen = right_screen[in_clip]
 
-    # Handle Y-coordinate inversion on hills
-    if not allow_invert and left_screen.shape[1] > 1:
-      y = left_screen[1, :]  # y-coordinates
-      keep = y == np.minimum.accumulate(y)
-      if not np.any(keep):
+    if not allow_invert and left_screen.shape[0] > 1:
+      y_vals = left_screen[:, 1]
+      keep = y_vals == np.minimum.accumulate(y_vals)
+      if not keep.any():
         return np.empty((0, 2), dtype=np.float32)
-      left_screen = left_screen[:, keep]
-      right_screen = right_screen[:, keep]
+      left_screen = left_screen[keep]
+      right_screen = right_screen[keep]
+
+    k = left_screen.shape[0]
+    tri_strip = np.empty((k * 2, 2), dtype=np.float32)
+    tri_strip[0::2] = left_screen
+    tri_strip[1::2] = right_screen
 
-    return np.vstack((left_screen.T, right_screen[:, ::-1].T)).astype(np.float32)
+    return tri_strip
 
   @staticmethod
   def _hsla_to_color(h, s, l, a):

system/ui/lib/shader_polygon.py

@@ -195,22 +195,6 @@ def _configure_shader_color(state: ShaderState, color: Optional[rl.Color],  # no
     state.fill_color_ptr[0:4] = [color.r / 255.0, color.g / 255.0, color.b / 255.0, color.a / 255.0]
     rl.set_shader_value(state.shader, state.locations['fillColor'], state.fill_color_ptr, UNIFORM_VEC4)
 
-
-def triangulate(pts: np.ndarray) -> list[tuple[float, float]]:
-  """Only supports simple polygons with two chains (ribbon)."""
-
-  # TODO: consider deduping close screenspace points
-  # interleave points to produce a triangle strip
-  assert len(pts) % 2 == 0, "Interleaving expects even number of points"
-
-  tri_strip = []
-  for i in range(len(pts) // 2):
-    tri_strip.append(pts[i])
-    tri_strip.append(pts[-i - 1])
-
-  return cast(list, np.array(tri_strip).tolist())
-
-
 def draw_polygon(origin_rect: rl.Rectangle, points: np.ndarray,
                  color: Optional[rl.Color] = None, gradient: Gradient | None = None):  # noqa: UP045
 
@@ -232,12 +216,9 @@ def draw_polygon(origin_rect: rl.Rectangle, points: np.ndarray,
   # Configure gradient shader
   _configure_shader_color(state, color, gradient, origin_rect)
 
-  # Triangulate via interleaving
-  tri_strip = triangulate(pts)
-
   # Draw strip, color here doesn't matter
   rl.begin_shader_mode(state.shader)
-  rl.draw_triangle_strip(tri_strip, len(tri_strip), rl.WHITE)
+  rl.draw_triangle_strip(points.tolist(), len(points), rl.WHITE)
   rl.end_shader_mode()