diff --git a/src/machinevisiontoolbox/Sources.py b/src/machinevisiontoolbox/Sources.py index efe1366a..349ecc97 100644 --- a/src/machinevisiontoolbox/Sources.py +++ b/src/machinevisiontoolbox/Sources.py @@ -7,6 +7,7 @@ import fnmatch import io import json +import logging import base64 import pathlib import shutil @@ -2883,6 +2884,9 @@ class ROSBag(ImageSource): :param colororder: override the color-plane order for all image, or by topic ``{topic: colororder}`` :type colororder: str or dict or None + :param depth_units: ``None`` preserves native pixels, ``"auto"`` converts + depth to metres, or a positive float supplies metres per unit + :param invalid_is_nan: replace raw zero depth with NaN during conversion :param kwargs: options applied to image frames, see :func:`~machinevisiontoolbox.base.imageio.convert` :raises ImportError: if the ``rosbags`` package is not installed @@ -2896,6 +2900,21 @@ class ROSBag(ImageSource): Each yielded object carries a ``timestamp`` attribute (ROS nanosecond epoch from the message header) and a ``topic`` attribute (the topic on which it was published). + Image frames carry ``depth_units`` metadata in metres per native unit. + The scale precedence is a user float, a per-message footer, the bag's + ``.../option/Depth_Units/value`` topic, REP 118 encoding (``16UC1`` is + 0.001 and ``32FC1`` is 1.0), then unknown (``None``). Color footers + are reported as 0.0 and never scale color pixels. + + With ``depth_units=None`` the native dtype and values are preserved. + With ``"auto"`` or a float, depth pixels become float32 metres and raw + zeros become NaN when ``invalid_is_nan`` is true. An unknown scale + raises :class:`ValueError` when conversion is requested. For ``32FC1`` + a user scale acts as a gain. This metres-per-unit convention is the + reciprocal of Open3D and TUM RGB-D ``depth_scale`` in units per metre + so 1000 corresponds to 0.001 and 5000 corresponds to 0.0002. + Unknown image trailer lengths raise :class:`ValueError`. + **Filters** The message and topic filters can be a single string, or a list of strings. @@ -2969,6 +2988,8 @@ def __init__( dtype: np.dtype | str | dict | None = None, colororder: str | dict | None = None, verbose: bool = False, + depth_units: Literal["auto"] | float | None = None, + invalid_is_nan: bool = True, **kwargs: Any, ) -> None: if not _rosbags_available: @@ -2977,6 +2998,19 @@ def __init__( "Install it with: pip install rosbags " "or pip install machinevision-toolbox-python[ros]" ) + if depth_units is not None and depth_units != "auto": + if ( + isinstance(depth_units, (str, bool)) + or not np.isfinite(depth_units) + or depth_units <= 0 + ): + raise ValueError( + "depth_units must be None, 'auto', or a positive finite scale" + ) + depth_units = float(depth_units) + self.depth_units = depth_units + self.invalid_is_nan = invalid_is_nan + self._bag_depth_units = None self.filename = filename self.release = release self._topic_filter = topicfilter @@ -2990,7 +3024,6 @@ def __init__( self.args = kwargs self._tmpfile: str | None = None self._is_ros2: bool = False - self._trim_bytes: int = 0 # trailing bytes to strip (e.g. RealSense SDK footer) self._is_realsense: bool = False def __repr__(self) -> str: @@ -3131,7 +3164,7 @@ def _open_reader(self) -> Any: if self._tmpfile is None: with tempfile.NamedTemporaryFile(suffix=".bag", delete=False) as tmp: self._tmpfile = tmp.name - print(f"Downloading {filename} ...") + logging.getLogger(__name__).info(f"Downloading {filename} ...") urllib.request.urlretrieve(filename, self._tmpfile) path = self._tmpfile else: @@ -3159,13 +3192,15 @@ def _open_reader(self) -> Any: ) if self.verbose: all_conns = list(self.reader.connections) - print( + logging.getLogger(__name__).info( f" bag: {len(all_conns)} total connections, {len(self.connections)} pass filter " f"(topicfilter={self._topic_filter!r}, msgfilter={self._msgfilter!r})" ) for c in all_conns: mark = "✓" if self._allowed(c) else "✗" - print(f" {mark} {c.topic!r} [{c.msgtype}]") + logging.getLogger(__name__).info( + f" {mark} {c.topic!r} [{c.msgtype}]" + ) return self.reader def topics(self) -> dict[str, str]: @@ -3344,6 +3379,144 @@ def __exit__(self, exc_type, exc, tb) -> None: "ROS2_DASHING", ] + @staticmethod + def _deserialize_message( + store: Any, rawdata: bytes, msgtype: str, is_cdr: bool + ) -> tuple[Any, bytes]: + """Deserialize a message and retain any Image trailer. + + :param store: typestore defining the message + :param rawdata: serialized message bytes + :param msgtype: ROS message type + :param is_cdr: whether the message uses CDR encoding + :returns: message and bytes following its Image data array + """ + from rosbags.serde import SerdeError + + deserialize = store.deserialize_cdr if is_cdr else store.deserialize_ros1 + if msgtype not in {"sensor_msgs/msg/Image", "sensor_msgs/msg/CompressedImage"}: + # Probe each message independently, including StreamInfo and CameraInfo. + for trim in (0, 4, 8): + try: + return deserialize( + rawdata[:-trim] if trim else rawdata, msgtype + ), b"" + except (AssertionError, SerdeError): + if trim == 8: + raise + + origin = 4 if is_cdr else 0 + byteorder = "<" + if is_cdr: + if len(rawdata) < 4 or bytes(rawdata[:2]) not in {b"\x00\x00", b"\x00\x01"}: + raise ValueError("Invalid or unsupported CDR encapsulation") + byteorder = "<" if rawdata[1] else ">" + pos = origin + + def skip(size: int) -> None: + nonlocal pos + pos += size + if pos > len(rawdata): + raise ValueError("Truncated Image layout or data array") + + def uint32() -> int: + nonlocal pos + if is_cdr: + skip((-(pos - origin)) % 4) + value = struct.unpack_from(byteorder + "I", rawdata, pos)[0] + skip(4) + return value + + def string() -> None: + skip(uint32()) # CDR length includes the terminating NUL. + + # ROS 1 Header has a sequence number; ROS 2 Header does not. + if not is_cdr: + uint32() + uint32() # stamp seconds + uint32() # stamp nanoseconds + string() # frame_id + if msgtype == "sensor_msgs/msg/Image": + height = uint32() + uint32() # width + string() # encoding + skip(1) # is_bigendian + step = uint32() + n = uint32() + if n != height * step: + raise ValueError("Image data length does not equal height * step") + else: + string() # compression format + n = uint32() + skip(n) + trailer = bytes(rawdata[pos:]) + # Some CDR encoders pad the final array to a four-byte boundary. + if is_cdr and len(trailer) == (-pos) % 4 and not any(trailer): + trailer = b"" + return deserialize(rawdata[:pos], msgtype), trailer + + def _register_bag_types(self, store: Any) -> None: + """Register selected types and depth-unit metadata using their definition format.""" + from rosbags.interfaces import MessageDefinitionFormat + from rosbags.typesys import TypesysError, get_types_from_idl, get_types_from_msg + + for connection in getattr(self.reader, "connections", ()): + if not ( + self._allowed(connection) + or connection.topic.endswith("/option/Depth_Units/value") + ): + continue + if connection.msgtype not in store.types and connection.msgdef.data: + try: + if connection.msgdef.format == MessageDefinitionFormat.IDL: + definition = connection.msgdef.data + separator = "=" * 80 + "\n" + types = {} + if definition.startswith(separator + "IDL: "): + for part in definition.split(separator)[1:]: + _, idl = part.split("\n", 1) + types.update(get_types_from_idl(idl)) + else: + types = get_types_from_idl(definition) + elif connection.msgdef.format == MessageDefinitionFormat.MSG: + types = get_types_from_msg(connection.msgdef.data, connection.msgtype) + else: + continue + store.register({ + name: definition + for name, definition in types.items() + if name not in store.types + }) + except (TypesysError, ValueError) as exc: + logging.getLogger(__name__).warning( + "Could not register %s on %r: %s", + connection.msgtype, connection.topic, exc, + ) + + def _depth_image( + self, pixels: np.ndarray, encoding: str, topic: str, footer: float | None + ) -> tuple[np.ndarray, float | None]: + """Attach a native scale and optionally convert depth pixels to metres.""" + is_depth = encoding in {"16UC1", "32FC1"} or ( + pixels.ndim == 2 + and ("depth" in topic.lower() or (footer is not None and footer > 0)) + ) + if not is_depth: + return pixels, 0.0 if footer is not None else None + units = self.depth_units if isinstance(self.depth_units, float) else footer + if units is None or units <= 0: + units = self._bag_depth_units + if units is None: + units = {"16UC1": 0.001, "32FC1": 1.0}.get(encoding) + if self.depth_units is not None: + if units is None: + raise ValueError(f"Unknown depth units for {topic!r}") + converted = pixels.astype(np.float32) * np.float32(units) + if self.invalid_is_nan: + converted[pixels == 0] = np.nan + pixels = converted + return pixels, units + def _autodetect_typestore(self, rawdata: bytes, msgtype: str) -> bool: """Try all known typestores and keep the first one that works. @@ -3351,101 +3524,117 @@ def _autodetect_typestore(self, rawdata: bytes, msgtype: str) -> bool: (standard ROS 1 bag wire format). Second pass: tries ROS 2 stores with ``deserialize_cdr`` (CDR wire - format). This handles bags recorded by the Intel RealSense SDK, which - use a ROS 1 ``.bag`` container but CDR-encoded messages with ROS - 2-style type names (e.g. ``sensor_msgs/msg/Image``). When CDR is + format), including CDR messages in a ROS 1 container. When CDR is detected ``self._is_ros2`` is set to ``True`` so that all subsequent messages are decoded with ``deserialize_cdr``. :return: ``True`` if a matching store was found, ``False`` otherwise :rtype: bool """ + from rosbags.serde import SerdeError + from rosbags.typesys import TypesysError + assert _Stores is not None assert _get_typestore is not None 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) for name in self._ROS1_STORES_ORDERED + self._ROS2_STORES_ORDERED: if name not in available: continue store = _get_typestore(next(s for s in _Stores if s.name == name)) - for trim in trim_order: - raw = rawdata if trim == 0 else rawdata[:-trim] - try: - store.deserialize_ros1(raw, msgtype) - self.typestore = store - self.release = name - self._trim_bytes = trim - suffix = f", {trim}-byte trailer stripped" if trim else "" - 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. - continue - except (KeyError, UnicodeDecodeError, struct.error) as e: - if self.verbose: - print( - f" auto-detect: {name} (ros1) failed: {type(e).__name__}: {e}" - ) - break # wrong store entirely; move on - except Exception as e: - if self.verbose: - print( - f" auto-detect: {name} (ros1) UNEXPECTED {type(e).__name__}: {e}" - ) - break - else: - # All trim amounts exhausted with AssertionError (trailing bytes) + try: + self._register_bag_types(store) + self._deserialize_message(store, rawdata, msgtype, False) + self.typestore = store + self.release = name + self._is_ros2 = False + logging.getLogger(__name__).info( + f" auto-detect: {name} (ros1 encoding)" + ) + return True + except ( + AssertionError, SerdeError, TypesysError, KeyError, + UnicodeDecodeError, struct.error, ValueError, + ) as e: if self.verbose: - print( - f" auto-detect: {name} (ros1) failed: AssertionError (trailing bytes)" + logging.getLogger(__name__).info( + f" auto-detect: {name} (ros1) failed: {type(e).__name__}: {e}" ) # Second pass: CDR wire format (ROS 2 stores only). - # Handles RealSense SDK bags: .bag container but CDR-encoded messages. + # Also supports CDR-encoded messages in a ROS 1 container. for name in self._ROS2_STORES_ORDERED: if name not in available: continue store = _get_typestore(next(s for s in _Stores if s.name == name)) try: - store.deserialize_cdr(rawdata, msgtype) + self._register_bag_types(store) + self._deserialize_message(store, rawdata, msgtype, True) self.typestore = store self.release = name self._is_ros2 = True # use CDR for all subsequent messages - self._trim_bytes = 0 - print(f" auto-detect: {name} (CDR encoding)") + logging.getLogger(__name__).info( + f" auto-detect: {name} (CDR encoding)" + ) return True - except (AssertionError, KeyError, UnicodeDecodeError, struct.error) as e: + except ( + AssertionError, SerdeError, TypesysError, KeyError, + UnicodeDecodeError, struct.error, ValueError, + ) as e: if self.verbose: - print( + logging.getLogger(__name__).info( f" auto-detect: {name} (CDR) failed: {type(e).__name__}: {e}" ) continue - except Exception as e: - if self.verbose: - print( - f" auto-detect: {name} (CDR) UNEXPECTED {type(e).__name__}: {e}" - ) - continue if self.verbose: - print(f" auto-detect: rawdata[:16] = {rawdata[:16].hex()}") - print( + logging.getLogger(__name__).info( + f" auto-detect: rawdata[:16] = {rawdata[:16].hex()}" + ) + logging.getLogger(__name__).info( f" auto-detect: FAILED for msgtype={msgtype!r} — no matching typestore found" ) return False def __iter__(self) -> Iterator[Any]: + from rosbags.serde import SerdeError + self._open_reader() assert self.reader is not None _auto = self.typestore is None # release='auto' try: + if self.typestore is not None: + self._register_bag_types(self.typestore) + self._bag_depth_units = None + options = [ + c + for c in self.reader.connections + if c.topic.endswith("/option/Depth_Units/value") + ] + if options and _auto: + # A scalar CDR option can also look like a ROS 1 scalar plus + # a footer. Determine the wire layout from an image first. + images = [c for c in self.reader.connections + if c.msgtype in {"sensor_msgs/msg/Image", "sensor_msgs/msg/CompressedImage"}] + if images: + for conn, _, raw in self.reader.messages(connections=images): + if self._autodetect_typestore(raw, conn.msgtype): + _auto = False + break + if options and not _auto: + for conn, _, raw in self.reader.messages(connections=options): + try: + option, _ = self._deserialize_message( + self.typestore, raw, conn.msgtype, self._is_ros2 + ) + value = float(option.data) + except (AssertionError, SerdeError, KeyError, UnicodeDecodeError, + struct.error, ValueError, TypeError, AttributeError): + continue + if np.isfinite(value) and value > 0: + self._bag_depth_units = value + break for connection, timestamp, rawdata in self.reader.messages( connections=self.connections ): @@ -3454,32 +3643,40 @@ def __iter__(self) -> Iterator[Any]: # probe until we find a typestore that works if not self._autodetect_typestore(rawdata, connection.msgtype): if self.verbose: - print( + logging.getLogger(__name__).info( f"{self.format_local_time(timestamp)} Could not auto-detect release for " f"{connection.msgtype} on {connection.topic!r} — skipping" ) continue _auto = False # found it, use for all subsequent messages - _deser = ( - self.typestore.deserialize_cdr - if self._is_ros2 - else self.typestore.deserialize_ros1 + msg, trailer = self._deserialize_message( + self.typestore, rawdata, connection.msgtype, self._is_ros2 ) - _raw = ( - rawdata - if self._trim_bytes == 0 - else rawdata[: -self._trim_bytes] - ) - msg = cast(Any, _deser(_raw, connection.msgtype)) + depth_units = None + if trailer: + if not ( + len(trailer) == 4 + and self._is_realsense + and isinstance(self.reader, _RosBagReader1) + and connection.msgtype == "sensor_msgs/msg/Image" + ): + raise ValueError( + f"Unexpected Image trailer on {connection.topic!r}: " + f"{len(trailer)} bytes; expected CDR alignment padding " + "or a 4-byte RealSense ROS1 Image footer" + ) + # librealsense appends a float32 after the Image data array. + byteorder = ">" if self._is_ros2 and rawdata[1] == 0 else "<" + depth_units = struct.unpack(byteorder + "f", trailer)[0] except KeyError as e: if self.verbose: - print( + logging.getLogger(__name__).info( f"{self.format_local_time(timestamp)} Error occurred while deserializing message: {e}" ) continue except AssertionError: if self.verbose: - print( + logging.getLogger(__name__).info( f"{self.format_local_time(timestamp)} Could not deserialize {connection.msgtype} on {connection.topic!r} " f"using release '{self.release}' — try --release auto or a specific release (melodic, noetic, …)" ) @@ -3511,6 +3708,15 @@ def __iter__(self) -> Iterator[Any]: else: colororder = self.colororder + img_array, depth_units = self._depth_image( + img_array, + getattr( + msg, "encoding", getattr(msg, "format", "").split(";")[0] + ), + connection.topic, + depth_units, + ) + if colororder is not None and not _wants_mono(self.args): img = _make_image( img_array, @@ -3518,14 +3724,21 @@ def __iter__(self) -> Iterator[Any]: **self.args, ) else: - img = _make_image(img_array, **self.args) + img = _make_image( + img_array, + dtype=True + if (depth_units is not None and depth_units > 0) + or "depth" in connection.topic.lower() + else None, + **self.args, + ) _set_sample_metadata( img, self._stamp_to_ns(msg.header.stamp), connection.topic ) + img.depth_units = depth_units yield img elif connection.msgtype.endswith("Image"): - if hasattr(msg, "encoding"): encoding = msg.encoding if encoding not in self._ROS_ENCODING_MAP: @@ -3563,6 +3776,15 @@ def __iter__(self) -> Iterator[Any]: else: img_array = data.reshape((msg.height, msg.width)) + img_array, depth_units = self._depth_image( + img_array, + getattr( + msg, "encoding", getattr(msg, "format", "").split(";")[0] + ), + connection.topic, + depth_units, + ) + if colororder is not None and not _wants_mono(self.args): img = _make_image( img_array, @@ -3570,10 +3792,18 @@ def __iter__(self) -> Iterator[Any]: **self.args, ) else: - img = _make_image(img_array, **self.args) + img = _make_image( + img_array, + dtype=True + if (depth_units is not None and depth_units > 0) + or "depth" in connection.topic.lower() + else None, + **self.args, + ) _set_sample_metadata( img, self._stamp_to_ns(msg.header.stamp), connection.topic ) + img.depth_units = depth_units yield img elif connection.msgtype.endswith("PointCloud2"): diff --git a/tests/create_ros_fixture.py b/tests/create_ros_fixture.py index edb6fae6..3655741f 100644 --- a/tests/create_ros_fixture.py +++ b/tests/create_ros_fixture.py @@ -10,6 +10,8 @@ tests/data/test_ros1.bag – ROS 1 Noetic bag tests/data/test_ros2/ – ROS 2 Humble bag directory + tests/data/test_realsense_trailer.bag – RealSense-style 4-byte footers + tests/data/test_realsense_plain.bag – RealSense-style without footers Contents (both bags carry the same messages): @@ -29,6 +31,7 @@ import io import shutil +import struct from pathlib import Path import numpy as np @@ -36,7 +39,7 @@ from rosbags.rosbag1 import Writer as Writer1 from rosbags.rosbag2 import Writer as Writer2 -from rosbags.typesys import Stores, get_typestore +from rosbags.typesys import Stores, get_typestore, get_types_from_msg HERE = Path(__file__).parent DATADIR = HERE / "data" @@ -50,7 +53,7 @@ def _jpeg_bytes(frame_idx: int) -> bytes: - """Return JPEG bytes for a small colour test image.""" + """Return JPEG bytes for a small color test image.""" arr = np.zeros((H, W, 3), dtype=np.uint8) arr[:, :, 0] = frame_idx * 40 # vary red channel per frame arr[:, :, 1] = 128 @@ -232,10 +235,135 @@ def _make_ros2() -> None: ) +def make_realsense( + path: Path, + trailer: bool = True, + *, + encoding: str = "16UC1", + footer: float = 0.002, + option: float | None = 0.003, + wire: str = "ros1", + other_trailer: int = 0, +) -> None: + """Write a tiny RealSense-style bag with depth, color and metadata. + + :param path: new bag path + :param trailer: append a float32 footer to image messages + :param encoding: depth image encoding + :param footer: per-message metres per unit + :param option: recorded depth-unit option, or None + :param wire: ros1, cdr_le or cdr_be + :param other_trailer: trailing bytes for non-image messages + """ + ts = get_typestore(Stores.ROS1_NOETIC if wire == "ros1" else Stores.ROS2_HUMBLE) + info_type = "realsense_msgs/msg/StreamInfo" + ts.register(get_types_from_msg("uint32 fps\nstring encoding\n", info_type)) + stamp = ts.types["builtin_interfaces/msg/Time"](sec=42, nanosec=123) + header_args = dict(stamp=stamp, frame_id="camera") + if wire == "ros1": + header_args["seq"] = 0 + header = ts.types["std_msgs/msg/Header"](**header_args) + + with Writer1(path) as writer: + + def write(topic: str, msgtype: str, msg: object, extra: bytes = b"") -> None: + conn = writer.add_connection(topic, msgtype, typestore=ts) + raw = ( + ts.serialize_ros1(msg, msgtype) + if wire == "ros1" + else ts.serialize_cdr(msg, msgtype, little_endian=wire == "cdr_le") + ) + writer.write(conn, STAMP0_NS, bytes(raw) + extra) + + write( + "/camera/info", + info_type, + ts.types[info_type](fps=30, encoding=encoding), + b"\0" * other_trailer, + ) + camera_type = "sensor_msgs/msg/CameraInfo" + camera = ( + ts.types[camera_type]( + header=header, + height=1, + width=2, + distortion_model="plumb_bob", + D=np.zeros(5), + K=np.eye(3).ravel(), + R=np.eye(3).ravel(), + P=np.zeros(12), + binning_x=0, + binning_y=0, + roi=ts.types["sensor_msgs/msg/RegionOfInterest"](0, 0, 0, 0, False), + ) + if wire == "ros1" + else ts.types[camera_type]( + header=header, + height=1, + width=2, + distortion_model="plumb_bob", + d=np.zeros(5), + k=np.eye(3).ravel(), + r=np.eye(3).ravel(), + p=np.zeros(12), + binning_x=0, + binning_y=0, + roi=ts.types["sensor_msgs/msg/RegionOfInterest"](0, 0, 0, 0, False), + ) + ) + write("/camera/calibration", camera_type, camera, b"\0" * other_trailer) + # Write the option after the images to exercise the independent scan. + for topic, enc, data in [ + ( + "/camera/depth", + encoding, + np.array( + [[0, 100]], + dtype={ + "16UC1": np.uint16, + "32FC1": np.float32, + "mono16": np.uint16, + }[encoding], + ), + ), + ("/camera/color", "rgb8", np.arange(6, dtype=np.uint8).reshape(1, 2, 3)), + ]: + msgtype = "sensor_msgs/msg/Image" + msg = ts.types[msgtype]( + header=header, + height=1, + width=2, + encoding=enc, + is_bigendian=0, + step=data.nbytes, + data=np.frombuffer(data.tobytes(), np.uint8), + ) + extra = ( + struct.pack( + ">f" if wire == "cdr_be" else " None: + """Preserve pixels and decode trailers independently for each message.""" + from rosbags.typesys import Stores, get_typestore + + is_cdr = wire != "ros1" + store = get_typestore(Stores.ROS2_HUMBLE if is_cdr else Stores.ROS1_NOETIC) + stamp = store.types["builtin_interfaces/msg/Time"](sec=42, nanosec=123) + header_args = dict(stamp=stamp, frame_id="camera_optical_frame") + if not is_cdr: + header_args["seq"] = 0 + header = store.types["std_msgs/msg/Header"](**header_args) + image_type = "sensor_msgs/msg/Image" + string_type = "std_msgs/msg/String" + connection = SimpleNamespace(msgtype=image_type, topic="/camera/raw") + other_connection = SimpleNamespace(msgtype=string_type, topic="/status") + pixels = np.arange(18, dtype=np.uint8).reshape(2, 3, 3) + depth = np.array([[0, 1, 255], [256, 1000, 65535]], dtype=np.uint16) + frames = [ + (pixels, "rgb8", None), + (pixels, "rgb8", 0.0), + (depth, "16UC1", 0.001), + (pixels, "rgb8", None), + ] + + frames = frames[first_frame:] + frames[:first_frame] + + def serialize(msg: object, msgtype: str) -> bytes: + if is_cdr: + return bytes( + store.serialize_cdr(msg, msgtype, little_endian=wire == "cdr_le") + ) + return bytes(store.serialize_ros1(msg, msgtype)) + + messages = [] + for data, encoding, units in frames: + msg = store.types[image_type]( + header=header, + height=2, + width=3, + encoding=encoding, + is_bigendian=0, + step=data.strides[0], + data=np.frombuffer(data.tobytes(), dtype=np.uint8), + ) + raw = serialize(msg, image_type) + if units is not None: + raw += struct.pack(">f" if wire == "cdr_be" else " None: + """Reject unknown trailers rather than silently dropping bytes.""" + from rosbags.typesys import Stores, get_typestore + + store = get_typestore(Stores.ROS1_NOETIC) + stamp = store.types["builtin_interfaces/msg/Time"](sec=42, nanosec=0) + header = store.types["std_msgs/msg/Header"](seq=0, stamp=stamp, frame_id="") + msgtype = "sensor_msgs/msg/Image" + msg = store.types[msgtype]( + header=header, + height=1, + width=2, + encoding="mono8", + is_bigendian=0, + step=2, + data=np.array([17, 29], dtype=np.uint8), + ) + raw = bytes(store.serialize_ros1(msg, msgtype)) + trailer + connection = SimpleNamespace(msgtype=msgtype, topic="/camera/raw") + bag = ROSBag("unused.bag", release=release) + reader = MagicMock() + reader.connections = [] + reader.messages.return_value = iter([(connection, STAMP0, raw)]) + bag.reader = reader + with pytest.raises(ValueError, match="Unexpected Image trailer.*expected CDR alignment padding"): + list(iter(bag)) + reader.close.assert_called_once() + + # --------------------------------------------------------------------------- # ROSMessage dataclass # --------------------------------------------------------------------------- @@ -1027,3 +1138,335 @@ def _run() -> None: if __name__ == "__main__": unittest.main() + + +@_skip_import +@pytest.mark.parametrize("wire", ["ros1", "cdr_le", "cdr_be"]) +@pytest.mark.parametrize("trailer", [False, True]) +@pytest.mark.parametrize("other_trailer", [0, 4, 8]) +def test_realsense_bag( + tmp_path: Path, wire: str, trailer: bool, other_trailer: int +) -> None: + """StreamInfo must not cause valid image bytes to be trimmed.""" + from create_ros_fixture import make_realsense + + path = tmp_path / "realsense.bag" + make_realsense(path, trailer, wire=wire, other_trailer=other_trailer, option=None) + bag = ROSBag(path, release="auto", msgfilter=None) + with bag: + assert bag._is_realsense + samples = list(bag) + assert len(samples) == 4 + assert samples[0].fps == 30 + assert samples[1].width == 2 + depth, color = samples[2:] + np.testing.assert_array_equal(depth.array, [[0, 100]]) + assert depth.array.dtype == np.uint16 + assert depth.depth_units == pytest.approx(0.002 if trailer else 0.001) + np.testing.assert_array_equal(color.array, np.arange(6).reshape(1, 2, 3)) + assert color.depth_units == (0.0 if trailer else None) + + +@_skip_import +@pytest.mark.parametrize("wire", ["ros1", "cdr_le", "cdr_be"]) +@pytest.mark.parametrize("mode", [None, "auto", 0.004]) +@pytest.mark.parametrize("invalid_is_nan", [False, True]) +@pytest.mark.parametrize( + "encoding,trailer,option,expected", + [ + ("16UC1", True, 0.003, 0.002), + ("16UC1", False, 0.003, 0.003), + ("16UC1", False, None, 0.001), + ("32FC1", False, None, 1.0), + ("mono16", False, None, None), + ], +) +def test_depth_units_api( + tmp_path: Path, + wire: str, + mode, + invalid_is_nan: bool, + encoding: str, + trailer: bool, + option, + expected, +) -> None: + """Scale precedence and conversion are independent of the recording SDK.""" + from create_ros_fixture import make_realsense + + path = tmp_path / "depth.bag" + make_realsense(path, trailer, encoding=encoding, option=option, wire=wire) + bag = ROSBag( + path, + release="auto", + depth_units=mode, + invalid_is_nan=invalid_is_nan, + topicfilter="/camera/depth", + ) + units = mode if isinstance(mode, float) else expected + if units is None and mode is not None: + with pytest.raises(ValueError, match="Unknown depth units"): + list(bag) + return + (frame,) = list(bag) + assert frame.depth_units == (None if units is None else pytest.approx(units)) + if mode is None: + np.testing.assert_array_equal(frame.array, [[0, 100]]) + assert frame.array.dtype == (np.float32 if encoding == "32FC1" else np.uint16) + else: + assert frame.array.dtype == np.float32 + assert frame.array[0, 1] == pytest.approx(100 * units) + if invalid_is_nan: + assert np.isnan(frame.array[0, 0]) + else: + assert frame.array[0, 0] == 0 + + +@_skip_import +@pytest.mark.parametrize("wire", ["ros1", "cdr_le", "cdr_be"]) +@pytest.mark.parametrize("msgtype", ["Image", "CompressedImage"]) +@pytest.mark.parametrize("footer", [b"", struct.pack(" None: + """Both layouts locate the array end without consulting generated internals.""" + from rosbags.typesys import Stores, get_typestore + + store = get_typestore(Stores.ROS1_NOETIC if wire == "ros1" else Stores.ROS2_HUMBLE) + args = dict(stamp=store.types["builtin_interfaces/msg/Time"](42, 0), frame_id="ab") + if wire == "ros1": + args["seq"] = 0 + header = store.types["std_msgs/msg/Header"](**args) + fulltype = "sensor_msgs/msg/" + msgtype + data = np.arange(6, dtype=np.uint8) + args = dict(header=header, data=data) + if msgtype == "Image": + args.update(height=2, width=3, encoding="mono8", is_bigendian=0, step=3) + else: + args["format"] = "jpeg" + msg = store.types[fulltype](**args) + + def serialize(message: object) -> bytes: + return bytes( + store.serialize_ros1(message, fulltype) + if wire == "ros1" + else store.serialize_cdr(message, fulltype, little_endian=wire == "cdr_le") + ) + + decoded, actual_footer = ROSBag._deserialize_message( + store, serialize(msg) + footer, fulltype, wire != "ros1" + ) + np.testing.assert_array_equal(decoded.data, data) + assert actual_footer == footer + with pytest.raises((ValueError, struct.error), match="Truncated|buffer"): + ROSBag._deserialize_message( + store, serialize(msg)[:-1], fulltype, wire != "ros1" + ) + if msgtype == "Image": + msg.step = 4 + with pytest.raises(ValueError, match="height \\* step"): + ROSBag._deserialize_message(store, serialize(msg), fulltype, wire != "ros1") + + +@_skip_import +@pytest.mark.parametrize("units", [0.0, -0.001, float("nan"), float("inf"), "bad", True]) +def test_invalid_depth_units(units) -> None: + """Reject invalid user scales before opening a bag.""" + with pytest.raises(ValueError, match="depth_units"): + ROSBag("unused.bag", depth_units=units) + + +@_skip_import +@pytest.mark.parametrize("mode", [None, "auto", 0.005]) +def test_depth_conversion_leaves_color_unchanged(tmp_path: Path, mode) -> None: + """Depth scale options never convert color images.""" + from create_ros_fixture import make_realsense + + path = tmp_path / "color.bag" + make_realsense(path) + color, = list(ROSBag(path, topicfilter="/camera/color", depth_units=mode)) + np.testing.assert_array_equal(color.array, np.arange(6).reshape(1, 2, 3)) + assert color.array.dtype == np.uint8 + assert color.depth_units == 0.0 + + +@_skip_import +@pytest.mark.parametrize("little_endian", [False, True]) +@pytest.mark.parametrize("msgtype", ["Image", "CompressedImage"]) +@pytest.mark.parametrize("padding", [0, 1, 2, 3]) +@pytest.mark.parametrize("trailer_kind", ["padding", "nonzero", "misaligned", "footer"]) +def test_cdr_image_padding(tmp_path, little_endian, msgtype, padding, trailer_kind) -> None: + """Final CDR padding is not depth metadata, including three-pixel mono8.""" + from rosbags.rosbag2 import Writer + from rosbags.typesys import Stores, get_typestore + + store = get_typestore(Stores.ROS2_HUMBLE) + header = store.types["std_msgs/msg/Header"]( + store.types["builtin_interfaces/msg/Time"](42, 0), "camera" + ) + pixels = np.arange(4 - padding, dtype=np.uint8).reshape(1, -1) + fulltype = "sensor_msgs/msg/" + msgtype + if msgtype == "Image": + msg = store.types[fulltype](header, 1, pixels.size, "mono8", 0, pixels.size, pixels.ravel()) + else: + import cv2 + + ok, data = cv2.imencode(".png", pixels) + assert ok + # PNG decoders permit bytes after IEND; vary array length to exercise + # every alignment independently of the compression implementation. + data = np.concatenate((data, np.zeros((-padding - data.size) % 4, np.uint8))) + msg = store.types[fulltype](header, "png", data) + raw = bytes(store.serialize_cdr(msg, fulltype, little_endian=little_endian)) + assert (-len(raw)) % 4 == padding + path = tmp_path / "padded" + with Writer(path, version=9) as writer: + conn = writer.add_connection("/camera/image", fulltype, typestore=store) + trailer = { + "padding": b"\0" * padding, + "nonzero": b"\1" * max(1, padding), + "misaligned": b"\0" * ((padding + 1) % 4 or 1), + "footer": struct.pack(" None: + """IDL definitions are registered only when selected and use the IDL parser.""" + from rosbags.rosbag2 import Writer + from rosbags.typesys import Stores, get_types_from_idl, get_typestore + + store = get_typestore(Stores.ROS2_HUMBLE) + msgtype = "custom_msgs/msg/Status" + dependencies = { + "none": "", + "custom": "module custom_msgs { module msg { struct Detail { long code; }; }; };\n", + "conflicting_builtin": ( + "module std_msgs { module msg { struct Header { " + "string frame_id; builtin_interfaces::msg::Time stamp; }; }; };\n" + ), + }[dependency] + fields = { + "none": "long value;", + "custom": "custom_msgs::msg::Detail detail; long value;", + "conflicting_builtin": "std_msgs::msg::Header header; long value;", + }[dependency] + definition = ( + "=" * 80 + "\nIDL: custom_msgs/msg/Status\n" + + "module custom_msgs { module msg { struct Status { " + fields + " }; }; };\n" + ) + if dependencies: + definition += "=" * 80 + "\nIDL: dependency\n" + dependencies + types = get_types_from_idl(definition.split("\n", 2)[2].split("=" * 80)[0]) + if dependency == "custom": + types.update(get_types_from_idl(dependencies)) + store.register(types) + header = store.types["std_msgs/msg/Header"]( + store.types["builtin_interfaces/msg/Time"](42, 0), "camera" + ) + image_type = "sensor_msgs/msg/Image" + msg = store.types[image_type](header, 1, 3, "mono8", 0, 3, np.arange(3, dtype=np.uint8)) + path = tmp_path / "idl" + with Writer(path, version=9) as writer: + image_conn = writer.add_connection("/camera/image", image_type, typestore=store) + conn = writer.add_connection( + "/status", msgtype, msgdef=definition, rihs01=store.hash_rihs01(msgtype) + ) + writer.write(image_conn, STAMP0, store.serialize_cdr(msg, image_type)) + if dependency == "custom": + status = store.types[msgtype](store.types["custom_msgs/msg/Detail"](23), 7) + elif dependency == "conflicting_builtin": + status = store.types[msgtype](header, 7) + else: + status = store.types[msgtype](7) + writer.write(conn, STAMP0 + 1, store.serialize_cdr(status, msgtype)) + bag = ROSBag(path, release=release, msgfilter=None, topicfilter="/camera" if filtered else None) + samples = list(bag) + np.testing.assert_array_equal(samples[0].array, [[0, 1, 2]]) + assert (msgtype in bag.typestore.types) == (not filtered) + if not filtered: + assert len(samples) == 2 + assert samples[1].value == 7 + if dependency == "custom": + assert samples[1].detail.code == 23 + assert "custom_msgs/msg/Detail" in bag.typestore.types + elif dependency == "conflicting_builtin": + assert samples[1].header.frame_id == "camera" + + +@_skip_import +@pytest.mark.parametrize("release", ["auto", "humble"]) +def test_bad_definition_preserves_readable_connections(tmp_path, release) -> None: + """One malformed definition must not block images or later custom types.""" + from rosbags.rosbag2 import Writer + from rosbags.typesys import Stores, get_types_from_msg, get_typestore + + store = get_typestore(Stores.ROS2_HUMBLE) + msgtype = "custom_msgs/msg/Status" + store.register(get_types_from_msg("int32 value", msgtype)) + path = tmp_path / "bad_definition" + with Writer(path, version=9) as writer: + writer.add_connection( + "/broken", "custom_msgs/msg/Broken", + msgdef="invalid definition", rihs01="0" * 64, + ) + header = store.types["std_msgs/msg/Header"]( + store.types["builtin_interfaces/msg/Time"](42, 0), "camera" + ) + image_type = "sensor_msgs/msg/Image" + image = store.types[image_type](header, 1, 1, "mono8", 0, 1, np.array([9], dtype=np.uint8)) + image_conn = writer.add_connection("/image", image_type, typestore=store) + writer.write(image_conn, STAMP0 - 1, store.serialize_cdr(image, image_type)) + conn = writer.add_connection("/status", msgtype, typestore=store) + writer.write(conn, STAMP0, store.serialize_cdr(store.types[msgtype](7), msgtype)) + samples = list(ROSBag(path, release=release, msgfilter=None)) + assert len(samples) == 2 + np.testing.assert_array_equal(samples[0].array, [[9]]) + assert samples[1].value == 7 + + +@_skip_import +def test_autodetect_propagates_unexpected_errors() -> None: + """Programming failures must not be reported as a typestore mismatch.""" + bag = ROSBag("unused.bag", release="auto") + with patch.object(bag, "_register_bag_types", side_effect=RuntimeError("sentinel")): + with pytest.raises(RuntimeError, match="sentinel"): + bag._autodetect_typestore(b"", "sensor_msgs/msg/Image") + + +@_skip_import +@pytest.mark.parametrize("release", ["auto", "noetic"]) +@pytest.mark.parametrize("bad_option", ["truncated", "non_numeric"]) +def test_unreadable_depth_option_is_ignored(tmp_path, release, bad_option) -> None: + """Optional metadata must not prevent readable images being returned.""" + from rosbags.rosbag1 import Writer + from rosbags.typesys import Stores, get_typestore + + store = get_typestore(Stores.ROS1_NOETIC) + header = store.types["std_msgs/msg/Header"]( + 0, store.types["builtin_interfaces/msg/Time"](42, 0), "camera" + ) + msgtype = "sensor_msgs/msg/Image" + msg = store.types[msgtype](header, 1, 1, "16UC1", 0, 2, np.array([1, 0], dtype=np.uint8)) + path = tmp_path / "bad_option.bag" + with Writer(path) as writer: + conn = writer.add_connection("/depth", msgtype, typestore=store) + writer.write(conn, STAMP0, store.serialize_ros1(msg, msgtype)) + option_type = "std_msgs/msg/String" + option = writer.add_connection("/camera/option/Depth_Units/value", option_type, typestore=store) + raw = store.serialize_ros1(store.types[option_type]("not a number"), option_type) + writer.write(option, STAMP0 + 1, raw[:2] if bad_option == "truncated" else raw) + bag = ROSBag(path, release=release) + (sample,) = list(bag) + assert bag._bag_depth_units is None + assert sample.depth_units == 0.001 + np.testing.assert_array_equal(sample.array, [[1]])