diff --git a/selfdrive/ui/mici/onroad/model_renderer.py b/selfdrive/ui/mici/onroad/model_renderer.py index c361db3f5b..0153796479 100644 --- a/selfdrive/ui/mici/onroad/model_renderer.py +++ b/selfdrive/ui/mici/onroad/model_renderer.py @@ -1,4 +1,5 @@ import colorsys +import math import numpy as np import pyray as rl from cereal import messaging, car @@ -37,6 +38,11 @@ LANE_LINE_COLORS = { **LANE_LINE_COLORS_SP, } +LEAD_TRACK_COLORS = ( + rl.Color(255, 215, 0, 255), + rl.Color(255, 140, 0, 220), +) + @dataclass class ModelPoints: @@ -49,6 +55,9 @@ class LeadVehicle: glow: list[float] = field(default_factory=list) chevron: list[float] = field(default_factory=list) fill_alpha: int = 0 + position: tuple[float, float] | None = None + radar_track_id: int = -1 + radar: bool = False class ModelRenderer(Widget, ModelRendererSP): @@ -131,7 +140,7 @@ class ModelRenderer(Widget, ModelRendererSP): model = sm['modelV2'] radar_state = sm['radarState'] if sm.valid['radarState'] else None lead_one = radar_state.leadOne if radar_state else None - render_lead_indicator = self._longitudinal_control and radar_state is not None + render_lead_indicator = ui_state.radar_tracks != 0 and radar_state is not None # Update model data when needed model_updated = sm.updated['modelV2'] @@ -154,11 +163,17 @@ class ModelRenderer(Widget, ModelRendererSP): self._draw_path(sm) if sm.valid['liveTracks'] and sm.recv_frame['liveTracks'] >= ui_state.started_frame: - self.radar_tracks.draw_radar_tracks(sm['liveTracks'], self._map_to_screen, self._path_offset_z, - screen_offset=(self._rect.x, self._rect.y), v_ego=sm['carState'].vEgo) + highlighted_tracks = self._lead_track_colors(radar_state) if render_lead_indicator else {} + matched_positions = self.radar_tracks.draw_radar_tracks( + sm['liveTracks'], self._map_to_screen, self._path_offset_z, + screen_offset=(self._rect.x, self._rect.y), v_ego=sm['carState'].vEgo, + highlighted_tracks=highlighted_tracks, + ) + if render_lead_indicator: + self._draw_lead_connectors(matched_positions, highlighted_tracks) - # if render_lead_indicator and radar_state: - # self._draw_lead_indicator() + if render_lead_indicator: + self._draw_lead_indicator() def _update_raw_points(self, model): """Update raw 3D points from model data""" @@ -188,7 +203,11 @@ class ModelRenderer(Widget, ModelRendererSP): z = self._path.raw_points[idx, 2] if idx < len(self._path.raw_points) else 0.0 point = self._map_to_screen(d_rel, -y_rel + self._camera_offset, z + self._path_offset_z) if point: - self._lead_vehicles[i] = self._update_lead_vehicle(d_rel, v_rel, point, self._rect) + lead_vehicle = self._update_lead_vehicle(d_rel, v_rel, point, self._rect) + lead_vehicle.position = lead_vehicle.chevron[1] + lead_vehicle.radar_track_id = int(lead_data.radarTrackId) + lead_vehicle.radar = lead_data.radar + self._lead_vehicles[i] = lead_vehicle def _update_model(self, lead, path_x_array): """Update model visualization data based on model message""" @@ -383,8 +402,34 @@ class ModelRenderer(Widget, ModelRendererSP): if not lead.glow or not lead.chevron: continue - rl.draw_triangle_fan(lead.glow, len(lead.glow), rl.Color(218, 202, 37, 255)) - rl.draw_triangle_fan(lead.chevron, len(lead.chevron), rl.Color(201, 34, 49, lead.fill_alpha)) + offset_glow = [(x + self._rect.x, y + self._rect.y) for x, y in lead.glow] + offset_chevron = [(x + self._rect.x, y + self._rect.y) for x, y in lead.chevron] + rl.draw_triangle_fan(offset_glow, len(offset_glow), rl.Color(218, 202, 37, 255)) + rl.draw_triangle_fan(offset_chevron, len(offset_chevron), rl.Color(201, 34, 49, lead.fill_alpha)) + + @staticmethod + def _lead_track_colors(radar_state): + highlighted_tracks = {} + if radar_state is None: + return highlighted_tracks + + for lead, color in zip((radar_state.leadOne, radar_state.leadTwo), LEAD_TRACK_COLORS, strict=True): + if lead.status and lead.radar and lead.radarTrackId >= 0: + highlighted_tracks.setdefault(int(lead.radarTrackId), color) + return highlighted_tracks + + def _draw_lead_connectors(self, matched_positions, highlighted_tracks): + for lead in self._lead_vehicles: + if not lead.radar or lead.position is None or lead.radar_track_id not in matched_positions: + continue + + radar_position = matched_positions[lead.radar_track_id] + lead_position = (lead.position[0] + self._rect.x, lead.position[1] + self._rect.y) + if math.dist(lead_position, radar_position) < 4: + continue + + rl.draw_line_ex(rl.Vector2(*lead_position), rl.Vector2(*radar_position), 2, + highlighted_tracks[lead.radar_track_id]) @staticmethod def _get_path_length_idx(pos_x_array: np.ndarray, path_height: float) -> int: diff --git a/selfdrive/ui/sunnypilot/onroad/radar_tracks.py b/selfdrive/ui/sunnypilot/onroad/radar_tracks.py index bb92da0e5e..4a843e6749 100644 --- a/selfdrive/ui/sunnypilot/onroad/radar_tracks.py +++ b/selfdrive/ui/sunnypilot/onroad/radar_tracks.py @@ -135,7 +135,11 @@ class RadarTracksStatus: class RadarTracks: - def draw_radar_tracks(self, live_tracks, map_to_screen, path_offset_z, track_size=7, screen_offset=(0, 0), v_ego=0.0): + def draw_radar_tracks(self, live_tracks, map_to_screen, path_offset_z, track_size=7, screen_offset=(0, 0), v_ego=0.0, + highlighted_tracks=None): + highlighted_tracks = highlighted_tracks or {} + highlighted_positions = {} + for track in live_tracks.points: d_rel, y_rel, v_rel = track.dRel, track.yRel, track.vRel if not (math.isfinite(d_rel) and math.isfinite(y_rel) and math.isfinite(v_rel)): @@ -148,4 +152,13 @@ class RadarTracks: x, y = pt[0] + screen_offset[0], pt[1] + screen_offset[1] color = radar_track_color(v_rel, v_ego) stationary = radar_track_is_stationary(v_rel, v_ego) - rl.draw_circle(int(x), int(y), max(1, track_size - 4) if stationary else track_size, color) + radius = max(1, track_size - 4) if stationary else track_size + track_id = int(track.trackId) + highlight_color = highlighted_tracks.get(track_id) + if highlight_color is not None: + center = rl.Vector2(int(x), int(y)) + rl.draw_ring(center, radius + 2, radius + 5, 0, 360, 24, highlight_color) + highlighted_positions[track_id] = (x, y) + rl.draw_circle(int(x), int(y), radius, color) + + return highlighted_positions diff --git a/selfdrive/ui/sunnypilot/onroad/test_radar_tracks.py b/selfdrive/ui/sunnypilot/onroad/test_radar_tracks.py index 8c5f638b9a..821d700a9e 100644 --- a/selfdrive/ui/sunnypilot/onroad/test_radar_tracks.py +++ b/selfdrive/ui/sunnypilot/onroad/test_radar_tracks.py @@ -126,3 +126,33 @@ def test_draw_radar_tracks_keeps_matched_speed_dots_large(monkeypatch): ) assert drawn_sizes == [6] + + +def test_draw_radar_tracks_highlights_and_returns_matched_track(monkeypatch): + live_tracks = car.RadarData.new_message() + points = live_tracks.init("points", 2) + for track_id, point in enumerate(points, start=10): + point.trackId = track_id + point.dRel = track_id + point.yRel = 1 + point.vRel = 2 + point.aRel = 0 + + highlight_color = radar_tracks.rl.Color(255, 215, 0, 255) + drawn_rings = [] + monkeypatch.setattr(radar_tracks.rl, "draw_circle", lambda *args: None) + monkeypatch.setattr( + radar_tracks.rl, + "draw_ring", + lambda center, inner, outer, start, end, segments, color: drawn_rings.append( + ((center.x, center.y), inner, outer, color_tuple(color)) + ), + ) + + matched_positions = radar_tracks.RadarTracks().draw_radar_tracks( + live_tracks, lambda d_rel, y_rel, z: (d_rel, 30), path_offset_z=1.2, + screen_offset=(100, 7), highlighted_tracks={11: highlight_color}, + ) + + assert drawn_rings == [((111, 37), 9, 12, (255, 215, 0, 255))] + assert matched_positions == {11: (111, 37)}