From e4191db37f6e788ea6e18266276dc4532931184b Mon Sep 17 00:00:00 2001 From: ktro2828 Date: Fri, 18 Oct 2024 02:22:27 +0900 Subject: [PATCH] feat: add `Tier4Viewer` Signed-off-by: ktro2828 --- t4_devkit/dataclass/pointcloud.py | 2 +- t4_devkit/tier4.py | 6 +- t4_devkit/viewer/__init__.py | 2 + t4_devkit/{common => viewer}/color.py | 0 t4_devkit/viewer/rendering_data.py | 107 ++++++++++ t4_devkit/viewer/viewer.py | 288 ++++++++++++++++++++++++++ tests/viewer/test_viewer.py | 76 +++++++ 7 files changed, 477 insertions(+), 4 deletions(-) create mode 100644 t4_devkit/viewer/__init__.py rename t4_devkit/{common => viewer}/color.py (100%) create mode 100644 t4_devkit/viewer/rendering_data.py create mode 100644 t4_devkit/viewer/viewer.py create mode 100644 tests/viewer/test_viewer.py diff --git a/t4_devkit/dataclass/pointcloud.py b/t4_devkit/dataclass/pointcloud.py index 6278813..ac37de3 100644 --- a/t4_devkit/dataclass/pointcloud.py +++ b/t4_devkit/dataclass/pointcloud.py @@ -12,7 +12,7 @@ from t4_devkit.typing import NDArrayFloat, NDArrayU8 -__all__ = ["LidarPointCloud", "RadarPointCloud", "SegmentationPointCloud"] +__all__ = ["PointCloud", "LidarPointCloud", "RadarPointCloud", "SegmentationPointCloud"] @dataclass diff --git a/t4_devkit/tier4.py b/t4_devkit/tier4.py index 8898b23..9fe16a1 100644 --- a/t4_devkit/tier4.py +++ b/t4_devkit/tier4.py @@ -12,7 +12,6 @@ from PIL import Image from pyquaternion import Quaternion -from t4_devkit.common.color import distance_color from t4_devkit.common.geometry import is_box_in_image, view_points from t4_devkit.common.timestamp import sec2us, us2sec from t4_devkit.dataclass import ( @@ -25,6 +24,7 @@ convert_label, ) from t4_devkit.schema import SchemaName, SensorModality, VisibilityLevel, build_schema +from t4_devkit.viewer import distance_color if TYPE_CHECKING: from rerun.blueprint.api import BlueprintLike, Container, SpaceView @@ -1287,8 +1287,8 @@ def _render_sensor_calibration(self, sample_data_token: str) -> None: f"world/ego_vehicle/{sensor_name}", rr.Pinhole( image_from_camera=calibrated_sensor.camera_intrinsic, - width=sample_data.width, - height=sample_data.height, + # width=sample_data.width, + # height=sample_data.height, ), static=True, ) diff --git a/t4_devkit/viewer/__init__.py b/t4_devkit/viewer/__init__.py new file mode 100644 index 0000000..65f9157 --- /dev/null +++ b/t4_devkit/viewer/__init__.py @@ -0,0 +1,2 @@ +from .color import * # noqa +from .viewer import * # noqa diff --git a/t4_devkit/common/color.py b/t4_devkit/viewer/color.py similarity index 100% rename from t4_devkit/common/color.py rename to t4_devkit/viewer/color.py diff --git a/t4_devkit/viewer/rendering_data.py b/t4_devkit/viewer/rendering_data.py new file mode 100644 index 0000000..3461031 --- /dev/null +++ b/t4_devkit/viewer/rendering_data.py @@ -0,0 +1,107 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +import rerun as rr + +if TYPE_CHECKING: + from t4_devkit.dataclass import Box2D, Box3D + from t4_devkit.typing import RoiType, SizeType, TranslationType, VelocityType + + +class BoxData3D: + """A class to store 3D boxes data for rendering.""" + + def __init__(self) -> None: + self._centers: list[TranslationType] = [] + self._rotations: list[rr.Quaternion] = [] + self._sizes: list[SizeType] = [] + self._class_ids: list[int] = [] + self._uuids: list[int] = [] + self._velocities: list[VelocityType] = [] + + def append(self, box: Box3D) -> None: + """Append a 3D box data. + + Args: + box (Box3D): `Box3D` object. + """ + self._centers.append(box.position) + + rotation_xyzw = np.roll(box.rotation.q, shift=-1) + self._rotations.append(rr.Quaternion(xyzw=rotation_xyzw)) + + width, length, height = box.size + self._sizes.append((length, width, height)) + + self._class_ids.append(box.semantic_label.label.value) + + if box.uuid is not None: + self._uuids.append(box.uuid[:6]) + + if box.velocity is not None: + self._velocities.append(box.velocity) + + def as_boxes3d(self) -> rr.Boxes3D: + """Return 3D boxes data as a `rr.Boxes3D`. + + Returns: + `rr.Boxes3D` object. + """ + labels = None if len(self._uuids) == 0 else self._uuids + return rr.Boxes3D( + sizes=self._sizes, + centers=self._centers, + rotations=self._rotations, + labels=labels, + class_ids=self._class_ids, + ) + + def as_arrows3d(self) -> rr.Arrows3D: + """Return velocities data as a `rr.Arrows3D`. + + Returns: + `rr.Arrows3D` object. + """ + return rr.Arrows3D( + vectors=self._velocities, + origins=self._centers, + class_ids=self._class_ids, + ) + + +class BoxData2D: + """A class to store 2D boxes data for rendering.""" + + def __init__(self) -> None: + self._rois: list[RoiType] = [] + self._uuids: list[str] = [] + self._class_ids: list[int] = [] + + def append(self, box: Box2D) -> None: + """Append a 2D box data. + + Args: + box (Box2D): `Box2D` object. + """ + self._rois.append(box.roi.roi) + + self._class_ids.append(box.semantic_label.label.value) + + if box.uuid is not None: + self._uuids.append(box.uuid) + + def as_boxes2d(self) -> rr.Boxes2D: + """Return 2D boxes data as a `rr.Boxes2D`. + + Returns: + `rr.Boxes2D` object. + """ + labels = None if len(self._uuids) == 0 else self._uuids + return rr.Boxes2D( + array=self._rois, + array_format=rr.Box2DFormat.XYXY, + labels=labels, + class_ids=self._class_ids, + ) diff --git a/t4_devkit/viewer/viewer.py b/t4_devkit/viewer/viewer.py new file mode 100644 index 0000000..21ebb08 --- /dev/null +++ b/t4_devkit/viewer/viewer.py @@ -0,0 +1,288 @@ +from __future__ import annotations + +import os.path as osp +import warnings +from functools import singledispatchmethod +from typing import TYPE_CHECKING, Sequence + +import numpy as np +import rerun as rr +import rerun.blueprint as rrb + +from t4_devkit.common.timestamp import us2sec +from t4_devkit.schema import CalibratedSensor, EgoPose, Sensor, SensorModality +from t4_devkit.typing import CamIntrinsicType, NDArrayU8, RotationType, TranslationType + +from .color import distance_color +from .rendering_data import BoxData2D, BoxData3D + +if TYPE_CHECKING: + from t4_devkit.dataclass import Box2D, Box3D, PointCloud + + +__all__ = ["Tier4Viewer"] + + +class Tier4Viewer: + # entity paths + map_entity = "map" + ego_entity = "map/base_link" + timeline = "timestamp" + + def __init__( + self, + app_id: str, + cameras: Sequence[str] | None = None, + *, + without_3d: bool = False, + spawn: bool = True, + ) -> None: + """Construct a new object. + + Args: + app_id (str): Application ID. + cameras (Sequence[str] | None, optional): Sequence of camera names. + without_3d (bool, optional): Whether to render objects without the 3D space. + spawn (bool, optional): Whether to spawn the viewer. + """ + self.app_id = app_id + self.without_3d = without_3d + self.cameras = cameras + + view_container = [] + if not without_3d: + view_container.append( + rrb.Horizontal( + rrb.Spatial3DView(name="3D", origin=self.map_entity), + column_shares=[3, 1], + ) + ) + + if self.cameras is not None: + camera_space_views = [ + rrb.Spatial2DView(name=name, origin=osp.join(self.ego_entity, name)) + for name in self.cameras + ] + view_container.append(rrb.Grid(*camera_space_views)) + + self.blueprint = rrb.Vertical(*view_container, row_shares=[4, 2]) + + rr.init( + application_id=self.app_id, + recording_id=None, + spawn=spawn, + default_enabled=True, + strict=True, + default_blueprint=self.blueprint, + ) + + def save(self, save_dir: str) -> None: + """Save recording result as `{app_id}.rrd`. + + Args: + save_dir (str): Directory path to save the result. + """ + filepath = osp.join(save_dir, f"{self.app_id}.rrd") + rr.save(filepath, default_blueprint=self.blueprint) + + def render_box3ds(self, seconds: float, boxes: Sequence[Box3D]) -> None: + """Render 3D boxes. + + Args: + seconds (float): Timestamp in [sec]. + boxes (Sequence[Box3D]): Sequence of `Box3D`s. + """ + rr.set_time_seconds(self.timeline, seconds) + + box_data: dict[str, BoxData3D] = {} + for box in boxes: + if box.frame_id not in box_data: + box_data[box.frame_id] = BoxData3D() + else: + box_data[box.frame_id].append(box) + + for frame_id, data in box_data.items(): + # record boxes 3d + rr.log( + osp.join(self.map_entity, frame_id, "box"), + data.as_boxes3d(), + ) + # record velocities + rr.log( + osp.join(self.map_entity, frame_id, "velocity"), + data.as_arrows3d(), + ) + + def render_box2ds(self, seconds: float, boxes: Sequence[Box2D]) -> None: + """Render 2D boxes. + + Args: + seconds (float): Timestamp in [sec]. + boxes (Sequence[Box2D]): Sequence of `Box2D`s. + """ + if self.cameras is None: + warnings.warn("There is no camera space.") + return + + rr.set_time_seconds(self.timeline, seconds) + + box_data: dict[str, BoxData2D] = {} + for box in boxes: + if box.frame_id not in box_data: + box_data[box.frame_id] = BoxData2D() + else: + box_data[box.frame_id].append(box) + + for frame_id, data in box_data.items(): + rr.log( + osp.join(self.ego_entity, frame_id, "box"), + data.as_boxes2d(), + ) + + def render_pointcloud(self, seconds: float, channel: str, pointcloud: PointCloud) -> None: + """Render pointcloud. + + Args: + seconds (float): Timestamp in [sec]. + channel (str): Name of the pointcloud sensor channel. + pointcloud (PointCloud): Inherence object of `PointCloud`. + """ + rr.set_time_seconds(self.timeline, seconds) + + colors = distance_color(np.linalg.norm(pointcloud.points[:3].T, axis=1)) + rr.log( + osp.join(self.ego_entity, channel), + rr.Points3D( + pointcloud.points[:3].T, + colors=colors, + ), + ) + + def render_image(self, seconds: float, camera: str, image: str | NDArrayU8) -> None: + """Render an image. + + Args: + seconds (float): Timestamp in [sec]. + camera (str): Name of the camera channel. + image (str | NDArrayU8): Image tensor or path of the image file. + """ + if self.cameras is None or camera not in self.cameras: + warnings.warn(f"There is no camera space: {camera}") + return + + rr.set_time_seconds(self.timeline, seconds) + + if isinstance(image, str): + rr.log(osp.join(self.ego_entity, camera), rr.ImageEncoded(image)) + else: + rr.log(osp.join(self.ego_entity, camera), rr.Image(image)) + + @singledispatchmethod + def render_ego(self, *args, **kwargs) -> None: + raise TypeError("Unexpected parameter types.") + + @render_ego.register + def _render_ego_with_schema(self, ego_pose: EgoPose) -> None: + """Render an ego pose. + + Args: + ego_pose (EgoPose): `EgoPose` object. + """ + rr.set_time_seconds(self.timeline, us2sec(ego_pose.timestamp)) + + rotation_xyzw = np.roll(ego_pose.rotation.q, shift=-1) + rr.log( + self.ego_entity, + rr.Transform3D( + translation=ego_pose.translation, + rotation=rr.Quaternion(xyzw=rotation_xyzw), + from_parent=False, + ), + ) + + @render_ego.register + def _render_ego_without_schema( + self, + seconds: float, + translation: TranslationType, + rotation: RotationType, + ) -> None: + rr.set_time_seconds(self.timeline, seconds) + + rotation_xyzw = np.roll(rotation.q, shift=-1) + rr.log( + self.ego_entity, + rr.Transform3D( + translation=translation, + rotation=rr.Quaternion(xyzw=rotation_xyzw), + from_parent=False, + ), + ) + + @singledispatchmethod + def render_calibration(self, *args, **kwargs) -> None: + raise TypeError("Unexpected parameter types.") + + @render_calibration.register + def _render_calibration_with_schema( + self, + sensor: Sensor, + calibration: CalibratedSensor, + ) -> None: + """Render a sensor calibration. + + Args: + sensor (Sensor): `Sensor` object. + calibration (CalibratedSensor): `CalibratedSensor` object. + """ + rotation_xyzw = np.roll(calibration.rotation.q, shift=-1) + + rr.log( + osp.join(self.ego_entity, sensor.channel), + rr.Transform3D( + translation=calibration.translation, + rotation=rr.Quaternion(xyzw=rotation_xyzw), + ), + static=True, + ) + + if sensor.modality == SensorModality.CAMERA: + rr.log( + osp.join(self.ego_entity, sensor.channel), + rr.Pinhole(image_from_camera=calibration.camera_intrinsic), + static=True, + ) + + @render_calibration.register + def _render_calibration_without_schema( + self, + channel: str, + modality: str | SensorModality, + translation: TranslationType, + rotation: RotationType, + camera_intrinsic: CamIntrinsicType | None = None, + ) -> None: + """Render a sensor calibration. + + Args: + channel (str): Name of the sensor channel. + modality (str | SensorModality): Sensor modality. + translation (TranslationType): Sensor translation in ego centric coords. + rotation (RotationType): Sensor rotation in ego centric coords. + camera_intrinsic (CamIntrinsicType | None, optional): Camera intrinsic matrix. + Defaults to None. + """ + rotation_xyzw = np.roll(rotation.q, shift=-1) + + rr.log( + osp.join(self.ego_entity, channel), + rr.Transform3D(translation=translation, rotation=rr.Quaternion(xyzw=rotation_xyzw)), + static=True, + ) + + if modality == SensorModality.CAMERA: + rr.log( + osp.join(self.ego_entity, channel), + rr.Pinhole(image_from_camera=camera_intrinsic), + static=True, + ) diff --git a/tests/viewer/test_viewer.py b/tests/viewer/test_viewer.py new file mode 100644 index 0000000..0f280ec --- /dev/null +++ b/tests/viewer/test_viewer.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +import numpy as np +from pyquaternion import Quaternion + +from t4_devkit.dataclass import LidarPointCloud +from t4_devkit.schema import CalibratedSensor, EgoPose, Sensor +from t4_devkit.viewer import Tier4Viewer + + +def test_tier4_viewer(dummy_box3ds, dummy_box2ds) -> None: + """Test `Tier4Viewer` class. + + Args: + dummy_box3ds (list[Box3D]): List of `Box3D`s. + dummy_box2ds (list[Box2D]): List of `Box2D`s. + """ + viewer = Tier4Viewer("test_viewer", cameras=["camera"], spawn=False) + + seconds = 1.0 # [sec] + + # test render_box3ds(...) + viewer.render_box3ds(seconds, dummy_box3ds) + + # test render_box2ds(...) + viewer.render_box2ds(seconds, dummy_box2ds) + + # test render_pointcloud(...) + dummy_pointcloud = LidarPointCloud(np.random.rand(4, 100)) + viewer.render_pointcloud(seconds, "lidar", dummy_pointcloud) + + # test render_ego(...) + # without `EgoPose` + ego_translation = [1, 0, 0] + ego_rotation = Quaternion([0, 0, 0, 1]) + viewer.render_ego(seconds, ego_translation, ego_rotation) + + # with `EgoPose` + ego_pose = EgoPose( + token="ego", + translation=ego_translation, + rotation=ego_rotation, + timestamp=1e6, + ) + viewer.render_ego(ego_pose) + + # test render_calibration(...) + # without `Sensor` and `CalibratedSensor` + channel = "camera" + modality = "camera" + camera_translation = [1, 0, 0] + camera_rotation = Quaternion([0, 0, 0, 1]) + camera_intrinsic = [ + [1000.0, 0.0, 100.0], + [0.0, 1000.0, 100.0], + [0.0, 0.0, 1.0], + ] + viewer.render_calibration( + channel, + modality, + camera_translation, + camera_rotation, + camera_intrinsic, + ) + + # with `Sensor` and `CalibratedSensor` + sensor = Sensor(token="sensor", channel="camera", modality="camera") + calibration = CalibratedSensor( + token="sensor_calibration", + sensor_token="sensor", + translation=camera_translation, + rotation=camera_rotation, + camera_intrinsic=camera_intrinsic, + camera_distortion=[0, 0, 0, 0, 0], + ) + viewer.render_calibration(sensor, calibration)