diff --git a/t4_devkit/viewer/lanelet.py b/t4_devkit/viewer/lanelet.py index 94c826fb..d0c070da 100644 --- a/t4_devkit/viewer/lanelet.py +++ b/t4_devkit/viewer/lanelet.py @@ -30,6 +30,189 @@ } +def _cuboid_mesh( + center: np.ndarray, + size: tuple[float, float, float], + color: list[float], +) -> tuple[np.ndarray, np.ndarray, list[list[float]]]: + half_size = np.asarray(size, dtype=float) / 2.0 + offsets = np.array( + [ + [-1, -1, -1], + [1, -1, -1], + [1, 1, -1], + [-1, 1, -1], + [-1, -1, 1], + [1, -1, 1], + [1, 1, 1], + [-1, 1, 1], + ], + dtype=float, + ) + vertices = center + offsets * half_size + triangles = np.array( + [ + [0, 1, 2], + [0, 2, 3], + [4, 6, 5], + [4, 7, 6], + [0, 4, 5], + [0, 5, 1], + [1, 5, 6], + [1, 6, 2], + [2, 6, 7], + [2, 7, 3], + [3, 7, 4], + [3, 4, 0], + ], + dtype=np.uint32, + ) + colors = [color] * len(vertices) + return vertices, triangles, colors + + +def _disc_mesh( + center: np.ndarray, + radius: float, + color: list[float], + *, + segments: int = 24, +) -> tuple[np.ndarray, np.ndarray, list[list[float]]]: + angles = np.linspace(0.0, 2.0 * np.pi, segments, endpoint=False) + vertices = np.vstack( + [ + center, + np.column_stack( + [ + center[0] + radius * np.cos(angles), + np.full(segments, center[1]), + center[2] + radius * np.sin(angles), + ] + ), + ] + ) + triangles = np.array( + [[0, i, 1 + (i % segments)] for i in range(1, segments + 1)], + dtype=np.uint32, + ) + colors = [color] * len(vertices) + return vertices, triangles, colors + + +def _orient_vertices( + vertices: np.ndarray, + center: np.ndarray, + direction: np.ndarray | None, +) -> np.ndarray: + if direction is None: + return vertices + center + + x_axis = np.array([direction[0], direction[1], 0.0], dtype=float) + norm = np.linalg.norm(x_axis) + if norm == 0.0: + return vertices + center + + x_axis /= norm + z_axis = np.array([0.0, 0.0, 1.0]) + y_axis = np.cross(z_axis, x_axis) + rotation = np.column_stack([x_axis, y_axis, z_axis]) + return vertices @ rotation.T + center + + +def _traffic_light_kind(way_subtype: str) -> str: + # red_green -> pedestrian, other (red_yellow_green) -> vehicle + if way_subtype == "red_green": + return "pedestrian" + return "vehicle" + + +def _traffic_light_mesh( + center: np.ndarray, + direction: np.ndarray | None = None, + *, + kind: str = "vehicle", +) -> rr.Mesh3D: + if kind == "pedestrian": + return _pedestrian_traffic_light_mesh(center, direction) + return _vehicle_traffic_light_mesh(center, direction) + + +def _vehicle_traffic_light_mesh( + center: np.ndarray, direction: np.ndarray | None = None +) -> rr.Mesh3D: + body_color = [0.05, 0.05, 0.05, 1.0] + visor_color = [0.02, 0.02, 0.02, 1.0] + lens_colors = [ + [0.1, 0.9, 0.2, 1.0], + [1.0, 0.8, 0.1, 1.0], + [1.0, 0.1, 0.1, 1.0], + ] + + parts = [] + parts.append(_cuboid_mesh(np.array([0.0, 0.0, 1.0]), (1.8, 0.25, 0.6), body_color)) + parts.append(_cuboid_mesh(np.array([0.0, 0.0, 0.55]), (0.12, 0.12, 0.3), body_color)) + for x_offset, color in zip([-0.55, 0.0, 0.55], lens_colors, strict=True): + lens_center = np.array([x_offset, -0.13, 1.0]) + parts.append( + _cuboid_mesh(lens_center + np.array([0.0, -0.02, 0.0]), (0.42, 0.05, 0.42), visor_color) + ) + parts.append(_disc_mesh(lens_center + np.array([0.0, -0.055, 0.0]), 0.18, color)) + + vertices = [] + triangles = [] + colors = [] + vertex_offset = 0 + for part_vertices, part_triangles, part_colors in parts: + vertices.append(part_vertices) + triangles.append(part_triangles + vertex_offset) + colors.extend(part_colors) + vertex_offset += len(part_vertices) + + return rr.Mesh3D( + vertex_positions=_orient_vertices(np.vstack(vertices), center, direction), + triangle_indices=np.vstack(triangles), + vertex_colors=colors, + ) + + +def _pedestrian_traffic_light_mesh( + center: np.ndarray, + direction: np.ndarray | None = None, +) -> rr.Mesh3D: + body_color = [0.05, 0.05, 0.05, 1.0] + visor_color = [0.02, 0.02, 0.02, 1.0] + lens_colors = [ + [1.0, 0.1, 0.1, 1.0], + [0.1, 0.9, 0.2, 1.0], + ] + + parts = [] + parts.append(_cuboid_mesh(np.array([0.0, 0.0, 1.0]), (0.55, 0.22, 1.0), body_color)) + parts.append(_cuboid_mesh(np.array([0.0, 0.0, 0.35]), (0.1, 0.1, 0.3), body_color)) + for z_offset, color in zip([1.25, 0.75], lens_colors, strict=True): + lens_center = np.array([0.0, -0.12, z_offset]) + parts.append( + _cuboid_mesh(lens_center + np.array([0.0, -0.02, 0.0]), (0.34, 0.05, 0.34), visor_color) + ) + parts.append(_disc_mesh(lens_center + np.array([0.0, -0.055, 0.0]), 0.14, color)) + + vertices = [] + triangles = [] + colors = [] + vertex_offset = 0 + for part_vertices, part_triangles, part_colors in parts: + vertices.append(part_vertices) + triangles.append(part_triangles + vertex_offset) + colors.extend(part_colors) + vertex_offset += len(part_vertices) + + return rr.Mesh3D( + vertex_positions=_orient_vertices(np.vstack(vertices), center, direction), + triangle_indices=np.vstack(triangles), + vertex_colors=colors, + ) + + def render_lanelets(parser: LaneletParser, root_entity: str) -> None: """Render lanelet polygons based on relations. @@ -95,23 +278,43 @@ def render_traffic_elements(parser: LaneletParser, root_entity: str) -> None: color = LANELET_COLORS["traffic_sign"] size = [0.8, 0.8, 0.8] element_type = "sign" - elif "light" in subtype: + elif ( + "light" in subtype + and member.role == "refers" + and way.tags.get("type") == "traffic_light" + ): color = LANELET_COLORS["traffic_light"] size = [0.6, 1.2, 0.6] element_type = "light" + elif "light" in subtype: + continue else: color = [0.8, 0.0, 0.8, 0.9] # Purple size = [0.5, 0.5, 0.5] element_type = "other" - for i, center in enumerate(coords): - entity_path = f"{root_entity}/traffic_elements/{element_type}/{relation.id}_{i}" - + if element_type == "light": + center = np.mean(np.asarray(coords), axis=0) + direction = ( + np.asarray(coords[-1]) - np.asarray(coords[0]) if len(coords) >= 2 else None + ) + kind = _traffic_light_kind(way.tags.get("subtype", "")) + entity_path = ( + f"{root_entity}/traffic_elements/{kind}_light/{relation.id}_{member.ref}" + ) rr.log( - entity_path, - rr.Boxes3D(sizes=[size], centers=[center], colors=[color]), - static=True, + entity_path, _traffic_light_mesh(center, direction, kind=kind), static=True ) + else: + for i, center in enumerate(coords): + entity_path = ( + f"{root_entity}/traffic_elements/{element_type}/{relation.id}_{i}" + ) + rr.log( + entity_path, + rr.Boxes3D(sizes=[size], centers=[center], colors=[color]), + static=True, + ) def render_ways(parser: LaneletParser, root_entity: str) -> None: diff --git a/tests/viewer/test_viewer.py b/tests/viewer/test_viewer.py index 2594a3fc..e7cb20d3 100644 --- a/tests/viewer/test_viewer.py +++ b/tests/viewer/test_viewer.py @@ -1,11 +1,14 @@ from __future__ import annotations import numpy as np +import rerun as rr from pyquaternion import Quaternion from t4_devkit.dataclass import LidarPointCloud +from t4_devkit.lanelet import LaneletParser from t4_devkit.schema import CalibratedSensor, EgoPose, Sensor from t4_devkit.viewer import EntityPath, format_entity +from t4_devkit.viewer.lanelet import _traffic_light_mesh, render_traffic_elements def test_format_entity() -> None: @@ -142,3 +145,98 @@ def test_render_calibration(dummy_viewer, dummy_camera_calibration) -> None: def test_render_map(dummy_viewer, dummy_lanelet_path) -> None: """Test rendering map with `RerunViewer`.""" dummy_viewer.render_map(dummy_lanelet_path) + + +def test_render_map_traffic_light_as_mesh(tmp_path, monkeypatch) -> None: + """Test rendering traffic light positions as meshes.""" + lanelet_path = tmp_path / "lanelet2_map.osm" + lanelet_path.write_text( + """ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +""", + encoding="utf-8", + ) + + logs = [] + + def log(entity_path, entity, *, static=False): + logs.append((entity_path, entity, static)) + + monkeypatch.setattr("t4_devkit.viewer.lanelet.rr.log", log) + + render_traffic_elements(LaneletParser(str(lanelet_path)), "map/vector_map") + + assert len(logs) == 2 + entities = {entity_path: entity for entity_path, entity, _ in logs} + assert set(entities) == { + "map/vector_map/traffic_elements/vehicle_light/200_100", + "map/vector_map/traffic_elements/pedestrian_light/201_101", + } + assert all(isinstance(entity, rr.Mesh3D) for entity in entities.values()) + assert all(static is True for _, _, static in logs) + + default_vertices = np.array( + _traffic_light_mesh(np.array([0.0, 0.0, 0.0])).vertex_positions.as_arrow_array().to_pylist() + ) + default_span = default_vertices.max(axis=0) - default_vertices.min(axis=0) + assert default_span[0] > default_span[2] + assert len(default_vertices) > 100 + + vehicle_vertices = np.array( + entities["map/vector_map/traffic_elements/vehicle_light/200_100"] + .vertex_positions.as_arrow_array() + .to_pylist() + ) + vehicle_span = vehicle_vertices.max(axis=0) - vehicle_vertices.min(axis=0) + assert vehicle_span[1] > vehicle_span[0] + + pedestrian_vertices = np.array( + entities["map/vector_map/traffic_elements/pedestrian_light/201_101"] + .vertex_positions.as_arrow_array() + .to_pylist() + ) + pedestrian_span = pedestrian_vertices.max(axis=0) - pedestrian_vertices.min(axis=0) + assert pedestrian_span[2] > pedestrian_span[0]