diff --git a/src/machinevisiontoolbox/Sources.py b/src/machinevisiontoolbox/Sources.py index efe1366a..9e85584a 100644 --- a/src/machinevisiontoolbox/Sources.py +++ b/src/machinevisiontoolbox/Sources.py @@ -49,6 +49,7 @@ try: from rosbags.rosbag1 import Reader as _RosBagReader1 from rosbags.rosbag2 import Reader as _RosBagReader2 + from rosbags.serde import SerdeError as _SerdeError from rosbags.typesys import Stores as _Stores, get_typestore as _get_typestore _rosbags_available = True @@ -57,8 +58,16 @@ _RosBagReader2 = None _Stores = None _get_typestore = None + _SerdeError = None _rosbags_available = False +# Errors raised while deserializing a message that do not rule out the store: +# the bytes may simply carry a trailer, so another trim amount is worth trying. +if _SerdeError is None: + _TRIM_RETRY_ERRORS: tuple[type[BaseException], ...] = (AssertionError,) +else: + _TRIM_RETRY_ERRORS = (AssertionError, _SerdeError) + try: import roslibpy @@ -106,10 +115,123 @@ def _make_image(*args: Any, **kwargs: Any) -> Image: return Image(*args, **kwargs) # pyright: ignore[reportAbstractUsage] -def _set_sample_metadata(sample: Any, timestamp: int, topic: str) -> None: +def _set_sample_metadata( + sample: Any, timestamp: int, topic: str, depth_units: float | None = None +) -> None: """Attach dynamic timestamp/topic metadata used by stream sources.""" setattr(sample, "timestamp", timestamp) setattr(sample, "topic", topic) + if depth_units is not None: + setattr(sample, "depth_units", depth_units) + + +def _ros1_image_payload_end(raw: bytes, msgtype: str) -> int | None: + """Return the end offset of the pixel payload of a ROS 1 image message. + + Uses the self-describing field layout of ``sensor_msgs/Image`` and + ``sensor_msgs/CompressedImage`` (all lengths are little-endian ``uint32`` + prefixes) so the payload boundary is computed rather than guessed. + + :return: offset of the first byte after the payload, or ``None`` if the + buffer does not match a valid image message + :rtype: int or None + """ + + def u32(pos: int) -> int: + if pos + 4 > len(raw): + raise ValueError("truncated") + return struct.unpack_from(" int | None: + """Return the end offset of the pixel payload of a CDR-encoded image message. + + Same idea as :func:`_ros1_image_payload_end` for ROS 2 CDR encoding: the + encapsulation header is skipped, every field is aligned relative to the + start of the CDR body and string lengths include the terminating NUL. + + :return: offset of the first byte after the payload, or ``None`` if the + buffer does not match a valid image message + :rtype: int or None + """ + if len(raw) < 4: + return None + order = "<" if raw[1] & 0x01 else ">" # 0x01 = CDR_LE, 0x00 = CDR_BE + body = raw[4:] + + def u32(pos: int) -> int: + if pos + 4 > len(body): + raise ValueError("truncated") + return struct.unpack_from(order + "I", body, pos)[0] + + try: + pos = 8 # Time: sec, nanosec + pos += 4 + u32(pos) # frame_id + if msgtype.endswith("CompressedImage"): + pos += (-pos) % 4 + pos += 4 + u32(pos) # format + pos += (-pos) % 4 + end = pos + 4 + u32(pos) # data + return end + 4 if end <= len(body) else None + + pos += (-pos) % 4 + height = u32(pos) + pos += 8 # height, width + pos += (-pos) % 4 + pos += 4 + u32(pos) # encoding + pos += 1 # is_bigendian + pos += (-pos) % 4 + step = u32(pos) + pos += 4 # step + pos += (-pos) % 4 + nbytes = u32(pos) + pos += 4 # data length + if nbytes != height * step: + return None + end = pos + nbytes + return end + 4 if end <= len(body) else None + except (ValueError, struct.error): + return None + + +def _image_payload_end(raw: bytes, msgtype: str, is_ros2: bool) -> int | None: + """Return the end offset of the pixel payload of a serialized image message. + + Intel RealSense SDK bags append a 4-byte footer to every image message: the + depth scale, e.g. ``0.001``. Because ``sensor_msgs`` image messages are + self-describing, the payload end can be derived from the field layout and + only the bytes that follow it are treated as a footer. + + :return: offset of the first trailing byte, or ``None`` if the layout does + not match a valid image message + :rtype: int or None + """ + if is_ros2: + return _cdr_image_payload_end(raw, msgtype) + return _ros1_image_payload_end(raw, msgtype) class ImageSource(ABC): @@ -3365,10 +3487,11 @@ def _autodetect_typestore(self, rawdata: bytes, msgtype: str) -> bool: available = {s.name for s in _Stores} # First pass: ROS 1 wire format (ros1 and ros2 typestores). - # RealSense SDK bags are positively identified by their realsense_msgs - # connection types, so the 4-byte trailer is certain — no need to probe - # other trim amounts. - trim_order = (4,) if self._is_realsense else (0, 4, 8) + # Try the untrimmed buffer first: RealSense SDK bags append a footer to + # some messages, but that must not make a self-contained message fail. + # The non-zero trim amounts remain as a fallback for producers that add + # trailing bytes which the untrimmed probe cannot accept. + trim_order = (0, 4, 8) for name in self._ROS1_STORES_ORDERED + self._ROS2_STORES_ORDERED: if name not in available: continue @@ -3381,12 +3504,12 @@ def _autodetect_typestore(self, rawdata: bytes, msgtype: str) -> bool: self.release = name self._trim_bytes = trim suffix = f", {trim}-byte trailer stripped" if trim else "" - print(f" auto-detect: {name} (ros1 encoding{suffix})") + if self.verbose: + print(f" auto-detect: {name} (ros1 encoding{suffix})") return True - except AssertionError: - # rosbags raises AssertionError when trailing bytes remain - # after parsing (assert pos == len(rawdata)). Try the next - # trim amount. + except _TRIM_RETRY_ERRORS: + # Trailing bytes remain after parsing, or the message cannot + # be decoded at this length. Try the next trim amount. continue except (KeyError, UnicodeDecodeError, struct.error) as e: if self.verbose: @@ -3419,9 +3542,16 @@ def _autodetect_typestore(self, rawdata: bytes, msgtype: str) -> bool: self.release = name self._is_ros2 = True # use CDR for all subsequent messages self._trim_bytes = 0 - print(f" auto-detect: {name} (CDR encoding)") + if self.verbose: + print(f" auto-detect: {name} (CDR encoding)") return True - except (AssertionError, KeyError, UnicodeDecodeError, struct.error) as e: + except (KeyError, UnicodeDecodeError, struct.error) as e: + if self.verbose: + print( + f" auto-detect: {name} (CDR) failed: {type(e).__name__}: {e}" + ) + continue + except _TRIM_RETRY_ERRORS as e: if self.verbose: print( f" auto-detect: {name} (CDR) failed: {type(e).__name__}: {e}" @@ -3436,9 +3566,9 @@ def _autodetect_typestore(self, rawdata: bytes, msgtype: str) -> bool: if self.verbose: print(f" auto-detect: rawdata[:16] = {rawdata[:16].hex()}") - print( - f" auto-detect: FAILED for msgtype={msgtype!r} — no matching typestore found" - ) + print( + f" auto-detect: FAILED for msgtype={msgtype!r} — no matching typestore found" + ) return False def __iter__(self) -> Iterator[Any]: @@ -3465,19 +3595,52 @@ def __iter__(self) -> Iterator[Any]: if self._is_ros2 else self.typestore.deserialize_ros1 ) - _raw = ( - rawdata - if self._trim_bytes == 0 - else rawdata[: -self._trim_bytes] + # RealSense SDK bags append a 4-byte footer to image + # messages: the depth scale. Resolve the payload boundary + # from the message layout so that only the bytes beyond the + # payload are treated as a footer. + _payload_end = ( + _image_payload_end(rawdata, connection.msgtype, self._is_ros2) + if connection.msgtype.endswith("Image") + else None + ) + _footer = b"" + if _payload_end is not None: + _candidates: list[bytes] = [rawdata[:_payload_end]] + _footer = rawdata[_payload_end:] + elif self._trim_bytes: + # The detector inferred a trailing footer from one + # message; other connections in the same bag (status + # topics, IMU, …) may be self-contained. + _candidates = [rawdata[: -self._trim_bytes], rawdata] + else: + _candidates = [rawdata] + + msg = None + _last_error: BaseException | None = None + for _candidate in _candidates: + try: + msg = cast(Any, _deser(_candidate, connection.msgtype)) + break + except _TRIM_RETRY_ERRORS as e: + _last_error = e + continue + if msg is None: + assert _last_error is not None + raise _last_error + + _depth_units = ( + struct.unpack(" Iterator[Any]: else: img = _make_image(img_array, **self.args) _set_sample_metadata( - img, self._stamp_to_ns(msg.header.stamp), connection.topic + img, + self._stamp_to_ns(msg.header.stamp), + connection.topic, + depth_units=_depth_units, ) yield img @@ -3572,7 +3738,10 @@ def __iter__(self) -> Iterator[Any]: else: img = _make_image(img_array, **self.args) _set_sample_metadata( - img, self._stamp_to_ns(msg.header.stamp), connection.topic + img, + self._stamp_to_ns(msg.header.stamp), + connection.topic, + depth_units=_depth_units, ) yield img diff --git a/tests/create_ros_fixture.py b/tests/create_ros_fixture.py index edb6fae6..ce81162e 100644 --- a/tests/create_ros_fixture.py +++ b/tests/create_ros_fixture.py @@ -1,241 +1,385 @@ -#!/usr/bin/env python -""" -Create synthetic ROS bag fixtures for the MVTB test suite. - -Usage:: - - $ python tests/create_ros_fixture.py - -Creates: - - tests/data/test_ros1.bag – ROS 1 Noetic bag - tests/data/test_ros2/ – ROS 2 Humble bag directory - -Contents (both bags carry the same messages): - - /camera/compressed sensor_msgs/CompressedImage 3 frames (JPEG, rgb) - /camera/raw sensor_msgs/Image 3 frames (rgb8 encoding) - /imu/data sensor_msgs/Imu 3 messages - -The image data are tiny (8 × 6 px) so the files stay small. The IMU topic -provides a non-image message for testing that the ``msgfilter`` correctly -accepts or rejects non-image types and that the generic pass-through path -attaches ``timestamp`` and ``topic`` attributes. - -Timestamps begin at 2023-11-14 22:13:20 UTC and step by 0.1 s. -""" - -from __future__ import annotations - -import io -import shutil -from pathlib import Path - -import numpy as np -from PIL import Image as PILImage - -from rosbags.rosbag1 import Writer as Writer1 -from rosbags.rosbag2 import Writer as Writer2 -from rosbags.typesys import Stores, get_typestore - -HERE = Path(__file__).parent -DATADIR = HERE / "data" -ROS1_BAG = DATADIR / "test_ros1.bag" -ROS2_BAG = DATADIR / "test_ros2" - -N = 3 # messages per topic -STAMP0_NS = 1_700_000_000_000_000_000 # 2023-11-14 22:13:20 UTC -STEP_NS = 100_000_000 # 0.1 s between messages -W, H = 8, 6 # image width, height (pixels) - - -def _jpeg_bytes(frame_idx: int) -> bytes: - """Return JPEG bytes for a small colour test image.""" - arr = np.zeros((H, W, 3), dtype=np.uint8) - arr[:, :, 0] = frame_idx * 40 # vary red channel per frame - arr[:, :, 1] = 128 - arr[:, :, 2] = 200 - buf = io.BytesIO() - PILImage.fromarray(arr, "RGB").save(buf, format="JPEG", quality=95) - return buf.getvalue() - - -def _rgb_bytes(frame_idx: int) -> bytes: - """Return raw rgb8 pixel bytes for a small test image.""" - arr = np.zeros((H, W, 3), dtype=np.uint8) - arr[:, :, 0] = frame_idx * 40 - arr[:, :, 1] = 100 - arr[:, :, 2] = 50 - return arr.tobytes() - - -def _make_ros1() -> None: - ts = get_typestore(Stores.ROS1_NOETIC) - - Header = ts.types["std_msgs/msg/Header"] - Time = ts.types["builtin_interfaces/msg/Time"] - Image = ts.types["sensor_msgs/msg/Image"] - CompressedImage = ts.types["sensor_msgs/msg/CompressedImage"] - Imu = ts.types["sensor_msgs/msg/Imu"] - Vec3 = ts.types["geometry_msgs/msg/Vector3"] - Quat = ts.types["geometry_msgs/msg/Quaternion"] - - ROS1_BAG.unlink(missing_ok=True) - with Writer1(ROS1_BAG) as writer: - conn_ci = writer.add_connection( - "/camera/compressed", - "sensor_msgs/msg/CompressedImage", - typestore=ts, - ) - conn_raw = writer.add_connection( - "/camera/raw", - "sensor_msgs/msg/Image", - typestore=ts, - ) - conn_imu = writer.add_connection( - "/imu/data", - "sensor_msgs/msg/Imu", - typestore=ts, - ) - - for i in range(N): - t_ns = STAMP0_NS + i * STEP_NS - sec = t_ns // 1_000_000_000 - nanosec = t_ns % 1_000_000_000 - stamp = Time(sec=sec, nanosec=nanosec) - header = Header(seq=i, stamp=stamp, frame_id="camera") - - # CompressedImage - ci = CompressedImage( - header=header, - format="jpeg", - data=np.frombuffer(_jpeg_bytes(i), dtype=np.uint8), - ) - writer.write( - conn_ci, - t_ns, - ts.serialize_ros1(ci, "sensor_msgs/msg/CompressedImage"), - ) - - # Raw Image (rgb8) - raw = Image( - header=header, - height=H, - width=W, - encoding="rgb8", - is_bigendian=0, - step=W * 3, - data=np.frombuffer(_rgb_bytes(i), dtype=np.uint8), - ) - writer.write( - conn_raw, - t_ns, - ts.serialize_ros1(raw, "sensor_msgs/msg/Image"), - ) - - # IMU - zero9 = np.zeros(9, dtype=np.float64) - imu = Imu( - header=header, - orientation=Quat(x=0.0, y=0.0, z=0.0, w=1.0), - orientation_covariance=zero9, - angular_velocity=Vec3(x=0.0, y=0.0, z=float(i) * 0.1), - angular_velocity_covariance=zero9, - linear_acceleration=Vec3(x=0.0, y=0.0, z=9.81), - linear_acceleration_covariance=zero9, - ) - writer.write( - conn_imu, - t_ns, - ts.serialize_ros1(imu, "sensor_msgs/msg/Imu"), - ) - - -def _make_ros2() -> None: - ts = get_typestore(Stores.ROS2_HUMBLE) - - Header = ts.types["std_msgs/msg/Header"] - Time = ts.types["builtin_interfaces/msg/Time"] - Image = ts.types["sensor_msgs/msg/Image"] - CompressedImage = ts.types["sensor_msgs/msg/CompressedImage"] - Imu = ts.types["sensor_msgs/msg/Imu"] - Vec3 = ts.types["geometry_msgs/msg/Vector3"] - Quat = ts.types["geometry_msgs/msg/Quaternion"] - - if ROS2_BAG.exists(): - shutil.rmtree(ROS2_BAG) - with Writer2(ROS2_BAG, version=9) as writer: - conn_ci = writer.add_connection( - "/camera/compressed", - "sensor_msgs/msg/CompressedImage", - typestore=ts, - ) - conn_raw = writer.add_connection( - "/camera/raw", - "sensor_msgs/msg/Image", - typestore=ts, - ) - conn_imu = writer.add_connection( - "/imu/data", - "sensor_msgs/msg/Imu", - typestore=ts, - ) - - for i in range(N): - t_ns = STAMP0_NS + i * STEP_NS - sec = t_ns // 1_000_000_000 - nanosec = t_ns % 1_000_000_000 - stamp = Time(sec=sec, nanosec=nanosec) - # ROS 2 Header has no seq field - header = Header(stamp=stamp, frame_id="camera") - - ci = CompressedImage( - header=header, - format="jpeg", - data=np.frombuffer(_jpeg_bytes(i), dtype=np.uint8), - ) - writer.write( - conn_ci, - t_ns, - ts.serialize_cdr(ci, "sensor_msgs/msg/CompressedImage"), - ) - - raw = Image( - header=header, - height=H, - width=W, - encoding="rgb8", - is_bigendian=0, - step=W * 3, - data=np.frombuffer(_rgb_bytes(i), dtype=np.uint8), - ) - writer.write( - conn_raw, - t_ns, - ts.serialize_cdr(raw, "sensor_msgs/msg/Image"), - ) - - zero9 = np.zeros(9, dtype=np.float64) - imu = Imu( - header=header, - orientation=Quat(x=0.0, y=0.0, z=0.0, w=1.0), - orientation_covariance=zero9, - angular_velocity=Vec3(x=0.0, y=0.0, z=float(i) * 0.1), - angular_velocity_covariance=zero9, - linear_acceleration=Vec3(x=0.0, y=0.0, z=9.81), - linear_acceleration_covariance=zero9, - ) - writer.write( - conn_imu, - t_ns, - ts.serialize_cdr(imu, "sensor_msgs/msg/Imu"), - ) - - -if __name__ == "__main__": - DATADIR.mkdir(exist_ok=True) - print(f"Writing ROS1 bag → {ROS1_BAG}") - _make_ros1() - print(f"Writing ROS2 bag → {ROS2_BAG}/") - _make_ros2() - print("Done.") +#!/usr/bin/env python +""" +Create synthetic ROS bag fixtures for the MVTB test suite. + +Usage:: + + $ python tests/create_ros_fixture.py + +Creates: + + tests/data/test_ros1.bag – ROS 1 Noetic bag + tests/data/test_ros1_realsense.bag – ROS 1 bag with RealSense SDK footers + tests/data/test_ros2/ – ROS 2 Humble bag directory + +Contents (both ``test_ros1`` and ``test_ros2`` carry the same messages): + + /camera/compressed sensor_msgs/CompressedImage 3 frames (JPEG, rgb) + /camera/raw sensor_msgs/Image 3 frames (rgb8 encoding) + /imu/data sensor_msgs/Imu 3 messages + +The image data are tiny (8 × 6 px) so the files stay small. The IMU topic +provides a non-image message for testing that the ``msgfilter`` correctly +accepts or rejects non-image types and that the generic pass-through path +attaches ``timestamp`` and ``topic`` attributes. + +``test_ros1_realsense.bag`` mimics a recording made by the Intel RealSense +SDK, which appends a 4-byte ``float32`` footer (the depth scale) to every +image message: + + /camera/depth/image_rect_raw sensor_msgs/Image 3 frames (16UC1) + /camera/color/image_raw sensor_msgs/Image 3 frames (rgb8) + /device_0/stream_info realsense_msgs/StreamInfo 1 message + +The ``realsense_msgs`` connection is what identifies the bag as a RealSense +recording. ``test_ros1_realsense_notrailer.bag`` holds the same topics without +the footer, matching bags such as Intel's ``single_depth_color_640x480.bag`` +where the messages are self-contained. + +Timestamps begin at 2023-11-14 22:13:20 UTC and step by 0.1 s. +""" + +from __future__ import annotations + +import io +import shutil +import struct +from pathlib import Path + +import numpy as np +from PIL import Image as PILImage + +from rosbags.rosbag1 import Writer as Writer1 +from rosbags.rosbag2 import Writer as Writer2 +from rosbags.typesys import Stores, get_typestore, get_types_from_msg + +HERE = Path(__file__).parent +DATADIR = HERE / "data" +ROS1_BAG = DATADIR / "test_ros1.bag" +ROS1_REALSENSE_BAG = DATADIR / "test_ros1_realsense.bag" +ROS1_REALSENSE_NOTRAILER_BAG = DATADIR / "test_ros1_realsense_notrailer.bag" +ROS2_BAG = DATADIR / "test_ros2" + +N = 3 # messages per topic +STAMP0_NS = 1_700_000_000_000_000_000 # 2023-11-14 22:13:20 UTC +STEP_NS = 100_000_000 # 0.1 s between messages +W, H = 8, 6 # image width, height (pixels) + +# Intel RealSense default depth scale, metres per depth unit +DEPTH_UNITS = 0.001 +COLOR_UNITS = 0.0 + +# Minimal realsense_msgs/StreamInfo definition: enough for the bag reader to +# recognise a RealSense recording without a RealSense installation. +STREAMINFO_MSGDEF = """\ +std_msgs/Header header +string name +uint32 width +uint32 height +uint32 encoding +uint32 fps +uint32 bytes_per_pixel +""" + + +def _jpeg_bytes(frame_idx: int) -> bytes: + """Return JPEG bytes for a small colour test image.""" + arr = np.zeros((H, W, 3), dtype=np.uint8) + arr[:, :, 0] = frame_idx * 40 # vary red channel per frame + arr[:, :, 1] = 128 + arr[:, :, 2] = 200 + buf = io.BytesIO() + PILImage.fromarray(arr, "RGB").save(buf, format="JPEG", quality=95) + return buf.getvalue() + + +def _rgb_bytes(frame_idx: int) -> bytes: + """Return raw rgb8 pixel bytes for a small test image.""" + arr = np.zeros((H, W, 3), dtype=np.uint8) + arr[:, :, 0] = frame_idx * 40 + arr[:, :, 1] = 100 + arr[:, :, 2] = 50 + return arr.tobytes() + + +def _depth_bytes(frame_idx: int) -> bytes: + """Return raw 16UC1 pixel bytes for a small test depth image.""" + arr = np.full((H, W), 1000 + frame_idx * 500, dtype=np.uint16) + arr[:, 0] = 900 + return arr.tobytes() + + +def _footer(depth_units: float) -> bytes: + """Return the 4-byte RealSense footer that follows an image message.""" + return struct.pack(" None: + """Write a ROS 1 bag that mimics a recording made by the RealSense SDK. + + The bag always carries a ``realsense_msgs/StreamInfo`` connection, which is + how the bag reader identifies RealSense recordings. When ``trailer`` is + true every image message is followed by the 4-byte little-endian + ``float32`` footer holding the depth scale, as written by the SDK. When it + is false the messages are self-contained, matching bag files that contain + RealSense topics but no trailer. + """ + ts = get_typestore(Stores.ROS1_NOETIC) + ts.register( + get_types_from_msg(STREAMINFO_MSGDEF, "realsense_msgs/msg/StreamInfo") + ) + + Header = ts.types["std_msgs/msg/Header"] + Time = ts.types["builtin_interfaces/msg/Time"] + Image = ts.types["sensor_msgs/msg/Image"] + StreamInfo = ts.types["realsense_msgs/msg/StreamInfo"] + + path.unlink(missing_ok=True) + with Writer1(path) as writer: + conn_depth = writer.add_connection( + "/camera/depth/image_rect_raw", + "sensor_msgs/msg/Image", + typestore=ts, + ) + conn_rgb = writer.add_connection( + "/camera/color/image_raw", + "sensor_msgs/msg/Image", + typestore=ts, + ) + conn_info = writer.add_connection( + "/device_0/stream_info", + "realsense_msgs/msg/StreamInfo", + typestore=ts, + ) + + for i in range(N): + t_ns = STAMP0_NS + i * STEP_NS + sec = t_ns // 1_000_000_000 + nanosec = t_ns % 1_000_000_000 + stamp = Time(sec=sec, nanosec=nanosec) + header = Header(seq=i, stamp=stamp, frame_id="camera") + + # Depth frame: 16-bit units, footer carries the depth scale + depth = Image( + header=header, + height=H, + width=W, + encoding="16UC1", + is_bigendian=0, + step=W * 2, + data=np.frombuffer(_depth_bytes(i), dtype=np.uint8), + ) + payload = bytes(ts.serialize_ros1(depth, "sensor_msgs/msg/Image")) + if trailer: + payload += _footer(DEPTH_UNITS) + writer.write(conn_depth, t_ns, payload) + + # Colour frame: footer is present but zero, as recorded by the SDK + rgb = Image( + header=header, + height=H, + width=W, + encoding="rgb8", + is_bigendian=0, + step=W * 3, + data=np.frombuffer(_rgb_bytes(i), dtype=np.uint8), + ) + payload = bytes(ts.serialize_ros1(rgb, "sensor_msgs/msg/Image")) + if trailer: + payload += _footer(COLOR_UNITS) + writer.write(conn_rgb, t_ns, payload) + + info = StreamInfo( + header=Header( + seq=0, + stamp=Time(sec=STAMP0_NS // 1_000_000_000, nanosec=0), + frame_id="camera", + ), + name="Depth", + width=W, + height=H, + encoding=0, + fps=30, + bytes_per_pixel=2, + ) + writer.write( + conn_info, + STAMP0_NS, + ts.serialize_ros1(info, "realsense_msgs/msg/StreamInfo"), + ) + + +def _make_ros1() -> None: + ts = get_typestore(Stores.ROS1_NOETIC) + + Header = ts.types["std_msgs/msg/Header"] + Time = ts.types["builtin_interfaces/msg/Time"] + Image = ts.types["sensor_msgs/msg/Image"] + CompressedImage = ts.types["sensor_msgs/msg/CompressedImage"] + Imu = ts.types["sensor_msgs/msg/Imu"] + Vec3 = ts.types["geometry_msgs/msg/Vector3"] + Quat = ts.types["geometry_msgs/msg/Quaternion"] + + ROS1_BAG.unlink(missing_ok=True) + with Writer1(ROS1_BAG) as writer: + conn_ci = writer.add_connection( + "/camera/compressed", + "sensor_msgs/msg/CompressedImage", + typestore=ts, + ) + conn_raw = writer.add_connection( + "/camera/raw", + "sensor_msgs/msg/Image", + typestore=ts, + ) + conn_imu = writer.add_connection( + "/imu/data", + "sensor_msgs/msg/Imu", + typestore=ts, + ) + + for i in range(N): + t_ns = STAMP0_NS + i * STEP_NS + sec = t_ns // 1_000_000_000 + nanosec = t_ns % 1_000_000_000 + stamp = Time(sec=sec, nanosec=nanosec) + header = Header(seq=i, stamp=stamp, frame_id="camera") + + # CompressedImage + ci = CompressedImage( + header=header, + format="jpeg", + data=np.frombuffer(_jpeg_bytes(i), dtype=np.uint8), + ) + writer.write( + conn_ci, + t_ns, + ts.serialize_ros1(ci, "sensor_msgs/msg/CompressedImage"), + ) + + # Raw Image (rgb8) + raw = Image( + header=header, + height=H, + width=W, + encoding="rgb8", + is_bigendian=0, + step=W * 3, + data=np.frombuffer(_rgb_bytes(i), dtype=np.uint8), + ) + writer.write( + conn_raw, + t_ns, + ts.serialize_ros1(raw, "sensor_msgs/msg/Image"), + ) + + # IMU + zero9 = np.zeros(9, dtype=np.float64) + imu = Imu( + header=header, + orientation=Quat(x=0.0, y=0.0, z=0.0, w=1.0), + orientation_covariance=zero9, + angular_velocity=Vec3(x=0.0, y=0.0, z=float(i) * 0.1), + angular_velocity_covariance=zero9, + linear_acceleration=Vec3(x=0.0, y=0.0, z=9.81), + linear_acceleration_covariance=zero9, + ) + writer.write( + conn_imu, + t_ns, + ts.serialize_ros1(imu, "sensor_msgs/msg/Imu"), + ) + + +def _make_ros2() -> None: + ts = get_typestore(Stores.ROS2_HUMBLE) + + Header = ts.types["std_msgs/msg/Header"] + Time = ts.types["builtin_interfaces/msg/Time"] + Image = ts.types["sensor_msgs/msg/Image"] + CompressedImage = ts.types["sensor_msgs/msg/CompressedImage"] + Imu = ts.types["sensor_msgs/msg/Imu"] + Vec3 = ts.types["geometry_msgs/msg/Vector3"] + Quat = ts.types["geometry_msgs/msg/Quaternion"] + + if ROS2_BAG.exists(): + shutil.rmtree(ROS2_BAG) + with Writer2(ROS2_BAG, version=9) as writer: + conn_ci = writer.add_connection( + "/camera/compressed", + "sensor_msgs/msg/CompressedImage", + typestore=ts, + ) + conn_raw = writer.add_connection( + "/camera/raw", + "sensor_msgs/msg/Image", + typestore=ts, + ) + conn_imu = writer.add_connection( + "/imu/data", + "sensor_msgs/msg/Imu", + typestore=ts, + ) + + for i in range(N): + t_ns = STAMP0_NS + i * STEP_NS + sec = t_ns // 1_000_000_000 + nanosec = t_ns % 1_000_000_000 + stamp = Time(sec=sec, nanosec=nanosec) + # ROS 2 Header has no seq field + header = Header(stamp=stamp, frame_id="camera") + + ci = CompressedImage( + header=header, + format="jpeg", + data=np.frombuffer(_jpeg_bytes(i), dtype=np.uint8), + ) + writer.write( + conn_ci, + t_ns, + ts.serialize_cdr(ci, "sensor_msgs/msg/CompressedImage"), + ) + + raw = Image( + header=header, + height=H, + width=W, + encoding="rgb8", + is_bigendian=0, + step=W * 3, + data=np.frombuffer(_rgb_bytes(i), dtype=np.uint8), + ) + writer.write( + conn_raw, + t_ns, + ts.serialize_cdr(raw, "sensor_msgs/msg/Image"), + ) + + zero9 = np.zeros(9, dtype=np.float64) + imu = Imu( + header=header, + orientation=Quat(x=0.0, y=0.0, z=0.0, w=1.0), + orientation_covariance=zero9, + angular_velocity=Vec3(x=0.0, y=0.0, z=float(i) * 0.1), + angular_velocity_covariance=zero9, + linear_acceleration=Vec3(x=0.0, y=0.0, z=9.81), + linear_acceleration_covariance=zero9, + ) + writer.write( + conn_imu, + t_ns, + ts.serialize_cdr(imu, "sensor_msgs/msg/Imu"), + ) + + +if __name__ == "__main__": + DATADIR.mkdir(exist_ok=True) + print(f"Writing ROS1 bag → {ROS1_BAG}") + _make_ros1() + print(f"Writing ROS1 RealSense bag → {ROS1_REALSENSE_BAG}") + _make_ros1_realsense(ROS1_REALSENSE_BAG) + print(f"Writing ROS1 RealSense bag (no trailer) → {ROS1_REALSENSE_NOTRAILER_BAG}") + _make_ros1_realsense(ROS1_REALSENSE_NOTRAILER_BAG, trailer=False) + print(f"Writing ROS2 bag → {ROS2_BAG}/") + _make_ros2() + print("Done.") diff --git a/tests/data/test_ros1_realsense.bag b/tests/data/test_ros1_realsense.bag new file mode 100644 index 00000000..f57aabbb Binary files /dev/null and b/tests/data/test_ros1_realsense.bag differ diff --git a/tests/data/test_ros1_realsense_notrailer.bag b/tests/data/test_ros1_realsense_notrailer.bag new file mode 100644 index 00000000..142bade5 Binary files /dev/null and b/tests/data/test_ros1_realsense_notrailer.bag differ diff --git a/tests/test_ros.py b/tests/test_ros.py index e4f8f4a1..74bae589 100644 --- a/tests/test_ros.py +++ b/tests/test_ros.py @@ -40,6 +40,8 @@ HERE = Path(__file__).parent ROS1_BAG = HERE / "data" / "test_ros1.bag" +ROS1_REALSENSE_BAG = HERE / "data" / "test_ros1_realsense.bag" +ROS1_REALSENSE_NOTRAILER_BAG = HERE / "data" / "test_ros1_realsense_notrailer.bag" ROS2_BAG = HERE / "data" / "test_ros2" _rosbags_available = ROS1_BAG.exists() and ROS2_BAG.exists() @@ -48,6 +50,11 @@ reason="ROS bag fixtures missing – run tests/create_ros_fixture.py first", ) +_skip_realsense_bags = pytest.mark.skipif( + not (ROS1_REALSENSE_BAG.exists() and ROS1_REALSENSE_NOTRAILER_BAG.exists()), + reason="RealSense bag fixtures missing – run tests/create_ros_fixture.py first", +) + _rosbags_importable = bool(getattr(_sources, "_rosbags_available", False)) if _rosbags_importable: from machinevisiontoolbox.Sources import ROSBag @@ -235,6 +242,74 @@ def test_format_local_time_returns_string(self): self.assertIn("2023", s) +# --------------------------------------------------------------------------- +# ROSBag – RealSense recordings (4-byte depth-units trailer) +# --------------------------------------------------------------------------- + + +@_skip_import +@_skip_realsense_bags +class TestRosBagRealSense(unittest.TestCase): + """RealSense SDK bags append a 4-byte footer to every image message. + + The footer holds the depth scale in metres, stored as a little-endian + float. Its position depends on the message layout, so it must be found + rather than assumed: other recordings of the same topics carry no footer + at all and would lose four bytes of pixel data if one were stripped. + """ + + DEPTH_TOPIC = "/camera/depth/image_rect_raw" + COLOR_TOPIC = "/camera/color/image_raw" + DEPTH_UNITS = 0.001 + + def _bag(self, path=ROS1_REALSENSE_BAG, **kwargs): + return ROSBag(str(path), release="noetic", **kwargs) + + # --- auto-detection and decoding --- + + def test_topics_are_detected(self): + topics = self._bag(msgfilter=None).topics() + self.assertIn(self.DEPTH_TOPIC, topics) + self.assertIn(self.COLOR_TOPIC, topics) + + def test_depth_images_decoded(self): + imgs = list(self._bag(topicfilter=self.DEPTH_TOPIC)) + self.assertEqual(len(imgs), N_FRAMES) + for img in imgs: + self.assertEqual(img.shape, (H, W)) + + def test_color_images_decoded(self): + imgs = list(self._bag(topicfilter=self.COLOR_TOPIC)) + self.assertEqual(len(imgs), N_FRAMES) + for img in imgs: + self.assertEqual(img.shape, (H, W, 3)) + + # --- footer handling --- + + def test_depth_units_read_from_footer(self): + for img in self._bag(topicfilter=self.DEPTH_TOPIC): + self.assertTrue(hasattr(img, "depth_units")) + self.assertAlmostEqual(img.depth_units, self.DEPTH_UNITS, places=7) + + def test_footer_of_colour_stream_is_zero(self): + for img in self._bag(topicfilter=self.COLOR_TOPIC): + self.assertEqual(img.depth_units, 0.0) + + # --- recordings whose messages are self-contained --- + + def test_untrailered_images_decoded(self): + bag = self._bag(ROS1_REALSENSE_NOTRAILER_BAG, topicfilter=self.DEPTH_TOPIC) + imgs = list(bag) + self.assertEqual(len(imgs), N_FRAMES) + for img in imgs: + self.assertEqual(img.shape, (H, W)) + + def test_untrailered_images_have_no_depth_units(self): + bag = self._bag(ROS1_REALSENSE_NOTRAILER_BAG, topicfilter=self.DEPTH_TOPIC) + for img in bag: + self.assertFalse(hasattr(img, "depth_units")) + + # --------------------------------------------------------------------------- # ROSBag – ROS 2 (same logical tests, different fixture path / release) # ---------------------------------------------------------------------------