feat(viewer): add support of rendering traffic light mesh

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:

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(
+        """<?xml version="1.0" encoding="UTF-8"?>
+<osm version="0.6">
+  <node id="1" lat="35.0" lon="139.0" ele="1.0">
+    <tag k="local_x" v="10.0"/>
+    <tag k="local_y" v="20.0"/>
+    <tag k="ele" v="1.0"/>
+  </node>
+  <node id="2" lat="35.0" lon="139.0" ele="1.0">
+    <tag k="local_x" v="10.0"/>
+    <tag k="local_y" v="22.0"/>
+    <tag k="ele" v="1.0"/>
+  </node>
+  <node id="3" lat="35.0" lon="139.0" ele="1.0">
+    <tag k="local_x" v="20.0"/>
+    <tag k="local_y" v="20.0"/>
+    <tag k="ele" v="1.0"/>
+  </node>
+  <node id="4" lat="35.0" lon="139.0" ele="1.0">
+    <tag k="local_x" v="21.0"/>
+    <tag k="local_y" v="20.0"/>
+    <tag k="ele" v="1.0"/>
+  </node>
+  <way id="100">
+    <nd ref="1"/>
+    <nd ref="2"/>
+    <tag k="type" v="traffic_light"/>
+    <tag k="subtype" v="red_yellow_green"/>
+  </way>
+  <way id="101">
+    <nd ref="3"/>
+    <nd ref="4"/>
+    <tag k="type" v="traffic_light"/>
+    <tag k="subtype" v="red_green"/>
+  </way>
+  <relation id="200">
+    <member type="way" ref="100" role="refers"/>
+    <tag k="type" v="regulatory_element"/>
+    <tag k="subtype" v="traffic_light"/>
+  </relation>
+  <relation id="201">
+    <member type="way" ref="101" role="refers"/>
+    <tag k="type" v="regulatory_element"/>
+    <tag k="subtype" v="traffic_light"/>
+  </relation>
+</osm>
+""",
+        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]