diff --git a/docs/tutorials/initialize.md b/docs/tutorials/initialize.md index 96822783..27737edb 100644 --- a/docs/tutorials/initialize.md +++ b/docs/tutorials/initialize.md @@ -19,6 +19,7 @@ │   ├── CAM_FRONT_RIGHT │   ├── LIDAR_CONCAT │   └── ...Other sensor channels + ├── input_bag (optional) ...rosbag2 files ... ``` @@ -36,6 +37,7 @@ │   ├── CAM_FRONT_RIGHT │   ├── LIDAR_CONCAT │   └── ...Other sensor channels + ├── input_bag (optional) ...rosbag2 files ... ``` diff --git a/docs/tutorials/render.md b/docs/tutorials/render.md index c6286115..afb72700 100644 --- a/docs/tutorials/render.md +++ b/docs/tutorials/render.md @@ -40,6 +40,51 @@ If `SampleData.info_filename` points to a pointcloud metainfo JSON file, `T4Devk ![Render PointCloud GIF](../assets/render_pointcloud.gif) +#### Rosbag Support + +If the dataset contains an `input_bag/` directory with rosbag2 files (db3 or mcap format), you can read LiDAR point clouds directly from the rosbag instead of the processed `.pcd.bin` files. + +Both `sensor_msgs/msg/PointCloud2` and `pandar_msgs/msg/PandarScan` (Hesai raw packet) topics are supported. PandarScan packets are decoded to point clouds using the specified sensor model's elevation angles. + +The `topic_mapping` parameter maps T4 dataset sensor channel names (e.g. `LIDAR_TOP`, `LIDAR_CONCAT`) to ROS topic names in the rosbag. This mapping is required so that `get_lidar_pointcloud()` can look up the correct rosbag topic for each `sample_data.channel`. + +For `PointCloud2` topics, a simple dict mapping is sufficient: + +```python +>>> from t4_devkit import T4Devkit + +>>> t4 = T4Devkit( +... "data/tier4/", +... use_rosbag=True, +... topic_mapping={"LIDAR_CONCAT": "/sensing/lidar/concatenated/pointcloud"}, +... ) +``` + +For `PandarScan` topics, use `TopicMapping` with `sensor_type` to specify the Hesai sensor model (`"XT32"` or `"OT128"`). Use `frame_id` to specify the TF frame so that the decoded point cloud is automatically transformed to `base_link`: + +```python +>>> from t4_devkit import T4Devkit +>>> from t4_devkit.rosbag.topic_mapping import TopicMapping + +>>> t4 = T4Devkit( +... "data/tier4/", +... use_rosbag=True, +... topic_mapping=[ +... TopicMapping( +... channel="LIDAR_CONCAT", +... topic="/sensing/lidar/top/pandar_packets", +... sensor_type="OT128", +... frame_id="hesai_top", +... ), +... ], +... ) + +>>> pc = t4.get_lidar_pointcloud(sample_data_token) +>>> t4.render_pointcloud() +``` + +If `topic_mapping` is omitted, `PointCloud2` topics are auto-detected from the rosbag. `PandarScan` topics always require explicit `topic_mapping` with `sensor_type`. The `frame_id` parameter is optional but recommended — without it, points remain in the sensor's native coordinate frame. + ### Save Recording You can save the rendering result as follows: diff --git a/pyproject.toml b/pyproject.toml index bcb9bc8a..65c09841 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,6 +22,7 @@ dependencies = [ "tqdm>=4.67.1", "returns>=0.26.0", "pyarrow<24.0.0", # NOTE: pyarrow==24.0.0 is only supported on macOS + "rosbags>=0.10.0", ] [dependency-groups] diff --git a/t4_devkit/cli/visualize.py b/t4_devkit/cli/visualize.py index a7c285df..ea9c4f6e 100644 --- a/t4_devkit/cli/visualize.py +++ b/t4_devkit/cli/visualize.py @@ -1,5 +1,6 @@ from __future__ import annotations +import json import os from typing import Annotated @@ -140,6 +141,41 @@ def pointcloud( t4.render_pointcloud(save_dir=output, show_map=show_map) +def _parse_topic_mapping(topic_mapping: str | None) -> list | None: + """Parse topic mapping from a JSON string or file path. + + Args: + topic_mapping (str | None): JSON string or path to a JSON file. + + Returns: + List of TopicMapping instances, or None. + """ + if topic_mapping is None: + return None + + try: + if os.path.isfile(topic_mapping): + from t4_devkit.common.io import load_json + + raw = load_json(topic_mapping) + else: + raw = json.loads(topic_mapping) + except Exception as e: + raise typer.BadParameter( + "--topic-mapping must be a JSON object (or a path to a JSON file) mapping channel -> topic" + ) from e + + if not isinstance(raw, dict): + raise typer.BadParameter("--topic-mapping must be a JSON object mapping channel -> topic") + + from t4_devkit.rosbag import TopicMapping + + try: + return TopicMapping.from_dict(raw) + except (TypeError, ValueError) as e: + raise typer.BadParameter(str(e)) from e + + def _create_dir(dir_path: str | None) -> None: """Create a directory with the specified path. diff --git a/t4_devkit/helper/rendering.py b/t4_devkit/helper/rendering.py index ec3147d0..95d0354c 100644 --- a/t4_devkit/helper/rendering.py +++ b/t4_devkit/helper/rendering.py @@ -12,7 +12,7 @@ import yaml from t4_devkit.common.timestamp import microseconds2seconds, seconds2microseconds -from t4_devkit.dataclass import LidarPointCloud, RadarPointCloud, SegmentationPointCloud +from t4_devkit.dataclass import RadarPointCloud, SegmentationPointCloud from t4_devkit.schema import SensorModality from t4_devkit.viewer import ( EntityPath, @@ -460,8 +460,8 @@ def _render_single_lidar(first_lidar_token: str) -> None: metainfo_filepath=metainfo_filepath, ) else: - pointcloud = LidarPointCloud.from_file( - osp.join(self._t4.data_root, sample_data.filename), + pointcloud = self._t4.get_lidar_pointcloud( + sample_data.token, metainfo_filepath=metainfo_filepath, ) diff --git a/t4_devkit/rosbag/__init__.py b/t4_devkit/rosbag/__init__.py new file mode 100644 index 00000000..e4679129 --- /dev/null +++ b/t4_devkit/rosbag/__init__.py @@ -0,0 +1,6 @@ +from __future__ import annotations + +from .reader import Rosbag2Reader +from .topic_mapping import TopicMapping + +__all__ = ["Rosbag2Reader", "TopicMapping"] diff --git a/t4_devkit/rosbag/pandar_decoder.py b/t4_devkit/rosbag/pandar_decoder.py new file mode 100644 index 00000000..4566e763 --- /dev/null +++ b/t4_devkit/rosbag/pandar_decoder.py @@ -0,0 +1,387 @@ +"""Decoder for Hesai PandarScan messages to LidarPointCloud. + +Supports decoding raw Hesai LiDAR UDP packets from ``pandar_msgs/msg/PandarScan`` +messages into t4-devkit ``LidarPointCloud`` format. + +Supported models: +- PandarXT-32 (``XT32``) +- OT128 (``OT128``) +""" + +from __future__ import annotations + +import logging +import struct +from dataclasses import dataclass, field + +import numpy as np +from rosbags.interfaces import Nodetype + +from t4_devkit.dataclass import LidarPointCloud + +__all__ = ["HESAI_MODELS", "pandarscan_to_lidar", "register_pandar_types"] + +logger = logging.getLogger(__name__) + +# pandar_msgs type definitions for rosbags typestore registration. +# PandarPacket uses a fixed uint8[1500] data buffer with a uint32 size field +# indicating the actual number of valid bytes. +_PANDAR_DATA_BUFFER_SIZE = 1500 + +PANDAR_FIELDDEFS: dict = { + "pandar_msgs/msg/PandarPacket": ( + [], + [ + ("stamp", (Nodetype.NAME, "builtin_interfaces/msg/Time")), + ( + "data", + (Nodetype.ARRAY, ((Nodetype.BASE, ("uint8", 0)), _PANDAR_DATA_BUFFER_SIZE)), + ), + ("size", (Nodetype.BASE, ("uint32", 0))), + ], + ), + "pandar_msgs/msg/PandarScan": ( + [], + [ + ("header", (Nodetype.NAME, "std_msgs/msg/Header")), + ( + "packets", + (Nodetype.SEQUENCE, ((Nodetype.NAME, "pandar_msgs/msg/PandarPacket"), 0)), + ), + ], + ), +} + +# Hesai packet magic bytes +_HESAI_SOF = (0xEE, 0xFF) + +# Pre-header + header size when SOF is present +_PRE_HEADER_SIZE = 6 +_HEADER_SIZE = 6 + +# Header field offsets (relative to header start, i.e. after pre-header) +_HEADER_LASER_NUM = 0 +_HEADER_BLOCK_NUM = 1 +_HEADER_DIS_UNIT = 3 + +_AZIMUTH_BYTES = 2 + +_EMPTY_POINTS = np.zeros((4, 0), dtype=np.float32) + + +@dataclass(frozen=True) +class _HesaiModelConfig: + """Configuration for a Hesai LiDAR model.""" + + name: str + elevation_deg: list[float] = field(repr=False) + azimuth_offset_deg: list[float] = field(repr=False) + udp_seq_offset_from_end: int = 4 + """Byte offset from the end of the raw packet to the start of the 4-byte + UDP Sequence field (u32 LE). XT32: 4 (last 4 bytes, in "Additional + information"). OT128: 30 (in Tail, followed by 26 bytes of IMU data).""" + cos_el: np.ndarray = field(init=False, repr=False, compare=False) + sin_el: np.ndarray = field(init=False, repr=False, compare=False) + azimuth_offset_rad: np.ndarray = field(init=False, repr=False, compare=False) + + def __post_init__(self) -> None: + el_rad = np.radians(np.array(self.elevation_deg, dtype=np.float32)) + object.__setattr__(self, "cos_el", np.cos(el_rad)) + object.__setattr__(self, "sin_el", np.sin(el_rad)) + object.__setattr__( + self, + "azimuth_offset_rad", + np.radians(np.array(self.azimuth_offset_deg, dtype=np.float32)), + ) + + +# XT32: 32 channels, 1° spacing from +15° to -16° +# https://www.hesaitech.com/product/xt16-32-32m/ +_XT32_ELEVATION_DEG: list[float] = [float(15 - i) for i in range(32)] +_XT32_AZIMUTH_OFFSET_DEG: list[float] = [0.0] * 32 + +# OT128: 128 channels, non-uniform spacing from +14.985° to -24.765° +# Per Hesai OT128 User Manual (O01-en-260410), Appendix A / Angle Correction File: +# https://www.hesaitech.com/product/ot128/ +# fmt: off +_OT128_ELEVATION_DEG: list[float] = [ + 14.985, 13.283, 11.758, 10.483, 9.836, 9.171, 8.496, 7.812, + 7.462, 7.115, 6.767, 6.416, 6.064, 5.71, 5.355, 4.998, + 4.643, 4.282, 3.921, 3.558, 3.194, 2.829, 2.463, 2.095, + 1.974, 1.854, 1.729, 1.609, 1.487, 1.362, 1.242, 1.12, + 0.995, 0.875, 0.75, 0.625, 0.5, 0.375, 0.25, 0.125, + 0.0, -0.125, -0.25, -0.375, -0.5, -0.626, -0.751, -0.876, + -1.001, -1.126, -1.251, -1.377, -1.502, -1.627, -1.751, -1.876, + -2.001, -2.126, -2.251, -2.376, -2.501, -2.626, -2.751, -2.876, + -3.001, -3.126, -3.251, -3.376, -3.501, -3.626, -3.751, -3.876, + -4.001, -4.126, -4.25, -4.375, -4.501, -4.626, -4.751, -4.876, + -5.001, -5.126, -5.252, -5.377, -5.502, -5.626, -5.752, -5.877, + -6.002, -6.378, -6.754, -7.13, -7.507, -7.882, -8.257, -8.632, + -9.003, -9.376, -9.749, -10.121, -10.493, -10.864, -11.234, -11.603, + -11.975, -12.343, -12.709, -13.075, -13.439, -13.803, -14.164, -14.525, + -14.879, -15.237, -15.593, -15.948, -16.299, -16.651, -17.0, -17.347, + -17.701, -18.386, -19.063, -19.73, -20.376, -21.653, -23.044, -24.765, +] +_OT128_AZIMUTH_OFFSET_DEG: list[float] = [ + 0.186, 0.185, 1.335, 1.343, 0.148, 0.147, 0.146, 0.146, + 1.335, 1.336, 1.337, 1.338, 1.339, 1.34, 1.341, 1.342, + 0.128, 0.128, 0.127, 0.127, 0.107, 0.106, 0.105, 0.105, + -3.118, 1.315, 4.529, -3.121, 1.316, 4.532, -3.124, 1.317, + 4.536, -3.127, 1.317, 4.539, -3.13, 1.318, 4.542, -3.133, + 0.103, 2.935, -1.517, 0.103, 2.937, -1.519, 0.103, 2.939, + -1.52, 0.103, 2.941, -1.521, 0.102, 2.943, -1.523, 0.102, + 2.945, -1.524, 0.102, 2.946, -1.526, 0.102, 2.948, -1.526, + 1.324, 4.57, -3.155, 1.325, 4.573, -3.157, 1.326, 4.575, + -3.159, 1.326, 4.578, -3.161, 1.327, 4.581, -3.163, 1.328, + 4.583, -3.165, 1.329, 4.586, -3.167, 1.329, 4.588, -3.168, + 0.102, 0.103, 0.103, 0.103, 0.104, 0.104, 0.104, 0.104, + 1.337, 1.337, 1.338, 1.339, 1.34, 1.341, 1.341, 1.342, + 0.108, 0.108, 0.109, 0.109, 0.13, 0.131, 0.131, 0.132, + 1.384, 1.384, 1.385, 1.385, 1.386, 1.386, 1.387, 1.387, + 0.151, 0.153, 0.154, 0.156, 1.388, 1.408, 0.196, 0.286, +] +# fmt: on + +# Model lookup by sensor type name. +HESAI_MODELS: dict[str, _HesaiModelConfig] = { + "XT32": _HesaiModelConfig( + name="XT32", + elevation_deg=_XT32_ELEVATION_DEG, + azimuth_offset_deg=_XT32_AZIMUTH_OFFSET_DEG, + ), + "OT128": _HesaiModelConfig( + name="OT128", + elevation_deg=_OT128_ELEVATION_DEG, + azimuth_offset_deg=_OT128_AZIMUTH_OFFSET_DEG, + udp_seq_offset_from_end=30, + ), +} + + +def register_pandar_types(typestore: object) -> None: + """Register ``pandar_msgs`` message types with a rosbags typestore. + + Args: + typestore: A rosbags ``Typestore`` instance. + """ + typestore.register(PANDAR_FIELDDEFS) + + +def _extract_udp_sequence(raw: bytes, config: _HesaiModelConfig) -> int | None: + """Extract the UDP Sequence number from a raw Hesai packet. + + The UDP Sequence is a monotonically incrementing u32 counter embedded + in each UDP packet. Its byte offset from the end of the packet varies + by model (see ``_HesaiModelConfig.udp_seq_offset_from_end``). + + Args: + raw: Raw packet bytes (trimmed to actual size). + config: Model config for the sensor type. + + Returns: + UDP Sequence number (u32), or ``None`` if the packet is too short. + """ + if len(raw) < config.udp_seq_offset_from_end + 4: + return None + offset = len(raw) - config.udp_seq_offset_from_end + return struct.unpack_from(" LidarPointCloud: + """Convert a deserialized ``PandarScan`` message to ``LidarPointCloud``. + + Decodes all packets in the scan, converts spherical coordinates to + Cartesian, and returns a ``LidarPointCloud`` with shape ``(4, N)``. + + When *min_completeness* > 0, the UDP Sequence numbers embedded in each + packet are checked. If the ratio ``actual_packets / expected_packets`` + is below *min_completeness*, a :class:`ValueError` is raised to indicate + that the scan is incomplete (packets were dropped during recording). + + Args: + msg: Deserialized ``pandar_msgs/msg/PandarScan`` message. + sensor_type: Hesai sensor model name (e.g. ``"OT128"``, ``"XT32"``). + min_completeness: Minimum fraction of expected packets that must be + present (0.0–1.0). Defaults to ``1.0`` (reject any scan with + dropped packets). Set to ``0.0`` to disable the check. + + Returns: + LidarPointCloud instance. + + Raises: + ValueError: If no packets, unsupported sensor type, channel count + mismatch, or scan completeness is below *min_completeness*. + """ + config = HESAI_MODELS.get(sensor_type) + if config is None: + raise ValueError( + f"Unsupported sensor type '{sensor_type}'. " + f"Supported types: {sorted(HESAI_MODELS.keys())}" + ) + + if not msg.packets: + raise ValueError("PandarScan message contains no packets") + + # Extract raw bytes and UDP sequences for each packet. + raw_packets: list[bytes] = [] + for packet in msg.packets: + data = packet.data[: packet.size] + raw = data.tobytes() if hasattr(data, "tobytes") else bytes(data) + if len(raw) < 20: + continue + raw_packets.append(raw) + + if not raw_packets: + raise ValueError("PandarScan message contains no valid packets") + + # Check scan completeness via UDP Sequence. + # Packets in a PandarScan are assembled by the ROS driver in + # chronological (azimuth) order, so first/last is sufficient. + if min_completeness > 0.0 and len(raw_packets) >= 2: + first_seq = _extract_udp_sequence(raw_packets[0], config) + last_seq = _extract_udp_sequence(raw_packets[-1], config) + if first_seq is not None and last_seq is not None: + expected = last_seq - first_seq + 1 + actual = len(raw_packets) + if expected > 0: + completeness = actual / expected + if completeness < min_completeness: + raise ValueError( + f"Incomplete PandarScan: {actual}/{expected} packets " + f"({completeness:.1%} completeness, " + f"threshold={min_completeness:.1%}). " + f"UDP Sequence range: {first_seq}–{last_seq}." + ) + if completeness < 1.0: + dropped = expected - actual + logger.debug( + "PandarScan has %d/%d packets (%.1f%% complete, " + "%d dropped). Proceeding (above %.0f%% threshold).", + actual, + expected, + completeness * 100, + dropped, + min_completeness * 100, + ) + + point_arrays: list[np.ndarray] = [] + for raw in raw_packets: + points = _decode_packet(raw, config) + if points.shape[1] > 0: + point_arrays.append(points) + + if not point_arrays: + return LidarPointCloud(points=_EMPTY_POINTS) + + combined = np.concatenate(point_arrays, axis=1) + return LidarPointCloud(points=combined) + + +def _decode_packet( + raw: bytes, + config: _HesaiModelConfig, +) -> np.ndarray: + """Decode a single Hesai UDP packet. + + Args: + raw: Raw packet bytes (trimmed to actual size via ``PandarPacket.size``). + config: Model config for the sensor type. + + Returns: + Points array with shape ``(4, N)`` in Hesai native frame + (x=d*cos*sin, y=d*cos*cos, z=d*sin). + + Raises: + ValueError: If the packet channel count doesn't match the config. + """ + has_sof = raw[0] == _HESAI_SOF[0] and raw[1] == _HESAI_SOF[1] + + if has_sof: + header_offset = _PRE_HEADER_SIZE + body_offset = _PRE_HEADER_SIZE + _HEADER_SIZE + else: + header_offset = 0 + body_offset = _HEADER_SIZE + + laser_num = raw[header_offset + _HEADER_LASER_NUM] + block_num = raw[header_offset + _HEADER_BLOCK_NUM] + dis_unit_mm = raw[header_offset + _HEADER_DIS_UNIT] + + if dis_unit_mm == 0: + dis_unit_mm = 4 # default 4mm + + if block_num == 0 or laser_num == 0: + raise ValueError(f"Invalid packet header: blocks={block_num}, lasers={laser_num}") + + expected_channels = len(config.elevation_deg) + if laser_num != expected_channels: + raise ValueError( + f"Packet channel count ({laser_num}) does not match " + f"sensor type '{config.name}' ({expected_channels} channels)" + ) + + # Hesai firmware uses either 3-byte (distance+reflectivity) or 4-byte + # (distance+reflectivity+reserved) channel layout with no explicit field + # in the header. Infer by checking which layout leaves a valid tail. + channel_bytes = 0 + for ch_bytes in (4, 3): + block_size = _AZIMUTH_BYTES + laser_num * ch_bytes + body_size = block_num * block_size + tail = len(raw) - body_offset - body_size + if tail >= 12: + channel_bytes = ch_bytes + break + + if channel_bytes == 0: + raise ValueError( + f"Cannot determine channel layout " + f"(packet_len={len(raw)}, blocks={block_num}, lasers={laser_num})" + ) + + block_size = _AZIMUTH_BYTES + laser_num * channel_bytes + + cos_el = config.cos_el + sin_el = config.sin_el + + # Build structured dtype for channel data + if channel_bytes == 3: + ch_dtype = np.dtype([("distance", " 0.0 + if not np.any(valid): + return _EMPTY_POINTS + + # Per-channel azimuth: block azimuth + channel-specific offset + block_az_rad = np.radians(azimuths_raw.astype(np.float32) / 100.0) + azimuths_rad = block_az_rad[:, np.newaxis] + config.azimuth_offset_rad + + # Hesai native frame: x = d*cos(el)*sin(az), y = d*cos(el)*cos(az), z = d*sin(el) + # The TF tree (hesai_top -> base_link) handles the frame conversion. + xy_dist = distances * cos_el + x = xy_dist * np.sin(azimuths_rad) + y = xy_dist * np.cos(azimuths_rad) + z = distances * sin_el + + return np.stack([x[valid], y[valid], z[valid], reflectivities[valid]], axis=0) diff --git a/t4_devkit/rosbag/pointcloud2.py b/t4_devkit/rosbag/pointcloud2.py new file mode 100644 index 00000000..32b7e6eb --- /dev/null +++ b/t4_devkit/rosbag/pointcloud2.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +import numpy as np + +from t4_devkit.dataclass import LidarPointCloud + +__all__ = ["pointcloud2_to_lidar"] + +# PointField datatype constants -> numpy dtype mapping +_DATATYPE_TO_NUMPY = { + 1: np.int8, # INT8 + 2: np.uint8, # UINT8 + 3: np.int16, # INT16 + 4: np.uint16, # UINT16 + 5: np.int32, # INT32 + 6: np.uint32, # UINT32 + 7: np.float32, # FLOAT32 + 8: np.float64, # FLOAT64 +} + + +def pointcloud2_to_lidar(msg: object) -> LidarPointCloud: + """Convert a deserialized PointCloud2 message to LidarPointCloud. + + Extracts x, y, z, intensity fields from the message and returns + a LidarPointCloud with shape (4, N), matching the format returned + by ``LidarPointCloud.from_file``. + + Args: + msg: Deserialized ``sensor_msgs/msg/PointCloud2`` message. + + Returns: + LidarPointCloud instance. + + Raises: + ValueError: If required fields (x, y, z) are not found. + """ + field_map: dict[str, tuple[int, int]] = {} # name -> (offset, datatype) + for f in msg.fields: + field_map[f.name] = (f.offset, f.datatype) + + for required in ("x", "y", "z"): + if required not in field_map: + raise ValueError(f"PointCloud2 message is missing required field: {required}") + + num_points = msg.height * msg.width + point_step = msg.point_step + + # Build a structured numpy dtype that covers all fields + padding + sorted_fields = sorted(msg.fields, key=lambda f: f.offset) + dt_list: list[tuple[str, type, int] | tuple[str, type]] = [] + offset = 0 + for f in sorted_fields: + if f.offset > offset: + dt_list.append((f"_pad{offset}", np.void, f.offset - offset)) + + np_dtype_base = _DATATYPE_TO_NUMPY.get(f.datatype) + if np_dtype_base is None: + raise ValueError(f"Unsupported PointField datatype: {f.datatype} for field '{f.name}'") + np_dtype = np.dtype(np_dtype_base).newbyteorder(">" if msg.is_bigendian else "<") + + count = int(getattr(f, "count", 1)) + if count < 1: + raise ValueError(f"Invalid PointField count for field '{f.name}': {count}") + + dt_list.append((f.name, np_dtype) if count == 1 else (f.name, np_dtype, count)) + offset = f.offset + np_dtype.itemsize * count + if offset < point_step: + dt_list.append(("_pad_end", np.void, point_step - offset)) + dtype = np.dtype(dt_list) + structured = np.frombuffer(msg.data, dtype=dtype, count=num_points) + + x = structured["x"].astype(np.float32) + y = structured["y"].astype(np.float32) + z = structured["z"].astype(np.float32) + intensity = ( + structured["intensity"].astype(np.float32) + if "intensity" in field_map + else np.zeros(num_points, dtype=np.float32) + ) + + points = np.stack([x, y, z, intensity], axis=0) + return LidarPointCloud(points=points) diff --git a/t4_devkit/rosbag/reader.py b/t4_devkit/rosbag/reader.py new file mode 100644 index 00000000..23fcf9cd --- /dev/null +++ b/t4_devkit/rosbag/reader.py @@ -0,0 +1,441 @@ +from __future__ import annotations + +import bisect +import logging +import threading +from pathlib import Path +from typing import TYPE_CHECKING + +import numpy as np +from rosbags.rosbag2 import Reader +from rosbags.typesys import Stores, get_typestore + +from t4_devkit.rosbag.pandar_decoder import ( + HESAI_MODELS, + pandarscan_to_lidar, + register_pandar_types, +) +from t4_devkit.rosbag.pointcloud2 import pointcloud2_to_lidar +from t4_devkit.rosbag.topic_mapping import TopicMapping + +if TYPE_CHECKING: + from rosbags.interfaces import Connection + + from t4_devkit.dataclass import LidarPointCloud + +__all__ = ["Rosbag2Reader"] + +_POINTCLOUD2_MSGTYPE = "sensor_msgs/msg/PointCloud2" +_PANDARSCAN_MSGTYPE = "pandar_msgs/msg/PandarScan" +_TF_STATIC_TOPIC = "/tf_static" + +_DEFAULT_TARGET_FRAME = "base_link" +_SUPPORTED_MSGTYPES = {_POINTCLOUD2_MSGTYPE, _PANDARSCAN_MSGTYPE} + +logger = logging.getLogger(__name__) + + +def _quat_to_matrix(x: float, y: float, z: float, w: float) -> np.ndarray: + """Convert quaternion (x, y, z, w) to a 3x3 rotation matrix.""" + return np.array( + [ + [1 - 2 * (y * y + z * z), 2 * (x * y - w * z), 2 * (x * z + w * y)], + [2 * (x * y + w * z), 1 - 2 * (x * x + z * z), 2 * (y * z - w * x)], + [2 * (x * z - w * y), 2 * (y * z + w * x), 1 - 2 * (x * x + y * y)], + ] + ) + + +def _make_transform(R: np.ndarray, t: np.ndarray) -> np.ndarray: + """Build a 4x4 homogeneous transform from a 3x3 rotation and 3-vector.""" + T = np.eye(4) + T[:3, :3] = R + T[:3, 3] = t + return T + + +def _read_tf_static(typestore: object, reader: Reader) -> dict[str, tuple[str, np.ndarray]]: + """Read ``/tf_static`` and return per-edge transforms. + + Returns: + Dict mapping ``child_frame_id`` to ``(parent_frame_id, T_4x4)`` + where ``T_4x4`` is the 4x4 homogeneous matrix satisfying + ``p_parent = T_4x4 @ p_child``. + """ + tf_conns = [c for c in reader.connections if c.topic == _TF_STATIC_TOPIC] + if not tf_conns: + return {} + + tree: dict[str, tuple[str, np.ndarray]] = {} + for conn, _ts, rawdata in reader.messages(connections=tf_conns): + msg = typestore.deserialize_cdr(rawdata, conn.msgtype) + for tf in msg.transforms: + parent = tf.header.frame_id + child = tf.child_frame_id + tr = tf.transform.translation + rot = tf.transform.rotation + R = _quat_to_matrix(rot.x, rot.y, rot.z, rot.w) + t = np.array([tr.x, tr.y, tr.z]) + tree[child] = (parent, _make_transform(R, t)) + + return tree + + +def _resolve_chain( + tree: dict[str, tuple[str, np.ndarray]], + frame_id: str, + target_frame: str = _DEFAULT_TARGET_FRAME, +) -> np.ndarray: + """Walk the TF tree from *frame_id* up to *target_frame* and compose transforms. + + Args: + tree: Per-edge TF dict from :func:`_read_tf_static`. + frame_id: Source sensor frame (e.g. ``"hesai_top"``). + target_frame: Destination frame. Defaults to ``"base_link"``. + + Returns: + 4x4 homogeneous matrix ``sensor2ego`` such that + ``p_target = sensor2ego @ p_sensor``. + + Raises: + ValueError: If no chain exists from *frame_id* to *target_frame*. + """ + if frame_id == target_frame: + return np.eye(4) + + visited: set[str] = set() + T = np.eye(4) + current = frame_id + while current != target_frame: + if current in visited: + break + visited.add(current) + if current not in tree: + break + parent, T_edge = tree[current] + T = T_edge @ T + current = parent + else: + return T + + raise ValueError( + f"No TF chain from '{frame_id}' to '{target_frame}'. " + f"Available frames: {sorted(tree.keys())}" + ) + + +class Rosbag2Reader: + """Reader for rosbag2 files that provides LiDAR point cloud data. + + Reads db3 or mcap format rosbag2 files from the ``input_bag/`` directory + and provides access to LiDAR point cloud data indexed by timestamp. + + Supports both ``sensor_msgs/msg/PointCloud2`` and + ``pandar_msgs/msg/PandarScan`` (Hesai raw packet) topics. + + Args: + bag_dir: Path to the rosbag2 directory (containing metadata.yaml). + topic_mapping: Optional list of ``TopicMapping`` instances. + If ``None``, supported LiDAR topics are auto-detected from the bag. + """ + + def __init__( + self, + bag_dir: str, + topic_mapping: list[TopicMapping] | None = None, + ) -> None: + bag_path = Path(bag_dir) + if not bag_path.is_dir(): + raise FileNotFoundError(f"Rosbag directory not found: {bag_dir}") + + self._typestore = get_typestore(Stores.EMPTY) + self._typestore.register(get_typestore(Stores.ROS2_HUMBLE).fielddefs) + + self._reader = Reader(bag_path) + self._reader.open() + self._lock = threading.Lock() + + # Find all supported LiDAR connections + lidar_connections: list[Connection] = [ + conn for conn in self._reader.connections if conn.msgtype in _SUPPORTED_MSGTYPES + ] + + if not lidar_connections: + self._reader.close() + raise ValueError( + f"No supported LiDAR topics found in rosbag at {bag_dir}. " + f"Supported types: {_SUPPORTED_MSGTYPES}. " + f"Available topics: {[c.topic for c in self._reader.connections]}" + ) + + # Register pandar_msgs types if any PandarScan connections exist + has_pandar = any(c.msgtype == _PANDARSCAN_MSGTYPE for c in lidar_connections) + if has_pandar: + register_pandar_types(self._typestore) + + # Build topic <-> channel mapping and sensor_type mapping + if topic_mapping is not None: + self._channel_to_topic = TopicMapping.to_channel_dict(topic_mapping) + self._channel_to_sensor_type: dict[str, str | None] = { + m.channel: m.sensor_type for m in topic_mapping + } + else: + # Auto-detect mode: only include PointCloud2 topics (PandarScan requires + # explicit sensor_type via TopicMapping). + auto_connections = [ + conn for conn in lidar_connections if conn.msgtype != _PANDARSCAN_MSGTYPE + ] + self._channel_to_topic = {conn.topic: conn.topic for conn in auto_connections} + self._channel_to_sensor_type = {} + + # Validate sensor_type for PandarScan topics + pandar_topics = {c.topic for c in lidar_connections if c.msgtype == _PANDARSCAN_MSGTYPE} + for channel, topic in self._channel_to_topic.items(): + if topic in pandar_topics: + sensor_type = self._channel_to_sensor_type.get(channel) + if sensor_type is None: + raise ValueError( + f"Topic '{topic}' (channel '{channel}') is a PandarScan topic. " + f"sensor_type must be specified in TopicMapping. " + f"Supported types: {sorted(HESAI_MODELS.keys())}" + ) + if sensor_type not in HESAI_MODELS: + raise ValueError( + f"Unsupported sensor_type '{sensor_type}' for channel '{channel}'. " + f"Supported types: {sorted(HESAI_MODELS.keys())}" + ) + + topic_to_channel = {v: k for k, v in self._channel_to_topic.items()} + + # Filter connections to only the mapped topics, indexed by topic for O(1) lookup + mapped_topics = set(self._channel_to_topic.values()) + self._connections = [conn for conn in lidar_connections if conn.topic in mapped_topics] + self._topic_connections: dict[str, list[Connection]] = {} + for conn in self._connections: + self._topic_connections.setdefault(conn.topic, []).append(conn) + + # Read /tf_static and build per-edge TF tree + self._tf_tree = _read_tf_static(self._typestore, self._reader) + + # Pre-compute sensor2ego per channel. Source priority: + # 1. Explicit ``sensor2ego_translation`` + ``sensor2ego_rotation`` + # on the TopicMapping (bypasses /tf_static — useful when the + # bag's TF tree is missing or holds an outdated calibration). + # 2. ``frame_id`` resolved against the bag's ``/tf_static`` chain. + # 3. No entry → points stay in sensor frame. + self._channel_sensor2ego: dict[str, np.ndarray] = {} + if topic_mapping is not None: + for m in topic_mapping: + if m.has_explicit_sensor2ego: + # ``sensor2ego_rotation`` is (w, x, y, z); ``_quat_to_matrix`` + # takes (x, y, z, w). Unpack accordingly. + qw, qx, qy, qz = m.sensor2ego_rotation + R = _quat_to_matrix(qx, qy, qz, qw) + t = np.asarray(m.sensor2ego_translation, dtype=np.float64) + self._channel_sensor2ego[m.channel] = _make_transform(R, t) + if m.frame_id is not None: + logger.debug( + "Channel '%s': explicit sensor2ego override is in " + "effect; frame_id='%s' is ignored for the sensor → " + "base_link step.", + m.channel, + m.frame_id, + ) + continue + if m.frame_id is None: + continue + try: + self._channel_sensor2ego[m.channel] = _resolve_chain( + self._tf_tree, + m.frame_id, + ) + except ValueError: + logger.warning( + "Channel '%s': frame_id='%s' not found in /tf_static. " + "Points will be in sensor frame.", + m.channel, + m.frame_id, + ) + + # Build timestamp index: channel -> sorted list of timestamp_ns + # Also build a cached list of timestamp_us per channel for bisect lookups + self._timestamp_ns: dict[str, list[int]] = {} + self._timestamp_us: dict[str, list[int]] = {} + self._build_timestamp_index(topic_to_channel) + + # Warn about mapped channels that have 0 messages + for channel, topic in self._channel_to_topic.items(): + if channel not in self._timestamp_ns: + logger.warning( + "Topic '%s' (channel '%s') has 0 messages in the rosbag. " + "LiDAR data will not be available for this channel.", + topic, + channel, + ) + + def _build_timestamp_index(self, topic_to_channel: dict[str, str]) -> None: + """Build timestamp index from the rosbag.""" + raw_ns: dict[str, list[int]] = {} + + for conn, timestamp_ns, _rawdata in self._reader.messages(self._connections): + channel = topic_to_channel[conn.topic] + raw_ns.setdefault(channel, []).append(timestamp_ns) + + for channel, ns_list in raw_ns.items(): + ns_list.sort() + self._timestamp_ns[channel] = ns_list + self._timestamp_us[channel] = [t // 1_000 for t in ns_list] + + @property + def channels(self) -> list[str]: + """Return list of available channel names.""" + return list(self._timestamp_ns.keys()) + + def has_channel(self, channel: str) -> bool: + """Check if the given channel has indexed messages. + + Args: + channel: Sensor channel name. + + Returns: + ``True`` if the channel has indexed messages. + """ + return channel in self._timestamp_ns + + def get_sensor2ego( + self, + frame_id: str, + target_frame: str = _DEFAULT_TARGET_FRAME, + ) -> np.ndarray: + """Return the 4x4 homogeneous transform from *frame_id* to *target_frame*. + + Resolves the ``/tf_static`` chain from the sensor frame to the target + frame (usually ``base_link``). The returned matrix satisfies + ``p_target = sensor2ego @ p_sensor``. + + Args: + frame_id: Source sensor frame (e.g. ``"hesai_top"``). + target_frame: Destination frame. Defaults to ``"base_link"``. + + Returns: + 4x4 ``np.ndarray`` (float64) homogeneous transformation matrix. + + Raises: + ValueError: If ``/tf_static`` is unavailable or no chain exists. + """ + if not self._tf_tree: + raise ValueError( + "No /tf_static data available in the rosbag. " + "Cannot compute sensor2ego transform." + ) + return _resolve_chain(self._tf_tree, frame_id, target_frame) + + def get_pointcloud( + self, + channel: str, + timestamp_us: int, + tolerance_us: int = 75_000, + min_completeness: float = 1.0, + ) -> LidarPointCloud: + """Retrieve a LiDAR point cloud from the rosbag matching the given timestamp. + + Automatically dispatches to the correct decoder based on the topic's + message type (PointCloud2 or PandarScan). + + When the channel has a ``frame_id`` in its ``TopicMapping`` and the + corresponding ``/tf_static`` chain exists, decoded PandarScan points + are automatically transformed to ``base_link``. + + Args: + channel: Sensor channel name. + timestamp_us: Target timestamp in microseconds (T4 format). + tolerance_us: Maximum allowed time difference in microseconds. + min_completeness: Minimum scan completeness for PandarScan topics + (0.0–1.0). If the fraction of received vs. expected UDP + packets is below this threshold the scan is rejected with + :class:`ValueError`. Defaults to ``1.0`` (reject any scan + with dropped packets). Set to ``0.0`` to disable the check. + + Returns: + LidarPointCloud with shape ``(4, N)`` matching ``LidarPointCloud.from_file`` format. + + Raises: + KeyError: If the channel is not available. + ValueError: If no message is found within the tolerance, or if + the PandarScan completeness is below *min_completeness*. + """ + if channel not in self._timestamp_us: + raise KeyError(f"Channel '{channel}' not found. Available: {self.channels}") + + ts_us_list = self._timestamp_us[channel] + ts_ns_list = self._timestamp_ns[channel] + + pos = bisect.bisect_left(ts_us_list, timestamp_us) + + # Find the closest timestamp (check pos and pos-1) + best_idx = None + best_diff = float("inf") + for candidate in (pos - 1, pos): + if 0 <= candidate < len(ts_us_list): + diff = abs(ts_us_list[candidate] - timestamp_us) + if diff < best_diff: + best_diff = diff + best_idx = candidate + + if best_idx is None or best_diff > tolerance_us: + raise ValueError( + f"No message found for channel '{channel}' within {tolerance_us}us " + f"of timestamp {timestamp_us}. " + f"Closest: {ts_us_list[best_idx] if best_idx is not None else 'N/A'}" + ) + + target_ns = ts_ns_list[best_idx] + topic = self._channel_to_topic[channel] + + conns_for_topic = self._topic_connections.get(topic) + if not conns_for_topic: + raise ValueError(f"No connections found for topic '{topic}' (channel '{channel}')") + + with self._lock: + # Workaround for a rosbags MCAP off-by-one (storage_mcap.py:591): + # `start < x.message_end_time` excludes the chunk when start equals + # the chunk's last-message ts (which is inclusive), causing 1ns-window + # queries at chunk boundaries to return empty. Widening start by 1ns + # selects the correct chunk; the [start, stop) interval still uniquely + # identifies the target packet. + for conn, ts_ns, rawdata in self._reader.messages( + connections=conns_for_topic, + start=target_ns - 1, + stop=target_ns + 1, + ): + msg = self._typestore.deserialize_cdr(rawdata, conn.msgtype) + if conn.msgtype == _PANDARSCAN_MSGTYPE: + pc = pandarscan_to_lidar( + msg, + self._channel_to_sensor_type[channel], + min_completeness=min_completeness, + ) + else: + pc = pointcloud2_to_lidar(msg) + # Apply sensor2ego (from /tf_static OR explicit TopicMapping + # override) uniformly across both decoder paths. + sensor2ego = self._channel_sensor2ego.get(channel) + if sensor2ego is not None: + R = sensor2ego[:3, :3] + t = sensor2ego[:3, 3] + pc.points[:3, :] = R @ pc.points[:3, :] + t[:, np.newaxis] + return pc + + raise ValueError( + f"Failed to read message for channel '{channel}' at timestamp {target_ns}ns" + ) + + def close(self) -> None: + """Close the rosbag reader.""" + self._reader.close() + + def __enter__(self) -> Rosbag2Reader: + return self + + def __exit__(self, *args: object) -> None: + self.close() diff --git a/t4_devkit/rosbag/topic_mapping.py b/t4_devkit/rosbag/topic_mapping.py new file mode 100644 index 00000000..c43aab44 --- /dev/null +++ b/t4_devkit/rosbag/topic_mapping.py @@ -0,0 +1,127 @@ +from __future__ import annotations + +from attrs import define, field, validators + +__all__ = ["TopicMapping"] + + +def _validate_topic_name(instance: TopicMapping, attribute: object, value: str) -> None: + """Validate that topic name starts with '/'.""" + if not value.startswith("/"): + raise ValueError(f"ROS topic must start with '/': got '{value}'") + + +def _validate_translation( + instance: TopicMapping, attribute: object, value: tuple[float, float, float] | None +) -> None: + """Translation must be a 3-tuple of finite floats, or ``None``.""" + if value is None: + return + if not isinstance(value, tuple) or len(value) != 3: + raise ValueError(f"sensor2ego_translation must be a tuple of 3 floats; got {value!r}") + for v in value: + if not isinstance(v, (int, float)) or v != v or abs(v) == float("inf"): + raise ValueError( + f"sensor2ego_translation entries must be finite numbers; got {value!r}" + ) + + +def _validate_rotation( + instance: TopicMapping, + attribute: object, + value: tuple[float, float, float, float] | None, +) -> None: + """Rotation must be a (w, x, y, z) quaternion or ``None``; near-unit norm enforced.""" + if value is None: + return + if not isinstance(value, tuple) or len(value) != 4: + raise ValueError( + f"sensor2ego_rotation must be a (w, x, y, z) tuple of 4 floats; got {value!r}" + ) + for v in value: + if not isinstance(v, (int, float)) or v != v or abs(v) == float("inf"): + raise ValueError(f"sensor2ego_rotation entries must be finite numbers; got {value!r}") + norm = sum(v * v for v in value) ** 0.5 + if not (0.98 < norm < 1.02): + raise ValueError(f"sensor2ego_rotation must be approximately unit-norm; got |q|={norm:.4f}") + + +@define(frozen=True) +class TopicMapping: + """Mapping from a T4 sensor channel name to a ROS topic name. + + Attributes: + channel (str): T4 sensor channel name (e.g. ``"LIDAR_CONCAT"``). + topic (str): ROS topic name (e.g. ``"/sensing/lidar/concatenated/pointcloud"``). + sensor_type (str | None): Hesai sensor model name for PandarScan topics + (e.g. ``"OT128"``, ``"XT32"``). + Required when the topic type is ``pandar_msgs/msg/PandarScan``. + frame_id (str | None): TF frame ID of the sensor (e.g. ``"hesai_top"``). + When specified and ``/tf_static`` is available in the rosbag, + the decoded point cloud is transformed from this frame to ``base_link``. + sensor2ego_translation (tuple[float, float, float] | None): Explicit + ``(x, y, z)`` translation from the sensor frame to ``base_link``, + in metres. When both ``sensor2ego_translation`` and + ``sensor2ego_rotation`` are set, the pair overrides the + ``/tf_static`` chain — the decoded point cloud is transformed + using these values and any ``frame_id`` is ignored for the + sensor → base_link step. Use this when the bag's ``/tf_static`` + is missing or holds an outdated calibration. + sensor2ego_rotation (tuple[float, float, float, float] | None): + Explicit ``(w, x, y, z)`` quaternion (T4 convention) paired with + ``sensor2ego_translation``. Both must be set together. + """ + + channel: str = field(validator=[validators.instance_of(str), validators.min_len(1)]) + topic: str = field(validator=[validators.instance_of(str), _validate_topic_name]) + sensor_type: str | None = field(default=None) + frame_id: str | None = field(default=None) + sensor2ego_translation: tuple[float, float, float] | None = field( + default=None, validator=_validate_translation + ) + sensor2ego_rotation: tuple[float, float, float, float] | None = field( + default=None, validator=_validate_rotation + ) + + def __attrs_post_init__(self) -> None: + # Cross-field validation: translation and rotation come as a pair. + t = self.sensor2ego_translation + r = self.sensor2ego_rotation + if (t is None) != (r is None): + raise ValueError( + "sensor2ego_translation and sensor2ego_rotation must be set " + "together (or both left as None)." + ) + + @property + def has_explicit_sensor2ego(self) -> bool: + """True iff an explicit sensor → base_link calibration is supplied.""" + return self.sensor2ego_translation is not None and self.sensor2ego_rotation is not None + + @staticmethod + def from_dict(mapping: dict[str, str]) -> list[TopicMapping]: + """Create a list of TopicMapping from a dict. + + Args: + mapping: Dict mapping channel names to ROS topic names. + + Returns: + List of TopicMapping instances. + + Raises: + TypeError: If keys or values are not strings. + ValueError: If a topic name does not start with ``/``. + """ + return [TopicMapping(channel=k, topic=v) for k, v in mapping.items()] + + @staticmethod + def to_channel_dict(mappings: list[TopicMapping]) -> dict[str, str]: + """Convert a list of TopicMapping to a channel-to-topic dict. + + Args: + mappings: List of TopicMapping instances. + + Returns: + Dict mapping channel names to topic names. + """ + return {m.channel: m.topic for m in mappings} diff --git a/t4_devkit/tier4.py b/t4_devkit/tier4.py index df439cfc..94371d87 100644 --- a/t4_devkit/tier4.py +++ b/t4_devkit/tier4.py @@ -13,7 +13,7 @@ from typing_extensions import deprecated from t4_devkit.common.geometry import is_box_in_image -from t4_devkit.dataclass import Box2D, Box3D, SemanticLabel, Shape, ShapeType +from t4_devkit.dataclass import Box2D, Box3D, LidarPointCloud, SemanticLabel, Shape, ShapeType from t4_devkit.helper import RenderingHelper, TimeseriesHelper from t4_devkit.schema import SchemaName, SensorModality, VisibilityLevel, build_schema from t4_devkit.schema.compatibility import fix_category_table @@ -126,6 +126,9 @@ def __init__( data_root: str, revision: str | None = None, verbose: bool = True, + *, + use_rosbag: bool = False, + topic_mapping: list | dict[str, str] | None = None, ) -> None: """Load database and creates reverse indexes and shortcuts. @@ -134,6 +137,13 @@ def __init__( revision (str | None, optional): You can specify any specific version if you want. If None, search the latest one. verbose (bool, optional): Whether to display status during load. + use_rosbag (bool, optional): Whether to read LiDAR data from rosbag + instead of processed .pcd.bin files. + topic_mapping (list[TopicMapping] | dict[str, str] | None, optional): + Mapping from sensor channel names to ROS topic names. + Accepts a list of ``TopicMapping`` instances or a dict + (e.g., ``{"LIDAR_TOP": "/sensing/lidar/top/pointcloud"}``). + If None and ``use_rosbag`` is True, topics are auto-detected. Examples: >>> from t4_devkit import T4Devkit @@ -214,10 +224,35 @@ def __init__( elapsed_time = time.time() - start_time print(f"Done loading in {elapsed_time:.3f} seconds.\n======") + # initialize rosbag reader if requested + self._rosbag_reader = None + if use_rosbag: + from t4_devkit.rosbag import Rosbag2Reader, TopicMapping + + if not osp.isdir(self.bag_dir): + raise FileNotFoundError( + f"Rosbag directory not found: {self.bag_dir}. " + "Cannot use use_rosbag=True without input_bag/ directory." + ) + if isinstance(topic_mapping, dict): + topic_mapping = TopicMapping.from_dict(topic_mapping) + elif topic_mapping is not None and not all( + isinstance(m, TopicMapping) for m in topic_mapping + ): + raise TypeError("topic_mapping must be dict[str, str] or list[TopicMapping]") + self._rosbag_reader = Rosbag2Reader(self.bag_dir, topic_mapping=topic_mapping) + if verbose: + print(f"Loaded rosbag reader with channels: {self._rosbag_reader.channels}") + # initialize helpers after finishing construction of T4Devkit self._timeseries_helper = TimeseriesHelper(self) self._rendering_helper = RenderingHelper(self) + def close(self) -> None: + """Close resources held by this instance.""" + if self._rosbag_reader is not None: + self._rosbag_reader.close() + @property def data_root(self) -> str: """Return the path to dataset root directory.""" @@ -385,6 +420,34 @@ def get_sample_data_path(self, sample_data_token: str) -> str: sd_record: SampleData = self.get("sample_data", sample_data_token) return osp.join(self.data_root, sd_record.filename) + def get_lidar_pointcloud( + self, + sample_data_token: str, + *, + metainfo_filepath: str | None = None, + ) -> LidarPointCloud: + """Return a LidarPointCloud for the given sample_data token. + + If ``use_rosbag`` is enabled and the channel exists in the rosbag, + reads from the rosbag. Otherwise, falls back to file-based loading. + + Args: + sample_data_token (str): Token of ``sample_data``. + metainfo_filepath (str | None, optional): Path to metainfo JSON file. + + Returns: + LidarPointCloud instance with shape (4, N). + """ + sd_record: SampleData = self.get("sample_data", sample_data_token) + + if self._rosbag_reader is not None and sd_record.channel: + if self._rosbag_reader.has_channel(sd_record.channel): + return self._rosbag_reader.get_pointcloud(sd_record.channel, sd_record.timestamp) + + # Fallback to file-based reading + filepath = osp.join(self.data_root, sd_record.filename) + return LidarPointCloud.from_file(filepath, metainfo_filepath=metainfo_filepath) + def get_sample_data( self, sample_data_token: str, diff --git a/t4_devkit/viewer/viewer.py b/t4_devkit/viewer/viewer.py index db69aa3a..701394d9 100644 --- a/t4_devkit/viewer/viewer.py +++ b/t4_devkit/viewer/viewer.py @@ -215,7 +215,7 @@ def render_box3ds(self, *args, **kwargs) -> None: self._render_box3ds_with_elements(*args, **kwargs) def _render_box3ds_with_boxes(self, seconds: float, boxes: Sequence[Box3D]) -> None: - rr.set_time_seconds(EntityPath.TIMELINE, seconds) + rr.set_time(EntityPath.TIMELINE, duration=seconds) batches: dict[str, BatchBox3D] = {} for box in boxes: @@ -288,7 +288,7 @@ def _render_box3ds_with_elements( future=future, ) - rr.set_time_seconds(EntityPath.TIMELINE, seconds) + rr.set_time(EntityPath.TIMELINE, duration=seconds) rr.log(format_entity(EntityPath.MAP, frame_id, EntityPath.BOX), batch.as_boxes3d()) @@ -344,7 +344,7 @@ def render_box2ds(self, *args, **kwargs) -> None: self._render_box2ds_with_elements(*args, **kwargs) def _render_box2ds_with_boxes(self, seconds: float, boxes: Sequence[Box2D]) -> None: - rr.set_time_seconds(EntityPath.TIMELINE, seconds) + rr.set_time(EntityPath.TIMELINE, duration=seconds) batches: dict[str, BatchBox2D] = {} for box in boxes: @@ -373,7 +373,7 @@ def _render_box2ds_with_elements( for roi, class_id, uuid in zip(rois, class_ids, uuids, strict=True): batch.append(roi=roi, class_id=class_id, uuid=uuid) - rr.set_time_seconds(EntityPath.TIMELINE, seconds) + rr.set_time(EntityPath.TIMELINE, duration=seconds) rr.log(format_entity(EntityPath.BASE_LINK, camera, EntityPath.BOX), batch.as_boxes2d()) @_check_spatial2d @@ -395,7 +395,7 @@ def render_segmentation2d( class_ids (Sequence[int]): Sequence of label ids. uuids (Sequence[str | None] | None, optional): Sequence of each instance ID. """ - rr.set_time_seconds(EntityPath.TIMELINE, seconds) + rr.set_time(EntityPath.TIMELINE, duration=seconds) batch = BatchSegmentation2D() if uuids is None: @@ -425,7 +425,7 @@ def render_pointcloud( color_mode (PointCloudColorMode, optional): Color mode for pointcloud. """ # TODO(ktro2828): add support of rendering pointcloud on images - rr.set_time_seconds(EntityPath.TIMELINE, seconds) + rr.set_time(EntityPath.TIMELINE, duration=seconds) entity_path = format_entity(EntityPath.BASE_LINK, channel) if color_mode == PointCloudColorMode.SEGMENTATION: @@ -450,10 +450,10 @@ def render_image(self, seconds: float, camera: str, image: str | NDArrayU8) -> N camera (str): Name of the camera channel. image (str | NDArrayU8): Image tensor or path of the image file. """ - rr.set_time_seconds(EntityPath.TIMELINE, seconds) + rr.set_time(EntityPath.TIMELINE, duration=seconds) entity_path = format_entity(EntityPath.BASE_LINK, camera) - entity = rr.ImageEncoded(path=image) if isinstance(image, str) else rr.Image(image) + entity = rr.EncodedImage(path=image) if isinstance(image, str) else rr.Image(image) rr.log(entity_path, entity) @@ -509,7 +509,7 @@ def _render_ego_without_schema( rotation: RotationLike, geocoordinate: Vector3Like | None = None, ) -> None: - rr.set_time_seconds(EntityPath.TIMELINE, seconds) + rr.set_time(EntityPath.TIMELINE, duration=seconds) rr.log( EntityPath.BASE_LINK, diff --git a/tests/rosbag/__init__.py b/tests/rosbag/__init__.py new file mode 100644 index 00000000..9d48db4f --- /dev/null +++ b/tests/rosbag/__init__.py @@ -0,0 +1 @@ +from __future__ import annotations diff --git a/tests/rosbag/test_pandar_decoder.py b/tests/rosbag/test_pandar_decoder.py new file mode 100644 index 00000000..32d3154b --- /dev/null +++ b/tests/rosbag/test_pandar_decoder.py @@ -0,0 +1,426 @@ +from __future__ import annotations + +import struct +from pathlib import Path + +import numpy as np +import pytest + +rosbags = pytest.importorskip("rosbags") + +from rosbags.rosbag2 import Writer # noqa: E402 +from rosbags.typesys import Stores, get_typestore # noqa: E402 + +from t4_devkit.rosbag.pandar_decoder import ( # noqa: E402 + HESAI_MODELS, + pandarscan_to_lidar, + register_pandar_types, +) +from t4_devkit.rosbag.reader import Rosbag2Reader # noqa: E402 +from t4_devkit.rosbag.topic_mapping import TopicMapping # noqa: E402 + +# XT32 factory byte +_XT32_FACTORY = 0x32 + + +def _build_xt32_packet( + azimuths_deg: list[float], + distances_mm: list[list[int]], + reflectivities: list[list[int]], + udp_sequence: int | None = None, +) -> bytes: + """Build a synthetic XT32 packet (pre-header + header + body + tail). + + Args: + azimuths_deg: Azimuth angle in degrees for each block. + distances_mm: Distance values (in raw units, 1 unit = 4mm) per block per channel. + reflectivities: Reflectivity values (0-255) per block per channel. + udp_sequence: Optional UDP Sequence number (u32 LE, appended as + "Additional information" after the tail, matching XT32 spec). + + Returns: + Raw packet bytes. + """ + block_num = len(azimuths_deg) + laser_num = 32 + dis_unit = 4 # 4mm + + # Pre-header (6 bytes) + pre_header = bytes([0xEE, 0xFF, 1, 0, 0, 0]) + + # Header (6 bytes) + header = bytes([laser_num, block_num, 0, dis_unit, 1, 0]) + + # Body + body = bytearray() + for blk in range(block_num): + azimuth_raw = int(azimuths_deg[blk] * 100) + body.extend(struct.pack(" None: + """Test decoding a packet with one block.""" + laser_num = 32 + # All channels at distance 1000 (= 4m), azimuth 0° + distances = [[1000] * laser_num] + reflectivities = [[128] * laser_num] + + packet = _build_xt32_packet( + azimuths_deg=[0.0], + distances_mm=distances, + reflectivities=reflectivities, + ) + + from t4_devkit.rosbag.pandar_decoder import _decode_packet + + config = HESAI_MODELS["XT32"] + points = _decode_packet(packet, config) + + assert points.shape[0] == 4 + assert points.shape[1] == laser_num + + # At azimuth=0°: x = d*cos(el)*sin(0) = 0, y = d*cos(el)*cos(0) = d*cos(el) + # All x values should be ~0 + np.testing.assert_array_almost_equal(points[0, :], 0.0, decimal=5) + # All intensities should be 128 + np.testing.assert_array_almost_equal(points[3, :], 128.0) + + def test_decode_zero_distance_filtered(self) -> None: + """Test that zero-distance points are filtered out.""" + laser_num = 32 + distances = [[0] * laser_num] + reflectivities = [[0] * laser_num] + + packet = _build_xt32_packet( + azimuths_deg=[90.0], + distances_mm=distances, + reflectivities=reflectivities, + ) + + from t4_devkit.rosbag.pandar_decoder import _decode_packet + + config = HESAI_MODELS["XT32"] + points = _decode_packet(packet, config) + assert points.shape[1] == 0 + + def test_decode_azimuth_90(self) -> None: + """Test decoding at azimuth=90° - x should be positive for 0° elevation.""" + laser_num = 32 + distances = [[2500] * laser_num] # 10m + reflectivities = [[100] * laser_num] + + packet = _build_xt32_packet( + azimuths_deg=[90.0], + distances_mm=distances, + reflectivities=reflectivities, + ) + + from t4_devkit.rosbag.pandar_decoder import _decode_packet + + config = HESAI_MODELS["XT32"] + points = _decode_packet(packet, config) + + # At azimuth=90°: x = d*cos(el)*sin(90°) = d*cos(el), y = d*cos(el)*cos(90°) ≈ 0 + # Channel 15 (elevation=0°): x = 10.0, y ≈ 0 + ch15_idx = None + for i in range(points.shape[1]): + if abs(points[2, i]) < 0.01: # z ≈ 0 means elevation ≈ 0 + ch15_idx = i + break + assert ch15_idx is not None + assert points[0, ch15_idx] == pytest.approx(10.0, abs=0.1) + assert points[1, ch15_idx] == pytest.approx(0.0, abs=0.1) + + def test_multiple_blocks(self) -> None: + """Test decoding a packet with multiple blocks.""" + laser_num = 32 + block_num = 4 + distances = [[500] * laser_num] * block_num + reflectivities = [[50] * laser_num] * block_num + azimuths = [0.0, 90.0, 180.0, 270.0] + + packet = _build_xt32_packet( + azimuths_deg=azimuths, + distances_mm=distances, + reflectivities=reflectivities, + ) + + from t4_devkit.rosbag.pandar_decoder import _decode_packet + + config = HESAI_MODELS["XT32"] + points = _decode_packet(packet, config) + assert points.shape[0] == 4 + assert points.shape[1] == laser_num * block_num + + def test_channel_count_mismatch_raises(self) -> None: + """Test that mismatched sensor type raises ValueError.""" + laser_num = 32 + distances = [[1000] * laser_num] + reflectivities = [[128] * laser_num] + + packet = _build_xt32_packet( + azimuths_deg=[0.0], + distances_mm=distances, + reflectivities=reflectivities, + ) + + from t4_devkit.rosbag.pandar_decoder import _decode_packet + + config = HESAI_MODELS["OT128"] # 128ch config for 32ch packet + with pytest.raises(ValueError, match="does not match"): + _decode_packet(packet, config) + + +class TestPandarScanToLidar: + """Test the full pandarscan_to_lidar conversion with mock messages.""" + + def test_basic_conversion(self) -> None: + """Test converting a mock PandarScan message.""" + laser_num = 32 + packet_data = _build_xt32_packet( + azimuths_deg=[45.0], + distances_mm=[[1000] * laser_num], + reflectivities=[[200] * laser_num], + ) + scan = _make_mock_scan([packet_data]) + pc = pandarscan_to_lidar(scan, "XT32") + assert pc.points.shape[0] == 4 + assert pc.points.shape[1] == laser_num + + def test_empty_packets_raises(self) -> None: + """Test that empty packets raise ValueError.""" + scan = _make_mock_scan([]) + with pytest.raises(ValueError, match="no packets"): + pandarscan_to_lidar(scan, "XT32") + + def test_unsupported_sensor_type_raises(self) -> None: + """Test that unsupported sensor type raises ValueError.""" + from types import SimpleNamespace + + scan = SimpleNamespace(header=None, packets=[object()]) + with pytest.raises(ValueError, match="Unsupported sensor type"): + pandarscan_to_lidar(scan, "UNKNOWN_SENSOR") + + def test_multiple_packets(self) -> None: + """Test combining multiple packets into one point cloud.""" + laser_num = 32 + packet_data_list = [ + _build_xt32_packet( + azimuths_deg=[az], + distances_mm=[[500] * laser_num], + reflectivities=[[100] * laser_num], + ) + for az in [0.0, 90.0, 180.0, 270.0] + ] + scan = _make_mock_scan(packet_data_list) + pc = pandarscan_to_lidar(scan, "XT32") + assert pc.points.shape[0] == 4 + assert pc.points.shape[1] == laser_num * 4 + + +def _create_pandar_rosbag( + bag_dir: Path, + topic: str, + timestamps_ns: list[int], +) -> None: + """Create a synthetic rosbag2 with PandarScan messages.""" + typestore = get_typestore(Stores.EMPTY) + typestore.register(get_typestore(Stores.ROS2_HUMBLE).fielddefs) + register_pandar_types(typestore) + + PandarScan = typestore.types["pandar_msgs/msg/PandarScan"] # noqa: N806 + PandarPacket = typestore.types["pandar_msgs/msg/PandarPacket"] # noqa: N806 + Header = typestore.types["std_msgs/msg/Header"] # noqa: N806 + Time = typestore.types["builtin_interfaces/msg/Time"] # noqa: N806 + + msgtype = "pandar_msgs/msg/PandarScan" + laser_num = 32 + + with Writer(bag_dir, version=9) as writer: + connection = writer.add_connection(topic, msgtype, typestore=typestore) + + for ts_ns in timestamps_ns: + sec = ts_ns // 1_000_000_000 + nanosec = ts_ns % 1_000_000_000 + + # Build a few packets per scan + packets = [] + for az in [0.0, 45.0, 90.0, 135.0]: + pkt_data = _build_xt32_packet( + azimuths_deg=[az], + distances_mm=[[1000] * laser_num], + reflectivities=[[128] * laser_num], + ) + # Pad to fixed 1500-byte buffer + padded = pkt_data + b"\x00" * (1500 - len(pkt_data)) + pkt = PandarPacket( + stamp=Time(sec=sec, nanosec=nanosec), + data=np.frombuffer(padded, dtype=np.uint8), + size=len(pkt_data), + ) + packets.append(pkt) + + msg = PandarScan( + header=Header( + stamp=Time(sec=sec, nanosec=nanosec), + frame_id="lidar_top", + ), + packets=packets, + ) + + serialized = typestore.serialize_cdr(msg, msgtype) + writer.write(connection, ts_ns, serialized) + + +class TestRosbag2ReaderPandarScan: + """Integration tests: read PandarScan from a synthetic rosbag.""" + + @pytest.fixture + def pandar_bag(self, tmp_path: Path) -> Path: + bag_dir = tmp_path / "input_bag" + timestamps_ns = [ + 1_704_067_200_000_000_000, + 1_704_067_200_100_000_000, + ] + _create_pandar_rosbag(bag_dir, "/sensing/lidar/top/pandar_packets", timestamps_ns) + return bag_dir + + def test_missing_sensor_type_raises(self, pandar_bag: Path) -> None: + """Test that PandarScan topic without sensor_type raises ValueError.""" + mapping = [TopicMapping(channel="LIDAR_TOP", topic="/sensing/lidar/top/pandar_packets")] + with pytest.raises(ValueError, match="sensor_type must be specified"): + Rosbag2Reader(str(pandar_bag), topic_mapping=mapping) + + def test_invalid_sensor_type_raises(self, pandar_bag: Path) -> None: + """Test that invalid sensor_type raises ValueError.""" + mapping = [ + TopicMapping( + channel="LIDAR_TOP", + topic="/sensing/lidar/top/pandar_packets", + sensor_type="INVALID", + ) + ] + with pytest.raises(ValueError, match="Unsupported sensor_type"): + Rosbag2Reader(str(pandar_bag), topic_mapping=mapping) + + def test_topic_mapping_pandar(self, pandar_bag: Path) -> None: + mapping = [ + TopicMapping( + channel="LIDAR_TOP", + topic="/sensing/lidar/top/pandar_packets", + sensor_type="XT32", + ) + ] + with Rosbag2Reader(str(pandar_bag), topic_mapping=mapping) as reader: + assert reader.has_channel("LIDAR_TOP") + pc = reader.get_pointcloud("LIDAR_TOP", 1_704_067_200_000_000) + assert pc.points.shape[0] == 4 + assert pc.points.shape[1] > 0 + + def test_get_pointcloud_pandar(self, pandar_bag: Path) -> None: + mapping = [ + TopicMapping( + channel="LIDAR_TOP", + topic="/sensing/lidar/top/pandar_packets", + sensor_type="XT32", + ) + ] + with Rosbag2Reader(str(pandar_bag), topic_mapping=mapping) as reader: + pc = reader.get_pointcloud("LIDAR_TOP", 1_704_067_200_000_000) + # 4 packets × 1 block × 32 channels = 128 points + assert pc.points.shape == (4, 128) + # All reflectivities should be 128 + np.testing.assert_array_almost_equal(pc.points[3, :], 128.0) + + +def _make_mock_scan(packet_data_list: list[bytes]) -> object: + """Create a mock PandarScan message from a list of raw packet bytes.""" + from types import SimpleNamespace + + stamp = SimpleNamespace(sec=0, nanosec=0) + packets = [SimpleNamespace(stamp=stamp, data=data, size=len(data)) for data in packet_data_list] + return SimpleNamespace(header=None, packets=packets) + + +class TestScanCompleteness: + """Tests for UDP Sequence-based packet drop detection.""" + + @staticmethod + def _make_packets( + seq_numbers: list[int], + ) -> list[bytes]: + """Build XT32 packets with given UDP Sequence numbers.""" + laser_num = 32 + return [ + _build_xt32_packet( + azimuths_deg=[float(i * 10 % 360)], + distances_mm=[[1000] * laser_num], + reflectivities=[[128] * laser_num], + udp_sequence=seq, + ) + for i, seq in enumerate(seq_numbers) + ] + + def test_complete_scan_passes(self) -> None: + """A scan with all packets present should pass any threshold.""" + packets = self._make_packets([100, 101, 102, 103, 104]) + scan = _make_mock_scan(packets) + pc = pandarscan_to_lidar(scan, "XT32", min_completeness=1.0) + assert pc.points.shape[0] == 4 + assert pc.points.shape[1] > 0 + + def test_incomplete_scan_rejected(self) -> None: + """A scan missing >50% of packets should be rejected at 90% threshold.""" + # Sequences 100, 101, 200 → expected 101 packets, actual 3 → 3% + packets = self._make_packets([100, 101, 200]) + scan = _make_mock_scan(packets) + with pytest.raises(ValueError, match="Incomplete PandarScan"): + pandarscan_to_lidar(scan, "XT32", min_completeness=0.9) + + def test_incomplete_scan_allowed_when_disabled(self) -> None: + """When min_completeness=0.0, incomplete scans should pass.""" + packets = self._make_packets([100, 101, 200]) + scan = _make_mock_scan(packets) + pc = pandarscan_to_lidar(scan, "XT32", min_completeness=0.0) + assert pc.points.shape[1] > 0 + + def test_minor_drop_below_threshold_passes(self) -> None: + """A scan with minor drops above the threshold should pass.""" + # 9 out of 10 packets present → 90% completeness + seqs = [100, 101, 102, 103, 104, 105, 106, 107, 109] + packets = self._make_packets(seqs) + scan = _make_mock_scan(packets) + pc = pandarscan_to_lidar(scan, "XT32", min_completeness=0.9) + assert pc.points.shape[1] > 0 + + def test_minor_drop_at_threshold_rejected(self) -> None: + """A scan exactly below the threshold should be rejected.""" + # 8 out of 10 packets → 80% < 90% threshold + seqs = [100, 101, 102, 103, 104, 105, 106, 109] + packets = self._make_packets(seqs) + scan = _make_mock_scan(packets) + with pytest.raises(ValueError, match="Incomplete PandarScan"): + pandarscan_to_lidar(scan, "XT32", min_completeness=0.9) + + def test_single_packet_skips_check(self) -> None: + """A scan with a single packet should skip the completeness check.""" + packets = self._make_packets([100]) + scan = _make_mock_scan(packets) + pc = pandarscan_to_lidar(scan, "XT32", min_completeness=1.0) + assert pc.points.shape[1] > 0 diff --git a/tests/rosbag/test_pointcloud2.py b/tests/rosbag/test_pointcloud2.py new file mode 100644 index 00000000..ad9318da --- /dev/null +++ b/tests/rosbag/test_pointcloud2.py @@ -0,0 +1,178 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +import pytest + +from t4_devkit.rosbag.pointcloud2 import pointcloud2_to_lidar + + +@dataclass +class MockPointField: + """Mock of sensor_msgs/msg/PointField.""" + + name: str + offset: int + datatype: int + count: int = 1 + + # PointField constants + INT8: int = 1 + UINT8: int = 2 + INT16: int = 3 + UINT16: int = 4 + INT32: int = 5 + UINT32: int = 6 + FLOAT32: int = 7 + FLOAT64: int = 8 + + +@dataclass +class MockPointCloud2: + """Mock of sensor_msgs/msg/PointCloud2.""" + + height: int + width: int + fields: list[MockPointField] + is_bigendian: bool + point_step: int + row_step: int + data: bytes + is_dense: bool = True + + +def _build_pointcloud2( + points: np.ndarray, + field_names: list[str] | None = None, +) -> MockPointCloud2: + """Build a mock PointCloud2 message from a numpy array. + + Args: + points: (N, D) float32 array. + field_names: Names for each column. Defaults to ["x", "y", "z", "intensity"]. + + Returns: + MockPointCloud2 message. + """ + if field_names is None: + field_names = ["x", "y", "z", "intensity"][: points.shape[1]] + + points = points.astype(np.float32) + n_points = points.shape[0] + n_fields = points.shape[1] + point_step = n_fields * 4 # float32 = 4 bytes + + fields = [ + MockPointField(name=name, offset=i * 4, datatype=MockPointField.FLOAT32) + for i, name in enumerate(field_names) + ] + + return MockPointCloud2( + height=1, + width=n_points, + fields=fields, + is_bigendian=False, + point_step=point_step, + row_step=point_step * n_points, + data=points.tobytes(), + ) + + +class TestPointCloud2ToLidar: + def test_basic_conversion(self) -> None: + points = np.array( + [ + [1.0, 2.0, 3.0, 0.5], + [4.0, 5.0, 6.0, 0.8], + [7.0, 8.0, 9.0, 1.0], + ], + dtype=np.float32, + ) + msg = _build_pointcloud2(points) + + result = pointcloud2_to_lidar(msg) + + assert result.points.shape == (4, 3) + np.testing.assert_array_almost_equal(result.points[0], [1.0, 4.0, 7.0]) # x + np.testing.assert_array_almost_equal(result.points[1], [2.0, 5.0, 8.0]) # y + np.testing.assert_array_almost_equal(result.points[2], [3.0, 6.0, 9.0]) # z + np.testing.assert_array_almost_equal(result.points[3], [0.5, 0.8, 1.0]) # intensity + + def test_missing_intensity_fills_zeros(self) -> None: + points = np.array( + [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], + dtype=np.float32, + ) + msg = _build_pointcloud2(points, field_names=["x", "y", "z"]) + + result = pointcloud2_to_lidar(msg) + + assert result.points.shape == (4, 2) + np.testing.assert_array_almost_equal(result.points[3], [0.0, 0.0]) + + def test_missing_xyz_raises(self) -> None: + points = np.array([[1.0, 2.0]], dtype=np.float32) + msg = _build_pointcloud2(points, field_names=["x", "y"]) + + with pytest.raises(ValueError, match="missing required field: z"): + pointcloud2_to_lidar(msg) + + def test_extra_fields_ignored(self) -> None: + points = np.array( + [[1.0, 2.0, 3.0, 0.5, 10.0, 20.0]], + dtype=np.float32, + ) + msg = _build_pointcloud2( + points, field_names=["x", "y", "z", "intensity", "ring", "timestamp"] + ) + + result = pointcloud2_to_lidar(msg) + + assert result.points.shape == (4, 1) + np.testing.assert_array_almost_equal(result.points[:, 0], [1.0, 2.0, 3.0, 0.5]) + + def test_empty_pointcloud(self) -> None: + points = np.zeros((0, 4), dtype=np.float32) + msg = _build_pointcloud2(points) + + result = pointcloud2_to_lidar(msg) + + assert result.points.shape == (4, 0) + + def test_uint8_intensity(self) -> None: + """Test with intensity as UINT8 (common in some LiDAR drivers).""" + n_points = 3 + # Build manually: x(f32), y(f32), z(f32), intensity(u8), padding(3 bytes) + point_step = 16 # 4+4+4+1+3 padding to align + fields = [ + MockPointField(name="x", offset=0, datatype=MockPointField.FLOAT32), + MockPointField(name="y", offset=4, datatype=MockPointField.FLOAT32), + MockPointField(name="z", offset=8, datatype=MockPointField.FLOAT32), + MockPointField(name="intensity", offset=12, datatype=MockPointField.UINT8), + ] + + data = bytearray() + xyz_values = [(1.0, 2.0, 3.0), (4.0, 5.0, 6.0), (7.0, 8.0, 9.0)] + intensities = [100, 200, 50] + for (x, y, z), intensity in zip(xyz_values, intensities): + data.extend(np.float32(x).tobytes()) + data.extend(np.float32(y).tobytes()) + data.extend(np.float32(z).tobytes()) + data.extend(np.uint8(intensity).tobytes()) + data.extend(b"\x00" * 3) # padding + + msg = MockPointCloud2( + height=1, + width=n_points, + fields=fields, + is_bigendian=False, + point_step=point_step, + row_step=point_step * n_points, + data=bytes(data), + ) + + result = pointcloud2_to_lidar(msg) + + assert result.points.shape == (4, 3) + np.testing.assert_array_almost_equal(result.points[3], [100.0, 200.0, 50.0]) diff --git a/tests/rosbag/test_reader.py b/tests/rosbag/test_reader.py new file mode 100644 index 00000000..59a5a5e1 --- /dev/null +++ b/tests/rosbag/test_reader.py @@ -0,0 +1,228 @@ +from __future__ import annotations + +import struct +from pathlib import Path + +import numpy as np +import pytest + +rosbags = pytest.importorskip("rosbags") + +from rosbags.rosbag2 import Writer # noqa: E402 +from rosbags.typesys import Stores, get_typestore # noqa: E402 + +from t4_devkit.rosbag.reader import Rosbag2Reader # noqa: E402 +from t4_devkit.rosbag.topic_mapping import TopicMapping # noqa: E402 + + +def _create_test_rosbag(bag_dir: Path, topic: str, timestamps_ns: list[int]) -> None: + """Create a synthetic rosbag2 with PointCloud2 messages. + + Args: + bag_dir: Directory to create the bag in. + topic: Topic name. + timestamps_ns: List of timestamps in nanoseconds. + """ + typestore = get_typestore(Stores.ROS2_HUMBLE) + PointCloud2 = typestore.types["sensor_msgs/msg/PointCloud2"] # noqa: N806 + PointField = typestore.types["sensor_msgs/msg/PointField"] # noqa: N806 + Header = typestore.types["std_msgs/msg/Header"] # noqa: N806 + Time = typestore.types["builtin_interfaces/msg/Time"] # noqa: N806 + + msgtype = "sensor_msgs/msg/PointCloud2" + connection = None + + with Writer(bag_dir, version=9) as writer: + connection = writer.add_connection(topic, msgtype, typestore=typestore) + + for ts_ns in timestamps_ns: + sec = ts_ns // 1_000_000_000 + nanosec = ts_ns % 1_000_000_000 + + # Create a simple PointCloud2 with 3 points + n_points = 3 + point_step = 16 # x(4) + y(4) + z(4) + intensity(4) + data = bytearray() + for i in range(n_points): + data.extend(struct.pack(" Path: + """Create a temporary rosbag with PointCloud2 messages.""" + bag_dir = tmp_path / "input_bag" + timestamps_ns = [ + 1_704_067_200_000_000_000, # 2024-01-01T00:00:00.000 + 1_704_067_200_100_000_000, # 2024-01-01T00:00:00.100 + 1_704_067_200_200_000_000, # 2024-01-01T00:00:00.200 + ] + _create_test_rosbag(bag_dir, "/sensing/lidar/top/pointcloud", timestamps_ns) + return bag_dir + + def test_init_auto_detect(self, bag_with_pointclouds: Path) -> None: + reader = Rosbag2Reader(str(bag_with_pointclouds)) + assert len(reader.channels) == 1 + assert reader.has_channel("/sensing/lidar/top/pointcloud") + reader.close() + + def test_init_with_topic_mapping(self, bag_with_pointclouds: Path) -> None: + mapping = [TopicMapping(channel="LIDAR_TOP", topic="/sensing/lidar/top/pointcloud")] + reader = Rosbag2Reader(str(bag_with_pointclouds), topic_mapping=mapping) + assert reader.has_channel("LIDAR_TOP") + assert not reader.has_channel("/sensing/lidar/top/pointcloud") + reader.close() + + def test_get_pointcloud_exact(self, bag_with_pointclouds: Path) -> None: + mapping = [TopicMapping(channel="LIDAR_TOP", topic="/sensing/lidar/top/pointcloud")] + with Rosbag2Reader(str(bag_with_pointclouds), topic_mapping=mapping) as reader: + pc = reader.get_pointcloud("LIDAR_TOP", 1_704_067_200_000_000) + + assert pc.points.shape == (4, 3) + # First point: x=0, y=1, z=2, intensity=0 + np.testing.assert_array_almost_equal(pc.points[:, 0], [0.0, 1.0, 2.0, 0.0]) + + def test_get_pointcloud_near_match(self, bag_with_pointclouds: Path) -> None: + mapping = [TopicMapping(channel="LIDAR_TOP", topic="/sensing/lidar/top/pointcloud")] + with Rosbag2Reader(str(bag_with_pointclouds), topic_mapping=mapping) as reader: + # Slightly off timestamp (1ms difference), within default 75ms tolerance + pc = reader.get_pointcloud("LIDAR_TOP", 1_704_067_200_001_000) + assert pc.points.shape == (4, 3) + + def test_get_pointcloud_out_of_tolerance(self, bag_with_pointclouds: Path) -> None: + mapping = [TopicMapping(channel="LIDAR_TOP", topic="/sensing/lidar/top/pointcloud")] + with Rosbag2Reader(str(bag_with_pointclouds), topic_mapping=mapping) as reader: + with pytest.raises(ValueError, match="No message found"): + # Timestamp far from any message (1 second off) + reader.get_pointcloud("LIDAR_TOP", 1_704_067_201_000_000) + + def test_get_pointcloud_unknown_channel(self, bag_with_pointclouds: Path) -> None: + mapping = [TopicMapping(channel="LIDAR_TOP", topic="/sensing/lidar/top/pointcloud")] + with Rosbag2Reader(str(bag_with_pointclouds), topic_mapping=mapping) as reader: + with pytest.raises(KeyError, match="LIDAR_BOTTOM"): + reader.get_pointcloud("LIDAR_BOTTOM", 1_704_067_200_000_000) + + def test_missing_directory(self) -> None: + with pytest.raises(FileNotFoundError): + Rosbag2Reader("/nonexistent/path") + + def test_context_manager(self, bag_with_pointclouds: Path) -> None: + with Rosbag2Reader(str(bag_with_pointclouds)) as reader: + assert len(reader.channels) > 0 + + def test_get_pointcloud_start_widened( + self, + bag_with_pointclouds: Path, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Regression: start is widened by 1ns to dodge rosbags MCAP chunk-filter off-by-one.""" + mapping = [TopicMapping(channel="LIDAR_TOP", topic="/sensing/lidar/top/pointcloud")] + with Rosbag2Reader(str(bag_with_pointclouds), topic_mapping=mapping) as reader: + captured_starts: list[int] = [] + original_messages = reader._reader.messages + + def patched_messages(*args, **kwargs): + if "start" in kwargs: + captured_starts.append(kwargs["start"]) + return original_messages(*args, **kwargs) + + monkeypatch.setattr(reader._reader, "messages", patched_messages) + + target_ts_us = 1_704_067_200_000_000 + target_ts_ns = target_ts_us * 1_000 + pc = reader.get_pointcloud("LIDAR_TOP", target_ts_us) + + assert pc.points.shape == (4, 3) + assert len(captured_starts) == 1 + assert captured_starts[0] == target_ts_ns - 1 + + def test_multiple_timestamps(self, bag_with_pointclouds: Path) -> None: + mapping = [TopicMapping(channel="LIDAR_TOP", topic="/sensing/lidar/top/pointcloud")] + with Rosbag2Reader(str(bag_with_pointclouds), topic_mapping=mapping) as reader: + # Read all three timestamps + for ts_us in [ + 1_704_067_200_000_000, + 1_704_067_200_100_000, + 1_704_067_200_200_000, + ]: + pc = reader.get_pointcloud("LIDAR_TOP", ts_us) + assert pc.points.shape == (4, 3) + + def test_explicit_sensor2ego_translation_only(self, bag_with_pointclouds: Path) -> None: + """Explicit ``sensor2ego_*`` translates points without /tf_static.""" + # Pure translation by (10, 20, 30); identity rotation. + mapping = [ + TopicMapping( + channel="LIDAR_TOP", + topic="/sensing/lidar/top/pointcloud", + sensor2ego_translation=(10.0, 20.0, 30.0), + sensor2ego_rotation=(1.0, 0.0, 0.0, 0.0), + ) + ] + with Rosbag2Reader(str(bag_with_pointclouds), topic_mapping=mapping) as reader: + pc = reader.get_pointcloud("LIDAR_TOP", 1_704_067_200_000_000) + # First point in sensor frame is (0, 1, 2); after translation: (10, 21, 32). + np.testing.assert_array_almost_equal(pc.points[:3, 0], [10.0, 21.0, 32.0]) + + def test_explicit_sensor2ego_overrides_frame_id(self, bag_with_pointclouds: Path) -> None: + """When the override is set, ``frame_id`` is ignored for the transform.""" + # ``frame_id`` resolves to identity (no /tf_static in this mock bag), + # so without the override points stay in sensor frame. With the + # override they should be shifted by the explicit translation. + mapping = [ + TopicMapping( + channel="LIDAR_TOP", + topic="/sensing/lidar/top/pointcloud", + frame_id="some_frame_not_in_tf", + sensor2ego_translation=(1.0, 2.0, 3.0), + sensor2ego_rotation=(1.0, 0.0, 0.0, 0.0), + ) + ] + with Rosbag2Reader(str(bag_with_pointclouds), topic_mapping=mapping) as reader: + pc = reader.get_pointcloud("LIDAR_TOP", 1_704_067_200_000_000) + np.testing.assert_array_almost_equal(pc.points[:3, 0], [1.0, 3.0, 5.0]) + + def test_explicit_sensor2ego_quaternion_rotation(self, bag_with_pointclouds: Path) -> None: + """Rotation quaternion is applied correctly (90° about +Z swaps x↔y).""" + # q = (cos(45°), 0, 0, sin(45°)) = 90° rotation about +z. + import math + + s = math.sin(math.pi / 4) + mapping = [ + TopicMapping( + channel="LIDAR_TOP", + topic="/sensing/lidar/top/pointcloud", + sensor2ego_translation=(0.0, 0.0, 0.0), + sensor2ego_rotation=(s, 0.0, 0.0, s), + ) + ] + with Rosbag2Reader(str(bag_with_pointclouds), topic_mapping=mapping) as reader: + pc = reader.get_pointcloud("LIDAR_TOP", 1_704_067_200_000_000) + # (0, 1, 2) → (-1, 0, 2) after 90° rotation about +z. + np.testing.assert_array_almost_equal(pc.points[:3, 0], [-1.0, 0.0, 2.0]) diff --git a/tests/rosbag/test_topic_mapping.py b/tests/rosbag/test_topic_mapping.py new file mode 100644 index 00000000..8428c0a6 --- /dev/null +++ b/tests/rosbag/test_topic_mapping.py @@ -0,0 +1,116 @@ +from __future__ import annotations + +import pytest + +from t4_devkit.rosbag.topic_mapping import TopicMapping + + +class TestTopicMapping: + def test_valid(self) -> None: + m = TopicMapping(channel="LIDAR_TOP", topic="/sensing/lidar/top/pointcloud") + assert m.channel == "LIDAR_TOP" + assert m.topic == "/sensing/lidar/top/pointcloud" + + def test_topic_must_start_with_slash(self) -> None: + with pytest.raises(ValueError, match="must start with '/'"): + TopicMapping(channel="LIDAR_TOP", topic="sensing/lidar/top/pointcloud") + + def test_channel_must_not_be_empty(self) -> None: + with pytest.raises(ValueError): + TopicMapping(channel="", topic="/topic") + + def test_channel_must_be_string(self) -> None: + with pytest.raises(TypeError): + TopicMapping(channel=123, topic="/topic") # type: ignore[arg-type] + + def test_topic_must_be_string(self) -> None: + with pytest.raises(TypeError): + TopicMapping(channel="LIDAR", topic=123) # type: ignore[arg-type] + + def test_from_dict(self) -> None: + raw = { + "LIDAR_TOP": "/sensing/lidar/top/pointcloud", + "LIDAR_CONCAT": "/sensing/lidar/concatenated/pointcloud", + } + mappings = TopicMapping.from_dict(raw) + assert len(mappings) == 2 + assert mappings[0].channel == "LIDAR_TOP" + assert mappings[1].channel == "LIDAR_CONCAT" + + def test_from_dict_invalid_topic(self) -> None: + with pytest.raises(ValueError, match="must start with '/'"): + TopicMapping.from_dict({"LIDAR": "bad_topic"}) + + def test_to_channel_dict(self) -> None: + mappings = [ + TopicMapping(channel="LIDAR_TOP", topic="/sensing/lidar/top/pointcloud"), + TopicMapping(channel="LIDAR_CONCAT", topic="/sensing/lidar/concatenated/pointcloud"), + ] + d = TopicMapping.to_channel_dict(mappings) + assert d == { + "LIDAR_TOP": "/sensing/lidar/top/pointcloud", + "LIDAR_CONCAT": "/sensing/lidar/concatenated/pointcloud", + } + + def test_frozen(self) -> None: + m = TopicMapping(channel="LIDAR_TOP", topic="/topic") + with pytest.raises(AttributeError): + m.channel = "OTHER" # type: ignore[misc] + + def test_sensor2ego_defaults_none(self) -> None: + m = TopicMapping(channel="LIDAR_TOP", topic="/topic") + assert m.sensor2ego_translation is None + assert m.sensor2ego_rotation is None + assert m.has_explicit_sensor2ego is False + + def test_sensor2ego_pair_valid(self) -> None: + m = TopicMapping( + channel="LIDAR_TOP", + topic="/topic", + sensor2ego_translation=(0.9, 0.0, 2.18), + sensor2ego_rotation=(1.0, 0.0, 0.0, 0.0), + ) + assert m.has_explicit_sensor2ego is True + assert m.sensor2ego_translation == (0.9, 0.0, 2.18) + assert m.sensor2ego_rotation == (1.0, 0.0, 0.0, 0.0) + + def test_sensor2ego_pair_required_together(self) -> None: + with pytest.raises(ValueError, match="together"): + TopicMapping( + channel="LIDAR_TOP", + topic="/topic", + sensor2ego_translation=(0.0, 0.0, 0.0), + ) + with pytest.raises(ValueError, match="together"): + TopicMapping( + channel="LIDAR_TOP", + topic="/topic", + sensor2ego_rotation=(1.0, 0.0, 0.0, 0.0), + ) + + def test_sensor2ego_translation_must_be_3tuple(self) -> None: + with pytest.raises(ValueError, match="3 floats"): + TopicMapping( + channel="LIDAR_TOP", + topic="/topic", + sensor2ego_translation=(0.0, 0.0), # type: ignore[arg-type] + sensor2ego_rotation=(1.0, 0.0, 0.0, 0.0), + ) + + def test_sensor2ego_rotation_must_be_4tuple(self) -> None: + with pytest.raises(ValueError, match="4 floats"): + TopicMapping( + channel="LIDAR_TOP", + topic="/topic", + sensor2ego_translation=(0.0, 0.0, 0.0), + sensor2ego_rotation=(1.0, 0.0, 0.0), # type: ignore[arg-type] + ) + + def test_sensor2ego_rotation_must_be_unit_norm(self) -> None: + with pytest.raises(ValueError, match="unit-norm"): + TopicMapping( + channel="LIDAR_TOP", + topic="/topic", + sensor2ego_translation=(0.0, 0.0, 0.0), + sensor2ego_rotation=(2.0, 0.0, 0.0, 0.0), + )