diff --git a/t4_devkit/helper/rendering.py b/t4_devkit/helper/rendering.py index 6f15bcc..b2df55b 100644 --- a/t4_devkit/helper/rendering.py +++ b/t4_devkit/helper/rendering.py @@ -15,7 +15,14 @@ from t4_devkit.common.timestamp import sec2us, us2sec from t4_devkit.dataclass import LidarPointCloud, RadarPointCloud from t4_devkit.schema import SensorModality -from t4_devkit.viewer import RerunViewer, ViewerBuilder, ViewerConfig, distance_color, format_entity +from t4_devkit.viewer import ( + PointCloudColorMode, + RerunViewer, + ViewerBuilder, + ViewerConfig, + format_entity, + pointcloud_color, +) if TYPE_CHECKING: from t4_devkit import Tier4 @@ -332,6 +339,7 @@ def _render_lidar_and_ego( viewer: RerunViewer, first_lidar_tokens: list[str], max_timestamp_us: float, + color_mode: PointCloudColorMode = PointCloudColorMode.DISTANCE, ) -> list[Future]: def _render_single_lidar(first_lidar_token: str) -> None: self._render_sensor_calibration(viewer=viewer, sample_data_token=first_lidar_token) @@ -353,6 +361,7 @@ def _render_single_lidar(first_lidar_token: str) -> None: seconds=us2sec(sample_data.timestamp), channel=sample_data.channel, pointcloud=pointcloud, + color_mode=color_mode, ) current_lidar_token = sample_data.next @@ -423,6 +432,7 @@ def _render_points_on_cameras( *, min_dist: float = 1.0, ignore_distortion: bool = True, + color_mode: PointCloudColorMode = PointCloudColorMode.DISTANCE, ) -> list[Future]: def _render_points_on_single_camera(camera: str) -> None: current_point_sample_data_token = first_point_sample_data_token @@ -441,11 +451,12 @@ def _render_points_on_single_camera(camera: str) -> None: if max_timestamp_us < sample.timestamp: break - points_on_image, depths, image = self._project_pointcloud( + points_on_image, colors, image = self._project_pointcloud( point_sample_data_token=current_point_sample_data_token, camera_sample_data_token=camera_sample_data_token, min_dist=min_dist, ignore_distortion=ignore_distortion, + color_mode=color_mode, ) rr.set_time_seconds(ViewerConfig.timeline, us2sec(sample.timestamp)) @@ -453,7 +464,7 @@ def _render_points_on_single_camera(camera: str) -> None: rr.log(format_entity(ViewerConfig.ego_entity, camera), rr.Image(image)) rr.log( format_entity(ViewerConfig.ego_entity, camera, "pointcloud"), - rr.Points2D(positions=points_on_image.T, colors=distance_color(depths)), + rr.Points2D(positions=points_on_image.T, colors=colors), ) current_point_sample_data_token = sample_data.next @@ -471,6 +482,7 @@ def _project_pointcloud( min_dist: float = 1.0, *, ignore_distortion: bool = True, + color_mode: PointCloudColorMode = PointCloudColorMode.DISTANCE, ) -> tuple[NDArrayF64, NDArrayF64, NDArrayU8]: """Project pointcloud on image plane. @@ -479,9 +491,10 @@ def _project_pointcloud( camera_sample_data_token (str): Sample data token of camera. min_dist (float, optional): Distance from the camera below which points are discarded. ignore_distortion (bool, optional): Whether to ignore distortion parameters. + color_mode (PointCloudColorMode, optional): Color mode for pointcloud. Returns: - Projected points [2, n], their normalized depths [n] and an image. + Projected points [2, n], their color values [n, 3], and an image. """ point_sample_data: SampleData = self._t4.get("sample_data", point_sample_data_token) pc_filepath = osp.join(self._t4.data_root, point_sample_data.filename) @@ -522,8 +535,6 @@ def _project_pointcloud( pointcloud.translate(-camera_cs_record.translation) pointcloud.rotate(camera_cs_record.rotation.rotation_matrix.T) - depths = pointcloud.points[2, :] - distortion = None if ignore_distortion else camera_cs_record.camera_distortion points_on_img = view_points( @@ -533,16 +544,20 @@ def _project_pointcloud( normalize=True, )[:2] - mask = np.ones(depths.shape[0], dtype=bool) - mask = np.logical_and(mask, depths > min_dist) + colors = pointcloud_color(pointcloud, color_mode) + depths = pointcloud.points[2, :] + + mask = np.ones(colors.shape[0], dtype=bool) + mask = np.logical_and(mask, depths > min_dist) # mask by depths + # mask by size of points on image mask = np.logical_and(mask, 1 < points_on_img[0]) mask = np.logical_and(mask, points_on_img[0] < img.size[0] - 1) mask = np.logical_and(mask, 1 < points_on_img[1]) mask = np.logical_and(mask, points_on_img[1] < img.size[1] - 1) points_on_img = points_on_img[:, mask] - depths = depths[mask] + colors = colors[mask] - return points_on_img, depths, np.array(img, dtype=np.uint8) + return points_on_img, colors, np.array(img, dtype=np.uint8) def _render_annotation3ds( self, diff --git a/t4_devkit/viewer/color.py b/t4_devkit/viewer/color.py index 6ebbf1b..ca30954 100644 --- a/t4_devkit/viewer/color.py +++ b/t4_devkit/viewer/color.py @@ -1,31 +1,61 @@ from __future__ import annotations +from enum import Enum, unique from typing import TYPE_CHECKING import matplotlib +import numpy as np if TYPE_CHECKING: - from t4_devkit.typing import ArrayLike, NDArrayF64, ScalarLike + from t4_devkit.dataclass import PointCloudLike + from t4_devkit.typing import ArrayLike, NDArrayF64 -def distance_color( - distances: ScalarLike | ArrayLike, - cmap: str | None = None, - v_min: float = 3.0, - v_max: float = 75.0, -) -> tuple[float, float, float] | NDArrayF64: - """Return color map depending on distance values. +__all__ = ["PointCloudColorMode", "pointcloud_color"] + + +@unique +class PointCloudColorMode(str, Enum): + """Color mode of point cloud.""" + + DISTANCE = "distance" + INTENSITY = "intensity" + + +def pointcloud_color( + pointcloud: PointCloudLike, + color_mode: PointCloudColorMode = PointCloudColorMode.DISTANCE, +) -> NDArrayF64: + """Return color map depending on the specified color mode. + + Args: + pointcloud (PointCloudLike): Any inheritance of `PointCloud` class. + color_mode (PointCloudColorMode, optional): Color mode for pointcloud. + """ + match color_mode: + case PointCloudColorMode.DISTANCE: + values = np.linalg.norm(pointcloud.points[:3].T, axis=1) + case PointCloudColorMode.INTENSITY: + values = pointcloud.points[3] + case _: + raise ValueError(f"Unsupported color mode: {color_mode}") + + return _normalize_color(values) + + +def _normalize_color(values: ArrayLike, cmap: str | None = None, alpha: float = 1.0) -> NDArrayF64: + """Return color map normalizing values. Args: - distances (ScalarLike | ArrayLike): Array of distances in the shape of (N,). + values (ArrayLike): Array of values in the shape of (N,). cmap (str | None, optional): Color map name in matplotlib. If None, `turbo_r` will be used. - v_min (float, optional): Min value to normalize. - v_max (float, optional): Max value to normalize. + alpha (float, optional): Alpha value of color map. Returns: - Color map in the shape of (N,). If input type is any number, returns a color as - `tuple[float, float, float]`. Otherwise, returns colors as `NDArrayF64`. + Color map in the shape of (N,). """ color_map = matplotlib.colormaps["turbo_r"] if cmap is None else matplotlib.colormaps[cmap] + v_min = np.min(values) + v_max = np.max(values) norm = matplotlib.colors.Normalize(v_min, v_max) - return color_map(norm(distances)) + return color_map(norm(values), alpha=alpha) diff --git a/t4_devkit/viewer/viewer.py b/t4_devkit/viewer/viewer.py index e2d7e63..a1d9a94 100644 --- a/t4_devkit/viewer/viewer.py +++ b/t4_devkit/viewer/viewer.py @@ -12,7 +12,7 @@ from t4_devkit.lanelet import LaneletParser from t4_devkit.schema import SensorModality -from .color import distance_color +from .color import PointCloudColorMode, pointcloud_color from .config import ViewerConfig, format_entity from .geography import calculate_geodetic_point from .lanelet import ( @@ -408,18 +408,25 @@ def render_segmentation2d( ) @_check_spatial3d - def render_pointcloud(self, seconds: float, channel: str, pointcloud: PointCloudLike) -> None: + def render_pointcloud( + self, + seconds: float, + channel: str, + pointcloud: PointCloudLike, + color_mode: PointCloudColorMode = PointCloudColorMode.DISTANCE, + ) -> None: """Render pointcloud. Args: seconds (float): Timestamp in [sec]. channel (str): Name of the pointcloud sensor channel. pointcloud (PointCloudLike): Inherence object of `PointCloud`. + color_mode (PointCloudColorMode, optional): Color mode for pointcloud. """ # TODO(ktro2828): add support of rendering pointcloud on images rr.set_time_seconds(self.config.timeline, seconds) - colors = distance_color(np.linalg.norm(pointcloud.points[:3].T, axis=1)) + colors = pointcloud_color(pointcloud, color_mode=color_mode) rr.log( format_entity(self.config.ego_entity, channel), rr.Points3D(pointcloud.points[:3].T, colors=colors),