diff --git a/t4_devkit/viewer/lanelet.py b/t4_devkit/viewer/lanelet.py
index d0c070d..6a46214 100644
--- a/t4_devkit/viewer/lanelet.py
+++ b/t4_devkit/viewer/lanelet.py
@@ -5,6 +5,8 @@
import numpy as np
import rerun as rr
+from .traffic_light import traffic_light_kind, traffic_light_mesh
+
if TYPE_CHECKING:
from t4_devkit.lanelet import LaneletParser
@@ -30,189 +32,6 @@
}
-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.
@@ -298,12 +117,12 @@ def render_traffic_elements(parser: LaneletParser, root_entity: str) -> None:
direction = (
np.asarray(coords[-1]) - np.asarray(coords[0]) if len(coords) >= 2 else None
)
- kind = _traffic_light_kind(way.tags.get("subtype", ""))
+ 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, _traffic_light_mesh(center, direction, kind=kind), static=True
+ entity_path, traffic_light_mesh(center, direction, kind=kind), static=True
)
else:
for i, center in enumerate(coords):
diff --git a/t4_devkit/viewer/traffic_light.py b/t4_devkit/viewer/traffic_light.py
new file mode 100644
index 0000000..f9c41d5
--- /dev/null
+++ b/t4_devkit/viewer/traffic_light.py
@@ -0,0 +1,193 @@
+from __future__ import annotations
+
+import numpy as np
+import rerun as rr
+
+__all__ = ["traffic_light_kind", "traffic_light_mesh"]
+
+
+def traffic_light_kind(way_subtype: str) -> str:
+ 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 _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 _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], # green
+ [1.0, 0.8, 0.1, 1.0], # yellow
+ [1.0, 0.1, 0.1, 1.0], # red
+ ]
+
+ parts = []
+ # add body
+ parts.append(_cuboid_mesh(np.array([0.0, 0.0, 1.0]), (1.8, 0.25, 0.6), body_color))
+ # add lens (green, yellow, red)
+ 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))
+
+ return _combine_mesh_parts(parts, center, direction)
+
+
+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], # red
+ [0.1, 0.9, 0.2, 1.0], # green
+ ]
+
+ parts = []
+ # add body
+ parts.append(_cuboid_mesh(np.array([0.0, 0.0, 1.0]), (0.55, 0.22, 1.0), body_color))
+ # add lens (red, green)
+ 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))
+
+ return _combine_mesh_parts(parts, center, direction)
+
+
+def _combine_mesh_parts(
+ parts: list[tuple[np.ndarray, np.ndarray, list[list[float]]]],
+ center: np.ndarray,
+ direction: np.ndarray | None,
+) -> rr.Mesh3D:
+ 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,
+ )
diff --git a/tests/viewer/test_traffic_light.py b/tests/viewer/test_traffic_light.py
new file mode 100644
index 0000000..728188a
--- /dev/null
+++ b/tests/viewer/test_traffic_light.py
@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+import numpy as np
+import rerun as rr
+
+from t4_devkit.lanelet import LaneletParser
+from t4_devkit.viewer.lanelet import render_traffic_elements
+from t4_devkit.viewer.traffic_light import traffic_light_mesh
+
+
+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]
diff --git a/tests/viewer/test_viewer.py b/tests/viewer/test_viewer.py
index e7cb20d..2594a3f 100644
--- a/tests/viewer/test_viewer.py
+++ b/tests/viewer/test_viewer.py
@@ -1,14 +1,11 @@
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:
@@ -145,98 +142,3 @@ 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]