feat(rosbag2_to_non_annotated_t4): improve point cloud conversion tim…

…e. (#105) * avoid using pointcloud_msg_to_numpy when checking the point num Signed-off-by: Shunsuke Miura <[email protected]> * add ros2-numpy, update lib versions Signed-off-by: Shunsuke Miura <[email protected]> * use ros2-numpy to convert pointcloud msg Signed-off-by: Shunsuke Miura <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix error Signed-off-by: Shunsuke Miura <[email protected]> * increment the version number Signed-off-by: Shunsuke Miura <[email protected]> * fix a bug Signed-off-by: Shunsuke Miura <[email protected]> --------- Signed-off-by: Shunsuke Miura <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
tier4 · May 23, 2024 · 22844d2 · 22844d2
1 parent e9c4e89
commit 22844d2
Show file tree

Hide file tree

Showing 4 changed files with 608 additions and 437 deletions.
diff --git a/perception_dataset/rosbag2/rosbag2_to_non_annotated_t4_converter.py b/perception_dataset/rosbag2/rosbag2_to_non_annotated_t4_converter.py
@@ -462,18 +462,21 @@ def _convert_pointcloud(
                         f"PointCloud message is dropped [{frame_index}]: cur={unix_timestamp} prev={prev_frame_unix_timestamp}"
                     )
 
-            points_arr = rosbag2_utils.pointcloud_msg_to_numpy(pointcloud_msg)
-            if len(points_arr) < max_num_points * self._lidar_points_ratio_threshold:
+            if hasattr(pointcloud_msg, "width"):
+                num_points = pointcloud_msg.width
+            else:
+                num_points = 0
+            if num_points < max_num_points * self._lidar_points_ratio_threshold:
                 if not self._accept_frame_drop:
                     raise ValueError(
                         f"PointCloud message is relatively lower than the maximum size, which is not acceptable. "
                         f"If you would like to accept, please change accept_frame_drop parameter. "
-                        f"frame_index: {frame_index}, stamp: {unix_timestamp}, # points: {len(points_arr)}"
+                        f"frame_index: {frame_index}, stamp: {unix_timestamp}, # points: {num_points}"
                     )
                 else:
                     warnings.warn(
                         f"PointCloud message is relatively lower than the maximum size. "
-                        f"May be encountering a LiDAR message drop. Skip frame_index: {frame_index}, stamp: {unix_timestamp}, # points: {len(points_arr)}"
+                        f"May be encountering a LiDAR message drop. Skip frame_index: {frame_index}, stamp: {unix_timestamp}, # points: {num_points}"
                     )
                     continue
 

diff --git a/perception_dataset/utils/rosbag2.py b/perception_dataset/utils/rosbag2.py
@@ -11,6 +11,7 @@
 import numpy as np
 from radar_msgs.msg import RadarTrack, RadarTracks
 from rclpy.time import Time
+import ros2_numpy as rnp
 from rosbag2_py import (
     ConverterOptions,
     Reindexer,
@@ -19,7 +20,6 @@
     StorageOptions,
 )
 from sensor_msgs.msg import CompressedImage, PointCloud2
-import sensor_msgs_py.point_cloud2
 import yaml
 
 from perception_dataset.utils.misc import unix_timestamp_to_nusc_timestamp
@@ -102,23 +102,33 @@ def get_default_storage_options(bag_dir: str) -> StorageOptions:
     return StorageOptions(uri=bag_dir, storage_id=storage_id)
 
 
-def pointcloud_msg_to_numpy(
-    pointcloud_msg: PointCloud2,
-) -> NDArray:
-    """numpy ver. of https://github.com/ros2/common_interfaces/blob/master/sensor_msgs_py/sensor_msgs_py/point_cloud2.py#L119"""
+def pointcloud_msg_to_numpy(pointcloud_msg: PointCloud2) -> NDArray:
+    """Convert ROS PointCloud2 message to numpy array using ros2-numpy."""
     NUM_DIMENSIONS = 5
 
     if not isinstance(pointcloud_msg, PointCloud2):
         return np.zeros((0, NUM_DIMENSIONS), dtype=np.float32)
 
-    points_arr = np.array(
-        [tuple(p) for p in sensor_msgs_py.point_cloud2.read_points(pointcloud_msg)],
-        dtype=np.float32,
-    )
-    if len(points_arr[0]) > NUM_DIMENSIONS:
-        points_arr = np.delete(points_arr, np.s_[NUM_DIMENSIONS:], axis=1)
-    while len(points_arr[0]) < NUM_DIMENSIONS:
-        points_arr = np.insert(points_arr, len(points_arr[0]), -1, axis=1)
+    # Convert the PointCloud2 message to a numpy structured array
+    points = rnp.point_cloud2.point_cloud2_to_array(pointcloud_msg)
+
+    # Extract the x, y, z coordinates and additional fields if available
+    xyz = points["xyz"]
+    if "intensity" in points.keys():
+        intensity = points["intensity"]
+        points_arr = np.hstack((xyz, intensity))
+    else:
+        points_arr = xyz
+
+    # Ensure the resulting array has exactly NUM_DIMENSIONS columns
+    if points_arr.shape[1] > NUM_DIMENSIONS:
+        points_arr = points_arr[:, :NUM_DIMENSIONS]
+    elif points_arr.shape[1] < NUM_DIMENSIONS:
+        padding = np.full(
+            (points_arr.shape[0], NUM_DIMENSIONS - points_arr.shape[1]), -1, dtype=np.float32
+        )
+        points_arr = np.hstack((points_arr, padding))
+
     return points_arr
 
 
@@ -167,7 +177,7 @@ def radar_tracks_msg_to_list(radar_tracks_msg: RadarTracks) -> List[Dict[str, An
     return radar_tracks
 
 
-def compressed_msg_to_numpy(compressed_image_msg: CompressedImage) -> np.ndarray:
+def compressed_msg_to_numpy(compressed_image_msg: CompressedImage) -> NDArray:
     if hasattr(compressed_image_msg, "_encoding"):
         try:
             np_arr = np.frombuffer(compressed_image_msg.data, np.uint8)