Replace ui.py with a Rerun visualizer (commaai#32850)

* Replace ui.py with rerun

* Visualizing radarpoints

* Visualizing all points

* Code clean-up

* Merging matrices into one

* Removing pygame depndency

* Replacing ui.py with rp_visualization.py

* Minor fix, changing color names

* Update README.md

tools/replay/README.md

@@ -19,8 +19,10 @@ tools/replay/replay --demo
 # watch the replay with the normal openpilot UI
 cd selfdrive/ui && ./ui
 
-# or try out a debug visualizer:
-python replay/ui.py
+# or try out radar point visualization in Rerun:
+python replay/rp_visualization.py
+
+# NOTE: To visualize radar points, make sure tools/replay/replay is running.
 ```
 
 ## usage

tools/replay/lib/rp_helpers.py

@@ -0,0 +1,108 @@
+import numpy as np
+from openpilot.selfdrive.controls.radard import RADAR_TO_CAMERA
+
+# Color palette used for rerun AnnotationContext
+rerunColorPalette = [(96, "red", (255, 0, 0)),
+                     (100, "pink", (255, 36, 0)),
+                     (124, "yellow", (255, 255, 0)),
+                     (230, "vibrantpink", (255, 36, 170)),
+                     (240, "orange", (255, 146, 0)),
+                     (255, "white", (255, 255, 255)),
+                     (110, "carColor", (255,0,127))]
+
+
+class UIParams:
+  lidar_x, lidar_y, lidar_zoom = 384, 960, 6
+  lidar_car_x, lidar_car_y = lidar_x / 2., lidar_y / 1.1
+  car_hwidth = 1.7272 / 2 * lidar_zoom
+  car_front = 2.6924 * lidar_zoom
+  car_back = 1.8796 * lidar_zoom
+  car_color = rerunColorPalette[6][0]
+UP = UIParams
+
+
+def to_topdown_pt(y, x):
+  px, py = x * UP.lidar_zoom + UP.lidar_car_x, -y * UP.lidar_zoom + UP.lidar_car_y
+  if px > 0 and py > 0 and px < UP.lidar_x and py < UP.lidar_y:
+    return int(px), int(py)
+  return -1, -1
+
+
+def draw_path(path, lid_overlay, lid_color=None):
+  x, y = np.asarray(path.x), np.asarray(path.y)
+  # draw lidar path point on lidar
+  if lid_color is not None and lid_overlay is not None:
+    for i in range(len(x)):
+      px, py = to_topdown_pt(x[i], y[i])
+      if px != -1:
+        lid_overlay[px, py] = lid_color
+
+
+def plot_model(m, lid_overlay):
+  if lid_overlay is None:
+    return
+  for lead in m.leadsV3:
+    if lead.prob < 0.5:
+      continue
+    x, y = lead.x[0], lead.y[0]
+    x_std = lead.xStd[0]
+    x -= RADAR_TO_CAMERA
+    _, py_top = to_topdown_pt(x + x_std, y)
+    px, py_bottom = to_topdown_pt(x - x_std, y)
+    lid_overlay[int(round(px - 4)):int(round(px + 4)), py_top:py_bottom] = rerunColorPalette[2][0]
+
+  for path in m.laneLines:
+    draw_path(path, lid_overlay, rerunColorPalette[2][0])
+  for edge in m.roadEdges:
+    draw_path(edge, lid_overlay, rerunColorPalette[0][0])
+  draw_path(m.position, lid_overlay, rerunColorPalette[0][0])
+
+
+def plot_lead(rs, lid_overlay):
+  for lead in [rs.leadOne, rs.leadTwo]:
+    if not lead.status:
+      continue
+    x = lead.dRel
+    px_left, py = to_topdown_pt(x, -10)
+    px_right, _ = to_topdown_pt(x, 10)
+    lid_overlay[px_left:px_right, py] = rerunColorPalette[0][0]
+
+
+def maybe_update_radar_points(lt, lid_overlay):
+  ar_pts = []
+  if lt is not None:
+    ar_pts = {}
+    for track in lt:
+      ar_pts[track.trackId] = [track.dRel, track.yRel, track.vRel, track.aRel, track.oncoming, track.stationary]
+  for ids, pt in ar_pts.items():
+    # negative here since radar is left positive
+    px, py = to_topdown_pt(pt[0], -pt[1])
+    if px != -1:
+      if pt[-1]:
+        color = rerunColorPalette[4][0]
+      elif pt[-2]:
+        color = rerunColorPalette[3][0]
+      else:
+        color = rerunColorPalette[5][0]
+      if int(ids) == 1:
+        lid_overlay[px - 2:px + 2, py - 10:py + 10] = rerunColorPalette[1][0]
+      else:
+        lid_overlay[px - 2:px + 2, py - 2:py + 2] = color
+
+
+def get_blank_lid_overlay(UP):
+  lid_overlay = np.zeros((UP.lidar_x, UP.lidar_y), 'uint8')
+  # Draw the car.
+  lid_overlay[int(round(UP.lidar_car_x - UP.car_hwidth)):int(
+    round(UP.lidar_car_x + UP.car_hwidth)), int(round(UP.lidar_car_y -
+                                                      UP.car_front))] = UP.car_color
+  lid_overlay[int(round(UP.lidar_car_x - UP.car_hwidth)):int(
+    round(UP.lidar_car_x + UP.car_hwidth)), int(round(UP.lidar_car_y +
+                                                      UP.car_back))] = UP.car_color
+  lid_overlay[int(round(UP.lidar_car_x - UP.car_hwidth)), int(
+    round(UP.lidar_car_y - UP.car_front)):int(round(
+      UP.lidar_car_y + UP.car_back))] = UP.car_color
+  lid_overlay[int(round(UP.lidar_car_x + UP.car_hwidth)), int(
+    round(UP.lidar_car_y - UP.car_front)):int(round(
+      UP.lidar_car_y + UP.car_back))] = UP.car_color
+  return lid_overlay