feat: add support of colorizing points on image by intensity

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

Author: ktro2828

t4_devkit/helper/rendering.py

@@ -16,7 +16,7 @@
 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, distance_color, format_entity
+from t4_devkit.viewer import PointCloudColorMode, RerunViewer, format_entity, normalize_color
 
 if TYPE_CHECKING:
     from t4_devkit import Tier4
@@ -468,7 +468,7 @@ 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,
@@ -480,7 +480,7 @@ def _render_points_on_single_camera(camera: str) -> None:
                 rr.log(format_entity(RerunViewer.ego_entity, camera), rr.Image(image))
                 rr.log(
                     format_entity(RerunViewer.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
@@ -498,6 +498,7 @@ def _project_pointcloud(
         min_dist: float = 1.0,
         *,
         ignore_distortion: bool = True,
+        color_mode: PointCloudColorMode = PointCloudColorMode.INTENSITY,
     ) -> tuple[NDArrayF64, NDArrayF64, NDArrayU8]:
         """Project pointcloud on image plane.
 
@@ -506,9 +507,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 normalized features [n] 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)
@@ -549,8 +551,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(
@@ -560,16 +560,24 @@ def _project_pointcloud(
             normalize=True,
         )[:2]
 
-        mask = np.ones(depths.shape[0], dtype=bool)
-        mask = np.logical_and(mask, depths > min_dist)
+        match color_mode:
+            case PointCloudColorMode.DISTANCE:
+                colors = normalize_color(pointcloud.points[2, :])
+            case _:
+                colors = normalize_color(pointcloud.points[3, :])
+        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

@@ -11,6 +11,9 @@
     from t4_devkit.typing import ArrayLike, NDArrayF64
 
 
+__all__ = ["PointCloudColorMode", "pointcloud_color", "normalize_color"]
+
+
 @unique
 class PointCloudColorMode(str, Enum):
     """Color mode of point cloud."""
@@ -35,25 +38,22 @@ def pointcloud_color(
         case _:
             values = pointcloud.points[3]
 
-    return distance_color(values)
+    return normalize_color(values)
 
 
-def distance_color(
-    distances: ArrayLike,
-    cmap: str | None = None,
-) -> tuple[float, float, float] | NDArrayF64:
-    """Return color map depending on distance values.
+def normalize_color(values: ArrayLike, cmap: str | None = None, alpha: float = 1.0) -> NDArrayF64:
+    """Return color map normalizing values.
 
     Args:
-        distances (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.
+        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(distances)
-    v_max = np.max(distances)
+    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)