feat: add support of rendering point cloud coloring with intensity (#189)

* feat: add support of rendering point cloud coloring with intensity

Signed-off-by: ktro2828 <kotaro.uetake@tier4.jp>

* feat: add support of colorizing points on image by intensity

Signed-off-by: ktro2828 <kotaro.uetake@tier4.jp>

* feat: set default value of PointCloudColorMode to DISTANCE

Signed-off-by: ktro2828 <kotaro.uetake@tier4.jp>

* refactor: modify normalize_color to private function

Signed-off-by: ktro2828 <kotaro.uetake@tier4.jp>

* Update t4_devkit/viewer/color.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update t4_devkit/helper/rendering.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

---------

Signed-off-by: ktro2828 <kotaro.uetake@tier4.jp>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: ktro2828

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,19 +451,20 @@ 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))
 
                 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,

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)

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),