From 5e0484b1db8108cd39191fa2bd47f70202599a33 Mon Sep 17 00:00:00 2001 From: cc Date: Fri, 23 Oct 2026 02:11:18 -0700 Subject: [PATCH 01/10] refactor(messages): isolate image, calibration and time adapters --- dimos/message_utils/camera_info.py | 178 +++++++++++++ dimos/message_utils/image.py | 176 ++++++++++++- dimos/message_utils/test_camera_info.py | 149 +++++++++++ dimos/message_utils/test_image_compression.py | 187 ++++++++++++++ dimos/message_utils/test_image_views.py | 233 ++++++++++++++++++ dimos/message_utils/test_time.py | 82 ++++++ dimos/message_utils/time.py | 77 ++++++ 7 files changed, 1074 insertions(+), 8 deletions(-) create mode 100644 dimos/message_utils/camera_info.py create mode 100644 dimos/message_utils/test_camera_info.py create mode 100644 dimos/message_utils/test_image_compression.py create mode 100644 dimos/message_utils/test_image_views.py create mode 100644 dimos/message_utils/test_time.py create mode 100644 dimos/message_utils/time.py diff --git a/dimos/message_utils/camera_info.py b/dimos/message_utils/camera_info.py new file mode 100644 index 00000000000..c8040990805 --- /dev/null +++ b/dimos/message_utils/camera_info.py @@ -0,0 +1,178 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Calibration construction and array access for generated ROS CameraInfo values.""" + +from copy import deepcopy +import math +from pathlib import Path +import re +from typing import Any, Literal + +import numpy as np +from numpy.typing import NDArray +import yaml + +from dimos.messages.builtin_interfaces.msg import Time +from dimos.messages.sensor_msgs.msg import CameraInfo, RegionOfInterest +from dimos.messages.std_msgs.msg import Header + + +def _dimensions(width: int, height: int) -> None: + if any( + isinstance(value, bool) or not isinstance(value, int) or value <= 0 + for value in (width, height) + ): + raise ValueError("Calibration width and height must be positive integers") + + +def camera_info_from_intrinsics( + fx: float, fy: float, cx: float, cy: float, width: int, height: int, *, header: Header +) -> CameraInfo: + """Construct an undistorted pinhole calibration with an explicit source header.""" + _dimensions(width, height) + if not all(math.isfinite(value) for value in (fx, fy, cx, cy)) or fx <= 0 or fy <= 0: + raise ValueError("Intrinsics must be finite with positive focal lengths") + return CameraInfo( + header=header, + width=width, + height=height, + distortion_model="plumb_bob", + d=np.asarray([0.0] * 5, dtype=np.float64), + k=np.array([fx, 0, cx, 0, fy, cy, 0, 0, 1], dtype=np.float64), + r=np.array([1, 0, 0, 0, 1, 0, 0, 0, 1], dtype=np.float64), + p=np.array([fx, 0, cx, 0, 0, fy, cy, 0, 0, 0, 1, 0], dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), + ) + + +def camera_info_from_fov( + fov_degrees: float, + width: int, + height: int, + *, + header: Header, + axis: Literal["horizontal", "vertical"] = "vertical", +) -> CameraInfo: + """Construct a centered pinhole model with square pixels from an axis's FOV.""" + _dimensions(width, height) + if not math.isfinite(fov_degrees) or not 0 < fov_degrees < 180: + raise ValueError("Field of view must be finite and between 0 and 180 degrees") + if axis not in ("horizontal", "vertical"): + raise ValueError("Field-of-view axis must be horizontal or vertical") + extent = width if axis == "horizontal" else height + focal = extent / (2 * math.tan(math.radians(fov_degrees) / 2)) + return camera_info_from_intrinsics( + focal, focal, width / 2, height / 2, width, height, header=header + ) + + +def _matrix(data: dict[str, Any], name: str, rows: int, columns: int | None) -> list[float]: + matrix = data[name] + if not isinstance(matrix, dict) or matrix.get("rows") != rows: + raise ValueError(f"{name} must declare {rows} rows") + actual_columns = matrix.get("cols") + if ( + not isinstance(actual_columns, int) + or actual_columns < 0 + or (columns is not None and actual_columns != columns) + ): + raise ValueError(f"{name} has invalid columns") + values = matrix.get("data") + if not isinstance(values, list) or len(values) != rows * actual_columns: + raise ValueError(f"{name} data length does not match its dimensions") + converted = [float(value) for value in values] + if not all(math.isfinite(value) for value in converted): + raise ValueError(f"{name} must contain finite values") + return converted + + +def camera_info_from_yaml(path: str | Path, *, header: Header) -> CameraInfo: + """Read a ROS calibration YAML file, validating declared matrix dimensions.""" + with Path(path).open() as source: + data = yaml.safe_load(source) + if not isinstance(data, dict): + raise TypeError("Calibration YAML must contain a mapping") + width, height = data["image_width"], data["image_height"] + _dimensions(width, height) + return CameraInfo( + header=header, + width=width, + height=height, + distortion_model=data["distortion_model"], + k=np.asarray(_matrix(data, "camera_matrix", 3, 3), dtype=np.float64), + d=np.asarray(_matrix(data, "distortion_coefficients", 1, None), dtype=np.float64), + r=np.asarray(_matrix(data, "rectification_matrix", 3, 3), dtype=np.float64), + p=np.asarray(_matrix(data, "projection_matrix", 3, 4), dtype=np.float64), + binning_x=0, + binning_y=0, + roi=RegionOfInterest(x_offset=0, y_offset=0, height=0, width=0, do_rectify=False), + ) + + +def intrinsic_matrix(message: CameraInfo) -> NDArray[np.float64]: + """Copy the 3x3 intrinsic matrix into independently mutable NumPy storage.""" + return np.array(message.k, dtype=np.float64, copy=True).reshape(3, 3) + + +def camera_info_with_stamp(message: CameraInfo, stamp: Time) -> CameraInfo: + """Copy a calibration with an exact source stamp, without mutating prior publishes.""" + result = deepcopy(message) + result.header.stamp = deepcopy(stamp) + return result + + +class CalibrationProvider: + """Lazily load named YAML calibrations as plain generated CameraInfo values. + + Both SingleWebcam and single_webcam resolve single_webcam.yaml. This + helper lives outside generated messages; calibration files stay local. + """ + + def __init__(self, calibration_dir: str | Path, *, frame_id: str = "camera_optical") -> None: + self._calibration_dir = Path(calibration_dir) + self._frame_id = frame_id + self._cache: dict[str, CameraInfo] = {} + + @staticmethod + def _to_snake_case(name: str) -> str: + first = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", name) + return re.sub("([a-z0-9])([A-Z])", r"\1_\2", first).lower() + + def __getattr__(self, name: str) -> CameraInfo: + if not re.fullmatch(r"[A-Za-z][A-Za-z0-9_]*", name): + raise AttributeError(f"No calibration file found for: {name}") + snake = self._to_snake_case(name) + for key in (name, snake): + if key in self._cache: + return self._cache[key] + for key in (name, snake): + path = self._calibration_dir / f"{key}.yaml" + if path.is_file(): + value = camera_info_from_yaml( + path, header=Header(frame_id=self._frame_id, stamp=Time(sec=0, nanosec=0)) + ) + self._cache[name] = value + self._cache[snake] = value + return value + raise AttributeError(f"No calibration file found for: {name}") + + def __dir__(self) -> list[str]: + names: set[str] = set() + for path in self._calibration_dir.glob("*.yaml"): + names.add(path.stem) + names.add("".join(word.capitalize() for word in path.stem.split("_"))) + return sorted(names) diff --git a/dimos/message_utils/image.py b/dimos/message_utils/image.py index 08cfb6296bf..16548a5f175 100644 --- a/dimos/message_utils/image.py +++ b/dimos/message_utils/image.py @@ -12,13 +12,18 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Read and convert generated image buffers.""" +"""Array views and compression for generated ROS image messages.""" +from __future__ import annotations + +from pathlib import Path from typing import Any +import cv2 import imagecodecs import numpy as np from numpy.typing import NDArray +from PIL import Image as PILImage from dimos.messages.builtin_interfaces.msg import Time from dimos.messages.sensor_msgs.msg import CompressedImage, Image @@ -39,9 +44,7 @@ } -def _image_from_array( - pixels: NDArray[Any], *, encoding: str, header: Header | None = None -) -> Image: +def image_from_array(pixels: NDArray[Any], *, encoding: str, header: Header | None = None) -> Image: """Copy an array into a generated image with an explicit pixel encoding.""" if encoding not in _FORMATS: raise ValueError(f"unsupported image encoding {encoding!r}") @@ -85,9 +88,88 @@ def image_view(msg: Image) -> NDArray[Any]: return view +def image_to_jpeg(msg: Image, quality: int = 75) -> bytes: + """Encode an 8-bit gray or color image using OpenCV's JPEG encoder.""" + + if isinstance(quality, bool) or not 0 <= quality <= 100: + raise ValueError("JPEG quality must be in 0..100") + pixels = image_view(msg) + conversions = { + "rgb8": cv2.COLOR_RGB2BGR, + "rgba8": cv2.COLOR_RGBA2BGR, + "bgra8": cv2.COLOR_BGRA2BGR, + } + if msg.encoding in conversions: + pixels = cv2.cvtColor(pixels, conversions[msg.encoding]) + elif msg.encoding not in ("bgr8", "mono8"): + raise ValueError(f"JPEG requires 8-bit color or gray, got {msg.encoding!r}") + ok, encoded = cv2.imencode(".jpg", pixels, [cv2.IMWRITE_JPEG_QUALITY, quality]) + if not ok: + raise ValueError("JPEG encoding failed") + return bytes(encoded) + + +def compressed_image_from_image( + message: Image, + quality: int = 75, + *, + format: str = "jpeg", + max_width: int | None = None, + effort: int | None = None, +) -> CompressedImage: + """Compress generated pixels, retaining the exact header. + + PNG preserves uint16 depth. JXL preserves uint16/float32 depth losslessly; + it requires imagecodecs and is an explicit codec extension for consumers. + """ + if not isinstance(message, Image): + raise TypeError("compression expects a generated Image") + if max_width is not None: + message, _ = image_resize_to_fit(message, max_width, max_width) + if format == "jpeg": + output_encoding = "mono8" if message.encoding == "mono8" else "bgr8" + data = image_to_jpeg(message, quality=quality) + elif format == "png": + pixels = image_view(message) + if pixels.dtype.kind != "u" or pixels.dtype.itemsize not in (1, 2): + raise ValueError("PNG cannot encode floating-point depth") + if message.encoding in ("rgb8", "rgba8"): + code = cv2.COLOR_RGB2BGR if message.encoding == "rgb8" else cv2.COLOR_RGBA2BGRA + pixels = cv2.cvtColor(pixels, code) + pixels = np.ascontiguousarray(pixels, dtype=pixels.dtype.newbyteorder("=")) + ok, encoded = cv2.imencode(".png", pixels) + if not ok: + raise ValueError("PNG encoding failed") + data = bytes(encoded) + output_encoding = {"rgb8": "bgr8", "rgba8": "bgra8"}.get(message.encoding, message.encoding) + elif format == "jxl": + pixels = image_view(message) + if pixels.dtype.kind not in ("u", "f") or ( + pixels.dtype.kind, + pixels.dtype.itemsize, + ) not in (("u", 1), ("u", 2), ("f", 4)): + raise ValueError(f"JXL cannot encode dtype {pixels.dtype}") + if message.encoding == "bgr8": + pixels = pixels[..., ::-1] + elif message.encoding == "bgra8": + pixels = pixels[..., [2, 1, 0, 3]] + pixels = np.ascontiguousarray(pixels, dtype=pixels.dtype.newbyteorder("=")) + if pixels.dtype == np.uint8: + data = bytes(imagecodecs.jpegxl_encode(pixels, level=quality, effort=effort or 3)) + else: + data = bytes(imagecodecs.jpegxl_encode(pixels, lossless=True, effort=effort or 1)) + output_encoding = {"bgr8": "rgb8", "bgra8": "rgba8"}.get(message.encoding, message.encoding) + else: + raise ValueError(f"unsupported compression format {format!r}") + return CompressedImage( + header=message.header, + format=f"{message.encoding}; {format} compressed {output_encoding}", + data=np.asarray(np.frombuffer(data, dtype=np.uint8).copy(), dtype=np.uint8), + ) + + def image_from_compressed(message: CompressedImage) -> Image: """Decode JPEG/PNG or the JXL extension into explicitly encoded generated pixels.""" - import cv2 format_name = message.format.lower() if "jxl" in format_name: @@ -106,7 +188,7 @@ def image_from_compressed(message: CompressedImage) -> Image: raise ValueError("unsupported decoded JXL layout") if encoding == "mono16" and format_name.startswith("16uc1;"): encoding = "16UC1" - return _image_from_array(jxl_pixels, encoding=encoding, header=message.header) + return image_from_array(jxl_pixels, encoding=encoding, header=message.header) if not any(name in format_name for name in ("jpeg", "jpg", "png")): raise ValueError(f"unsupported compressed image format {message.format!r}") pixels = cv2.imdecode(np.frombuffer(bytes(message.data), dtype=np.uint8), cv2.IMREAD_UNCHANGED) @@ -122,7 +204,64 @@ def image_from_compressed(message: CompressedImage) -> Image: raise ValueError("unsupported decoded image layout") if encoding == "mono16" and format_name.startswith("16uc1;"): encoding = "16UC1" - return _image_from_array(pixels, encoding=encoding, header=message.header) + return image_from_array(pixels, encoding=encoding, header=message.header) + + +def image_brightness(message: Image) -> float: + """Sample mean pixel intensity in [0, 1] for 8/16-bit image encodings. + + Respect row padding and endian. Floating-point depth has no normalized + intensity scale and is rejected rather than treated as a color image. + """ + pixels = image_view(message) + if pixels.dtype.kind != "u": + raise ValueError("brightness requires unsigned integer pixels") + stride = max(1, max(message.height, message.width) // 256) + return float(pixels[::stride, ::stride].mean() / np.iinfo(pixels.dtype).max) + + +def image_sharpness(message: Image) -> float: + """Laplacian variance of an 8-bit visual image, downsampled to 160 pixels wide.""" + + pixels = image_view(message) + codes = { + "rgb8": cv2.COLOR_RGB2GRAY, + "bgr8": cv2.COLOR_BGR2GRAY, + "rgba8": cv2.COLOR_RGBA2GRAY, + "bgra8": cv2.COLOR_BGRA2GRAY, + } + if message.encoding == "mono8": + gray = pixels + elif message.encoding in codes: + gray = cv2.cvtColor(pixels, codes[message.encoding]) + else: + raise ValueError("Sharpness requires an 8-bit visual image encoding") + if gray.size == 0: + raise ValueError("Sharpness requires a nonempty image") + if message.width > 160: + gray = cv2.resize( + gray, + (160, max(1, round(message.height * 160 / message.width))), + interpolation=cv2.INTER_AREA, + ) + return float(cv2.Laplacian(gray, cv2.CV_64F).var()) + + +def image_to_bgr(message: Image) -> NDArray[np.uint8]: + """Return an independent BGR8 array for drawing or OpenCV color operations.""" + + pixels = image_view(message) + if message.encoding == "bgr8": + return pixels.copy() + codes = { + "rgb8": cv2.COLOR_RGB2BGR, + "rgba8": cv2.COLOR_RGBA2BGR, + "bgra8": cv2.COLOR_BGRA2BGR, + "mono8": cv2.COLOR_GRAY2BGR, + } + if message.encoding not in codes: + raise ValueError(f"Cannot convert {message.encoding!r} to BGR8") + return np.asarray(cv2.cvtColor(pixels, codes[message.encoding]), dtype=np.uint8) def image_to_rgb(message: Image) -> NDArray[np.uint8]: @@ -131,7 +270,6 @@ def image_to_rgb(message: Image) -> NDArray[np.uint8]: Alpha is discarded. Floating-point depth needs an explicit visualization scale and is rejected rather than assigning pixel values to meters. """ - import cv2 pixels = image_view(message) if message.encoding == "rgb8": @@ -149,3 +287,25 @@ def image_to_rgb(message: Image) -> NDArray[np.uint8]: if message.encoding not in codes: raise ValueError(f"Cannot convert {message.encoding!r} to RGB8") return np.asarray(cv2.cvtColor(pixels, codes[message.encoding]), dtype=np.uint8) + + +def image_resize_to_fit(message: Image, max_width: int, max_height: int) -> tuple[Image, float]: + """Downscale a generated image while preserving encoding, aspect and exact header.""" + + if min(max_width, max_height, message.width, message.height) <= 0: + raise ValueError("Image and target dimensions must be positive") + if message.width <= max_width and message.height <= max_height: + return message, 1.0 + scale = min(max_width / message.width, max_height / message.height) + pixels = cv2.resize( + image_view(message), + (max(1, int(message.width * scale)), max(1, int(message.height * scale))), + interpolation=cv2.INTER_LINEAR, + ) + return image_from_array(pixels, encoding=message.encoding, header=message.header), scale + + +def image_from_file(path: str | Path, *, header: Header | None = None) -> Image: + """Decode an image file as generated RGB8 with an explicit optional source header.""" + with PILImage.open(path) as image: + return image_from_array(np.asarray(image.convert("RGB")), encoding="rgb8", header=header) diff --git a/dimos/message_utils/test_camera_info.py b/dimos/message_utils/test_camera_info.py new file mode 100644 index 00000000000..bcef99230c3 --- /dev/null +++ b/dimos/message_utils/test_camera_info.py @@ -0,0 +1,149 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Calibration helpers preserve generated ROS fields and reject malformed input.""" + +import math +from pathlib import Path + +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +import numpy as np +import pytest +import yaml + +from dimos.message_utils.camera_info import ( + camera_info_from_fov, + camera_info_from_intrinsics, + camera_info_from_yaml, + camera_info_with_stamp, + intrinsic_matrix, +) +from dimos.messages.builtin_interfaces.msg import Time +from dimos.messages.sensor_msgs.msg import CameraInfo +from dimos.messages.std_msgs.msg import Header + +_CALIBRATION = Path(__file__).parents[1] / "robot/unitree/go2/front_camera_720.yaml" + + +def test_real_go2_calibration_survives_cdr_with_exact_header(): + header = Header(frame_id="camera_optical", stamp=Time(sec=1700000000, nanosec=123456789)) + source = yaml.safe_load(_CALIBRATION.read_text()) + message = camera_info_from_yaml(_CALIBRATION, header=header) + decoded = cdr_decode(cdr_encode(message), CameraInfo) + assert decoded.header == header + assert (decoded.width, decoded.height) == (1280, 720) + assert decoded.distortion_model == "equidistant" + for field, key in [ + ("k", "camera_matrix"), + ("d", "distortion_coefficients"), + ("r", "rectification_matrix"), + ("p", "projection_matrix"), + ]: + np.testing.assert_array_equal(getattr(decoded, field), source[key]["data"]) + matrix = intrinsic_matrix(decoded) + assert matrix.shape == (3, 3) + matrix[0, 0] = 0 + assert decoded.k[0] == source["camera_matrix"]["data"][0] + decoded.header.frame_id = "changed" + assert header.frame_id == "camera_optical" + + +@pytest.mark.parametrize("axis,focal", [("vertical", 240), ("horizontal", 320)]) +def test_fov_converts_selected_axis_and_builds_projection(axis, focal): + message = camera_info_from_fov( + 90, 640, 480, header=Header(frame_id="optical", stamp=Time(sec=0, nanosec=0)), axis=axis + ) + np.testing.assert_allclose(message.k, [focal, 0, 320, 0, focal, 240, 0, 0, 1]) + np.testing.assert_allclose(message.p, [focal, 0, 320, 0, 0, focal, 240, 0, 0, 0, 1, 0]) + np.testing.assert_array_equal(message.r, np.eye(3).reshape(-1)) + assert message.header.frame_id == "optical" + + +@pytest.mark.parametrize("fov", [0, -1, 180, 181, math.nan, math.inf]) +def test_fov_rejects_nonphysical_angles(fov): + with pytest.raises(ValueError, match="Field of view"): + camera_info_from_fov( + fov, 640, 480, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="") + ) + + +@pytest.mark.parametrize("width,height", [(0, 480), (640, -1), (True, 480), (1.5, 480)]) +def test_intrinsics_reject_invalid_dimensions(width, height): + with pytest.raises(ValueError, match="positive integers"): + camera_info_from_intrinsics( + 300, + 300, + 320, + 240, + width, + height, + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + + +@pytest.mark.parametrize("fx", [0, -1, math.nan, math.inf]) +def test_intrinsics_reject_invalid_focal_length(fx): + with pytest.raises(ValueError, match="positive focal"): + camera_info_from_intrinsics( + fx, 300, 320, 240, 640, 480, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="") + ) + + +@pytest.mark.parametrize("mutation", ["rows", "length", "nonfinite"]) +def test_yaml_rejects_invalid_camera_matrix(tmp_path, mutation): + data = yaml.safe_load(_CALIBRATION.read_text()) + matrix = data["camera_matrix"] + if mutation == "rows": + matrix["rows"] = 1 + elif mutation == "length": + matrix["data"].pop() + else: + matrix["data"][0] = math.nan + path = tmp_path / "invalid.yaml" + path.write_text(yaml.safe_dump(data)) + with pytest.raises(ValueError, match="camera_matrix"): + camera_info_from_yaml(path, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")) + + +def test_stamp_copy_preserves_calibration_and_does_not_mutate_prior_publish(): + original = camera_info_from_intrinsics( + 500, + 510, + 320, + 240, + 640, + 480, + header=Header(frame_id="optical", stamp=Time(sec=1, nanosec=987654321)), + ) + original.binning_x = 2 + original.roi.x_offset = 12 + stamp = Time(sec=2, nanosec=123456789) + copied = camera_info_with_stamp(original, stamp) + assert copied.header.stamp == stamp + assert copied.header.frame_id == original.header.frame_id + assert list(copied.k) == list(original.k) + assert copied.binning_x == 2 and copied.roi.x_offset == 12 + copied.k[0] = 100 + copied.roi.x_offset = 24 + stamp.nanosec = 0 + assert original.k[0] == 500 and original.roi.x_offset == 12 + assert original.header.stamp == Time(sec=1, nanosec=987654321) + assert copied.header.stamp.nanosec == 123456789 + + +def test_calibration_yaml_rejects_non_mapping_input(tmp_path): + source = tmp_path / "calibration.yaml" + source.write_text("- width\n- height\n") + with pytest.raises(TypeError, match="Calibration YAML must contain a mapping"): + camera_info_from_yaml(source, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")) diff --git a/dimos/message_utils/test_image_compression.py b/dimos/message_utils/test_image_compression.py new file mode 100644 index 00000000000..0514d681424 --- /dev/null +++ b/dimos/message_utils/test_image_compression.py @@ -0,0 +1,187 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import cv2 +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +import numpy as np +import pytest + +from dimos.message_utils.image import ( + compressed_image_from_image, + image_from_array, + image_from_compressed, + image_to_bgr, + image_to_rgb, + image_view, +) +from dimos.message_utils.time import time_from_seconds, to_nanoseconds +from dimos.messages.builtin_interfaces.msg import Time +from dimos.messages.sensor_msgs.msg import CompressedImage +from dimos.messages.std_msgs.msg import Header + + +@pytest.fixture() +def rgb_image(): + rng = np.random.RandomState(42) + gradient = np.linspace(0, 255, 640, dtype=np.uint8) + data = np.broadcast_to(gradient, (480, 640)).copy() + data = np.stack([data, data // 2, rng.randint(0, 50, (480, 640), dtype=np.uint8)], axis=-1) + return image_from_array( + data, encoding="rgb8", header=Header(frame_id="cam", stamp=time_from_seconds(1234.5678)) + ) + + +def test_jpeg_roundtrip_preserves_meta_and_pixels(rgb_image) -> None: + ci = compressed_image_from_image(rgb_image, quality=90) + assert "; jpeg compressed " in ci.format + assert ci.header.frame_id == "cam" + assert to_nanoseconds(ci.header.stamp) == to_nanoseconds(rgb_image.header.stamp) + assert 0 < len(ci.data) < image_view(rgb_image).nbytes // 4 + img = image_from_compressed(ci) + assert img.encoding in ("rgb8", "bgr8") + assert img.header.frame_id == "cam" + assert to_nanoseconds(img.header.stamp) == to_nanoseconds(rgb_image.header.stamp) + assert image_view(img).shape == image_view(rgb_image).shape + diff = np.abs(image_to_rgb(img).astype(int) - image_view(rgb_image).astype(int)).mean() + assert diff < 5, f"JPEG q90 mean pixel error too high: {diff}" + + +def test_png_roundtrip_is_lossless_bgr(rgb_image) -> None: + bgr = image_from_array(image_to_bgr(rgb_image), encoding="bgr8", header=rgb_image.header) + ci = compressed_image_from_image(bgr, format="png") + img = image_from_compressed(ci) + assert img.encoding == "bgr8" + assert np.array_equal(image_view(img), image_view(bgr)) + assert to_nanoseconds(img.header.stamp) == to_nanoseconds(bgr.header.stamp) + + +def test_png_roundtrip_is_lossless_gray16() -> None: + data = np.arange(100 * 80, dtype=np.uint16).reshape(100, 80) + src = image_from_array( + data, encoding="mono16", header=Header(frame_id="d", stamp=time_from_seconds(1.0)) + ) + img = image_from_compressed(compressed_image_from_image(src, format="png")) + assert img.encoding == "mono16" + assert np.array_equal(image_view(img), data) + + +def test_png_decode_alpha_is_bgra() -> None: + arr = np.zeros((10, 10, 4), dtype=np.uint8) + arr[..., 3] = 128 + ok, buf = cv2.imencode(".png", arr) + assert ok + img = image_from_compressed( + CompressedImage( + data=np.frombuffer(buf.tobytes(), dtype=np.uint8), + format="png", + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + ) + assert img.encoding == "bgra8" + assert np.array_equal(image_view(img), arr) + + +def test_lcm_roundtrip(rgb_image) -> None: + ci = compressed_image_from_image(rgb_image) + wire = cdr_encode(ci) + out = cdr_decode(wire, CompressedImage) + np.testing.assert_array_equal(out.data, ci.data) + assert "; jpeg compressed " in out.format + assert out.header.frame_id == "cam" + assert abs(to_nanoseconds(out.header.stamp) - to_nanoseconds(ci.header.stamp)) < 1e-06 + + +def test_jpeg_rejects_depth_formats() -> None: + depth = image_from_array( + np.zeros((10, 10), dtype=np.uint16), + encoding="mono16", + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + with pytest.raises(ValueError, match="JPEG requires 8-bit"): + compressed_image_from_image(depth) + + +def test_png_rejects_float_depth() -> None: + depth = image_from_array( + np.zeros((10, 10), dtype=np.float32), + encoding="32FC1", + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + with pytest.raises(ValueError, match="PNG cannot encode"): + compressed_image_from_image(depth, format="png") + + +def test_from_image_rejects_non_image() -> None: + with pytest.raises(TypeError): + compressed_image_from_image(b"not an image") + + +def test_max_width_resizes(rgb_image) -> None: + img = image_from_compressed(compressed_image_from_image(rgb_image, max_width=320)) + assert img.width <= 320 + assert img.height <= 320 + + +def test_jxl_roundtrip_is_lossless_gray16() -> None: + data = np.arange(100 * 80, dtype=np.uint16).reshape(100, 80) + src = image_from_array( + data, encoding="mono16", header=Header(frame_id="d", stamp=time_from_seconds(1.0)) + ) + ci = compressed_image_from_image(src, format="jxl") + assert "; jxl compressed " in ci.format + img = image_from_compressed(ci) + assert img.encoding == "mono16" + assert np.array_equal(image_view(img), data) + + +def test_jxl_roundtrip_is_lossless_float_depth() -> None: + data = np.linspace(0.1, 10.0, 100 * 80, dtype=np.float32).reshape(100, 80) + src = image_from_array( + data, encoding="32FC1", header=Header(frame_id="d", stamp=time_from_seconds(2.0)) + ) + ci = compressed_image_from_image(src, format="jxl") + img = image_from_compressed(ci) + assert img.encoding == "32FC1" + assert np.array_equal(image_view(img), data) + assert to_nanoseconds(img.header.stamp) == to_nanoseconds(src.header.stamp) + assert img.header.frame_id == "d" + + +def test_jxl_lossy_rgb_honors_quality(rgb_image) -> None: + ci = compressed_image_from_image(rgb_image, format="jxl", quality=90) + assert 0 < len(ci.data) < image_view(rgb_image).nbytes // 4 + img = image_from_compressed(ci) + assert img.encoding in ("rgb8", "bgr8") + diff = np.abs(image_to_rgb(img).astype(int) - image_view(rgb_image).astype(int)).mean() + assert diff < 5, f"JXL q90 mean pixel error too high: {diff}" + + +def test_jxl_effort_trades_cpu_for_size(rgb_image) -> None: + fast = compressed_image_from_image(rgb_image, format="jxl", effort=1) + thorough = compressed_image_from_image(rgb_image, format="jxl", effort=7) + assert len(thorough.data) <= len(fast.data) + + +def test_jxl_rejects_float64() -> None: + depth = image_from_array(np.zeros((10, 10), dtype=np.float64), encoding="64FC1") + with pytest.raises(ValueError, match="JXL cannot encode"): + compressed_image_from_image(depth, format="jxl") + + +@pytest.mark.parametrize("codec", ["png", "jxl"]) +def test_compressed_big_endian_depth(codec: str) -> None: + pixels = np.arange(80, dtype=">u2").reshape(10, 8) * 257 + source = image_from_array(pixels.astype(">u2"), encoding="16UC1") + decoded = image_from_compressed(compressed_image_from_image(source, format=codec)) + np.testing.assert_array_equal(image_view(decoded), pixels) diff --git a/dimos/message_utils/test_image_views.py b/dimos/message_utils/test_image_views.py new file mode 100644 index 00000000000..a4d0680bcaf --- /dev/null +++ b/dimos/message_utils/test_image_views.py @@ -0,0 +1,233 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Borrowed buffers preserve padding, byte order, lifetime, and ROS value semantics.""" + +import gc + +import cv2 +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +import numpy as np +from PIL import Image as PILImage +import pytest + +from dimos.message_utils.image import ( + image_brightness, + image_from_array, + image_from_file, + image_sharpness, + image_to_jpeg, + image_to_rgb, + image_view, +) +from dimos.message_utils.time import time_from_nanoseconds +from dimos.messages.builtin_interfaces.msg import Time +from dimos.messages.sensor_msgs.msg import Image +from dimos.messages.std_msgs.msg import Header + + +def test_padded_image_view_retains_owner_and_is_readonly() -> None: + msg = Image( + width=2, + height=2, + step=4, + encoding="mono8", + data=np.array([1, 2, 99, 99, 3, 4, 99, 99], dtype=np.uint8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + is_bigendian=0, + ) + pixels = image_view(msg) + np.testing.assert_array_equal(pixels, [[1, 2], [3, 4]]) + assert pixels.strides == (4, 1) + with pytest.raises(ValueError): + pixels[0, 0] = 9 + del msg + gc.collect() + np.testing.assert_array_equal(pixels, [[1, 2], [3, 4]]) + copied = pixels.copy() + copied[0, 0] = 9 + assert pixels[0, 0] == 1 + + +def test_big_endian_depth_view() -> None: + msg = Image( + width=2, + height=1, + step=4, + encoding="16UC1", + is_bigendian=1, + data=np.array([0x01, 0x02, 0x03, 0x04], dtype=np.uint8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + np.testing.assert_array_equal(image_view(msg), [[258, 772]]) + + +def test_image_array_copy_preserves_endian_and_handles_strides() -> None: + pixels = np.arange(24, dtype=">u2").reshape(4, 6)[::2, ::2] + msg = image_from_array(pixels, encoding="16UC1") + assert (msg.width, msg.height, msg.step, msg.is_bigendian) == (3, 2, 6, 1) + expected = pixels.copy() + pixels[:] = 0 + np.testing.assert_array_equal(image_view(msg), expected) + + +@pytest.mark.parametrize( + "pixels,encoding", + [(np.zeros((2, 2), dtype=np.float32), "mono8"), (np.zeros((2, 2), dtype=np.uint8), "rgb8")], +) +def test_image_array_requires_matching_encoding(pixels: np.ndarray, encoding: str) -> None: + with pytest.raises(ValueError, match="does not match"): + image_from_array(pixels, encoding=encoding) + + +@pytest.mark.parametrize("step,data", [(1, [1]), (2, [1]), (2, [1, 2, 3])]) +def test_invalid_image_layout(step: int, data: list[int]) -> None: + with pytest.raises(ValueError, match="dimensions"): + image_view( + Image( + width=2, + height=1, + step=step, + data=np.asarray(data, dtype=np.uint8), + encoding="mono8", + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + is_bigendian=0, + ) + ) + + +def test_jpeg_preserves_rgb_color() -> None: + msg = Image( + width=16, + height=16, + step=48, + encoding="rgb8", + data=np.asarray([255, 0, 0] * 256, dtype=np.uint8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + is_bigendian=0, + ) + jpeg = image_to_jpeg(msg, quality=95) + bgr = cv2.imdecode(np.frombuffer(jpeg, np.uint8), cv2.IMREAD_COLOR) + assert bgr.shape == (16, 16, 3) + assert bgr[8, 8, 2] > 250 and bgr[8, 8, 0] < 5 + + +@pytest.mark.parametrize("encoding", ["rgb8", "bgr8", "rgba8", "bgra8", "mono8"]) +def test_sharpness_prefers_edges_to_blurred_pixels(encoding): + mono = ((np.indices((64, 64)).sum(axis=0) % 2) * 255).astype(np.uint8) + channels = 1 if encoding == "mono8" else (4 if "a" in encoding else 3) + pixels = mono if channels == 1 else np.repeat(mono[:, :, None], channels, axis=2) + blurred = cv2.GaussianBlur(pixels, (5, 5), 0) + sharp = image_from_array(pixels, encoding=encoding) + soft = image_from_array(blurred, encoding=encoding) + assert image_sharpness(sharp) > image_sharpness(soft) + + +def test_sharpness_rejects_depth_and_empty_images(): + depth = image_from_array(np.zeros((2, 2), dtype=np.float32), encoding="32FC1") + with pytest.raises(ValueError, match="8-bit visual"): + image_sharpness(depth) + with pytest.raises(ValueError, match="nonempty"): + image_sharpness( + Image( + encoding="mono8", + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + height=0, + width=0, + is_bigendian=0, + step=0, + data=np.array([], dtype=np.uint8), + ) + ) + + +@pytest.mark.parametrize( + "encoding,pixel", + [ + ("rgb8", [255, 0, 17]), + ("bgr8", [17, 0, 255]), + ("rgba8", [255, 0, 17, 23]), + ("bgra8", [17, 0, 255, 23]), + ], +) +def test_rgb_model_input_is_color_correct_and_independent(encoding: str, pixel: list[int]) -> None: + source = image_from_array(np.array([[pixel]], dtype=np.uint8), encoding=encoding) + restored = cdr_decode(cdr_encode(source), Image) + rgb = image_to_rgb(restored) + np.testing.assert_array_equal(rgb, [[[255, 0, 17]]]) + rgb[0, 0] = 0 + np.testing.assert_array_equal(image_view(restored), [[pixel]]) + + +def test_rgb_model_input_respects_big_endian_padded_grayscale() -> None: + msg = Image( + width=2, + height=2, + step=6, + encoding="mono16", + is_bigendian=1, + data=np.array([0x12, 0x34, 0xFF, 0xFF, 99, 99, 0x01, 0x00, 0, 0, 99, 99], dtype=np.uint8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + np.testing.assert_array_equal( + image_to_rgb(msg), [[[18, 18, 18], [255, 255, 255]], [[1, 1, 1], [0, 0, 0]]] + ) + + +def test_rgb_model_input_requires_explicit_float_depth_scale() -> None: + msg = image_from_array(np.array([[2.5]], dtype=np.float32), encoding="32FC1") + with pytest.raises(ValueError, match="RGB8"): + image_to_rgb(msg) + + +def test_brightness_ignores_padding_and_normalizes_big_endian_u16() -> None: + message = Image( + width=2, + height=1, + step=6, + encoding="mono16", + is_bigendian=1, + data=np.array([0, 0, 255, 255, 255, 255], dtype=np.uint8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + assert image_brightness(message) == pytest.approx(0.5) + depth = image_from_array(np.array([[1.0]], dtype=np.float32), encoding="32FC1") + with pytest.raises(ValueError, match="unsigned integer"): + image_brightness(depth) + + +@pytest.mark.parametrize("mode", ["L", "RGBA"]) +def test_image_file_decodes_rgb_and_preserves_source_header(tmp_path, mode): + path = tmp_path / "frame.png" + pixels = np.array([[7, 19]], dtype=np.uint8) + source = PILImage.fromarray(pixels).convert(mode) + source.save(path) + header = Header(stamp=time_from_nanoseconds(1234567890), frame_id="camera") + message = image_from_file(path, header=header) + assert message.encoding == "rgb8" + assert cdr_encode(message.header) == cdr_encode(header) + np.testing.assert_array_equal(image_view(message), [[[7, 7, 7], [19, 19, 19]]]) + + +@pytest.mark.parametrize( + "encoding,dtype,shape", + [("64FC1", ">f8", (3, 4)), ("16SC1", ">i2", (3, 4)), ("32FC3", ">f4", (3, 4, 3))], +) +def test_remaining_raw_encodings_keep_signed_float_and_endian_values(encoding, dtype, shape): + pixels = (np.arange(np.prod(shape)).reshape(shape) - 5).astype(dtype) + source = image_from_array(pixels, encoding=encoding) + decoded = cdr_decode(cdr_encode(source), Image) + assert decoded.encoding == encoding + assert decoded.is_bigendian == 1 + np.testing.assert_array_equal(image_view(decoded), pixels) diff --git a/dimos/message_utils/test_time.py b/dimos/message_utils/test_time.py new file mode 100644 index 00000000000..9a5f2e8da5a --- /dev/null +++ b/dimos/message_utils/test_time.py @@ -0,0 +1,82 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +import pytest + +from dimos.message_utils.time import ( + duration_from_nanoseconds, + duration_from_seconds, + header_now, + time_from_nanoseconds, + time_from_seconds, + to_nanoseconds, + to_seconds, +) + + +@pytest.mark.parametrize("factory", [time_from_nanoseconds, duration_from_nanoseconds]) +@pytest.mark.parametrize("value", [0, -1, -1_250_000_000, 1_700_000_000_123_456_789]) +def test_integer_time_round_trip_preserves_every_nanosecond(factory, value): + stamp = factory(value) + + assert 0 <= stamp.nanosec < 1_000_000_000 + assert to_nanoseconds(stamp) == value + assert to_nanoseconds(cdr_decode(cdr_encode(stamp), stamp.__msgtype__)) == value + + +@pytest.mark.parametrize("factory", [time_from_seconds, duration_from_seconds]) +@pytest.mark.parametrize("value,expected", [(-1.25, (-2, 750_000_000)), (0.9999999996, (1, 0))]) +def test_float_conversion_normalizes_negative_values_and_fractional_carry(factory, value, expected): + stamp = factory(value) + + assert (stamp.sec, stamp.nanosec) == expected + assert to_seconds(stamp) == pytest.approx(value, abs=1e-9) + + +@pytest.mark.parametrize("value", [float("nan"), float("inf"), -float("inf")]) +def test_nonfinite_seconds_are_rejected(value): + with pytest.raises(ValueError, match="finite"): + time_from_seconds(value) + + +def test_invalid_ros_fraction_is_rejected(): + stamp = time_from_nanoseconds(0) + stamp.nanosec = 1_000_000_000 + + with pytest.raises(ValueError, match="nanosec"): + to_nanoseconds(stamp) + + +def test_now_reads_integer_clock_and_preserves_frame(monkeypatch): + monkeypatch.setattr("dimos.message_utils.time.time.time_ns", lambda: 1_700_000_000_123_456_789) + + header = header_now("camera") + + assert to_nanoseconds(header.stamp) == 1_700_000_000_123_456_789 + assert header.frame_id == "camera" + + +@pytest.mark.parametrize("factory", [time_from_nanoseconds, duration_from_nanoseconds]) +@pytest.mark.parametrize("seconds", [-(2**31), 2**31 - 1]) +def test_ros_seconds_range_boundaries(factory, seconds): + value = seconds * 1_000_000_000 + 999_999_999 + assert to_nanoseconds(factory(value)) == value + + +@pytest.mark.parametrize("factory", [time_from_nanoseconds, duration_from_nanoseconds]) +@pytest.mark.parametrize("value", [-(2**31) * 1_000_000_000 - 1, 2**31 * 1_000_000_000]) +def test_ros_seconds_overflow_is_explicit(factory, value): + with pytest.raises(ValueError, match="signed 32-bit"): + factory(value) diff --git a/dimos/message_utils/time.py b/dimos/message_utils/time.py new file mode 100644 index 00000000000..3d9197cb58a --- /dev/null +++ b/dimos/message_utils/time.py @@ -0,0 +1,77 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Explicit conversions for generated ROS time, duration, and header fields.""" + +import math +import time +from typing import Any + +from dimos.messages.builtin_interfaces.msg import Duration, Time +from dimos.messages.std_msgs.msg import Header + +NANOSECONDS_PER_SECOND = 1_000_000_000 + + +def to_nanoseconds(stamp: Time | Duration) -> int: + """Convert a normalized ROS time or duration without losing integer precision.""" + if not 0 <= stamp.nanosec < NANOSECONDS_PER_SECOND: + raise ValueError("nanosec must be in [0, 1000000000)") + return stamp.sec * NANOSECONDS_PER_SECOND + stamp.nanosec + + +def to_seconds(stamp: Time | Duration) -> float: + """Convert to floating seconds for APIs that require them.""" + return to_nanoseconds(stamp) / NANOSECONDS_PER_SECOND + + +def time_from_nanoseconds(value: int) -> Time: + sec, nanosec = divmod(value, NANOSECONDS_PER_SECOND) + if not -(2**31) <= sec < 2**31: + raise ValueError("seconds must fit the ROS signed 32-bit field") + return Time(sec=sec, nanosec=nanosec) + + +def duration_from_nanoseconds(value: int) -> Duration: + sec, nanosec = divmod(value, NANOSECONDS_PER_SECOND) + if not -(2**31) <= sec < 2**31: + raise ValueError("seconds must fit the ROS signed 32-bit field") + return Duration(sec=sec, nanosec=nanosec) + + +def _seconds_to_nanoseconds(value: float) -> int: + if not math.isfinite(value): + raise ValueError("seconds must be finite") + sec = math.floor(value) + return sec * NANOSECONDS_PER_SECOND + round((value - sec) * NANOSECONDS_PER_SECOND) + + +def time_from_seconds(value: float) -> Time: + """Round to the nearest nanosecond, normalizing negative and carried fractions.""" + return time_from_nanoseconds(_seconds_to_nanoseconds(value)) + + +def duration_from_seconds(value: float) -> Duration: + return duration_from_nanoseconds(_seconds_to_nanoseconds(value)) + + +def header_now(frame_id: str = "") -> Header: + """Construct a wall-clock header; generated defaults themselves remain zero.""" + return Header(stamp=time_from_nanoseconds(time.time_ns()), frame_id=frame_id) + + +def message_header(message: Any) -> Header | None: + """Return source metadata from a generated stamped value.""" + header = getattr(message, "header", None) + return header if isinstance(header, Header) else None From ec07ea0e773dcdf474e151aa99d08186a2b21ee0 Mon Sep 17 00:00:00 2001 From: cc Date: Fri, 23 Oct 2026 02:11:19 -0700 Subject: [PATCH 02/10] refactor(messages): isolate generated geometry value operations --- dimos/message_utils/geometry.py | 376 +++++++++++++++++++++++++++ dimos/message_utils/test_geometry.py | 283 ++++++++++++++++++++ 2 files changed, 659 insertions(+) create mode 100644 dimos/message_utils/geometry.py create mode 100644 dimos/message_utils/test_geometry.py diff --git a/dimos/message_utils/geometry.py b/dimos/message_utils/geometry.py new file mode 100644 index 00000000000..2ab85421267 --- /dev/null +++ b/dimos/message_utils/geometry.py @@ -0,0 +1,376 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Geometry operations on generated ROS value types.""" + +from collections.abc import Sequence +from copy import deepcopy +import math + +import numpy as np +from numpy.typing import NDArray +from scipy.spatial.transform import Rotation + +from dimos.messages.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Vector3, + Wrench, +) +from dimos.messages.nav_msgs.msg import Odometry +from dimos.messages.std_msgs.msg import Header + + +def yaw(rotation: Quaternion) -> float: + """Return rotation about Z in radians (ROS XYZW quaternion convention).""" + return math.atan2( + 2 * (rotation.w * rotation.z + rotation.x * rotation.y), + 1 - 2 * (rotation.y * rotation.y + rotation.z * rotation.z), + ) + + +def transform_from_odometry(message: Odometry) -> TransformStamped: + """Copy the pose and declared frames into TF, preserving the exact source stamp.""" + position = message.pose.pose.position + return TransformStamped( + header=deepcopy(message.header), + child_frame_id=message.child_frame_id, + transform=Transform( + translation=Vector3(x=position.x, y=position.y, z=position.z), + rotation=deepcopy(message.pose.pose.orientation), + ), + ) + + +def quaternion_from_euler(roll: float, pitch: float, yaw: float) -> Quaternion: + """Convert fixed-axis XYZ angles in radians to a normalized ROS quaternion.""" + if not all(math.isfinite(value) for value in (roll, pitch, yaw)): + raise ValueError("Euler angles must be finite") + x, y, z, w = Rotation.from_euler("xyz", [roll, pitch, yaw]).as_quat() + return Quaternion(x=x, y=y, z=z, w=w) + + +def _rotation(quaternion: Quaternion) -> Rotation: + values = [quaternion.x, quaternion.y, quaternion.z, quaternion.w] + if not all(math.isfinite(value) for value in values): + raise ValueError("Quaternion components must be finite") + # SciPy rejects zero norm and normalizes non-unit input. + return Rotation.from_quat(values) + + +def _translation(transform: Transform) -> NDArray[np.float64]: + translation = transform.translation + values = np.array([translation.x, translation.y, translation.z], dtype=np.float64) + if not np.isfinite(values).all(): + raise ValueError("Translation components must be finite") + return values + + +def _transform(translation: NDArray[np.float64], rotation: Rotation) -> Transform: + x, y, z, w = rotation.as_quat() + return Transform( + translation=Vector3(x=translation[0], y=translation[1], z=translation[2]), + rotation=Quaternion(x=x, y=y, z=z, w=w), + ) + + +def transform_matrix(transform: Transform) -> NDArray[np.float64]: + """Copy a generated transform into a homogeneous 4-by-4 matrix.""" + matrix = np.eye(4, dtype=np.float64) + matrix[:3, :3] = _rotation(transform.rotation).as_matrix() + matrix[:3, 3] = _translation(transform) + return matrix + + +def compose_transforms(first: TransformStamped, second: TransformStamped) -> TransformStamped: + """Compose A←B and B←C, preserving A's exact source stamp in the A←C result. + + Inputs are copied into a new value. The shared frame must match; callers + querying a time-varying chain are responsible for choosing each edge's time. + """ + if first.child_frame_id != second.header.frame_id: + raise ValueError( + f"Cannot compose frames {first.child_frame_id!r} and {second.header.frame_id!r}" + ) + rotation = _rotation(first.transform.rotation) + translation = _translation(first.transform) + rotation.apply(_translation(second.transform)) + return TransformStamped( + header=deepcopy(first.header), + child_frame_id=second.child_frame_id, + transform=_transform(translation, rotation * _rotation(second.transform.rotation)), + ) + + +def inverse_transform(message: TransformStamped) -> TransformStamped: + """Return B←A from A←B, swapping frame names and preserving the exact stamp.""" + rotation = _rotation(message.transform.rotation).inv() + return TransformStamped( + header=Header(stamp=deepcopy(message.header.stamp), frame_id=message.child_frame_id), + child_frame_id=message.header.frame_id, + transform=_transform(-rotation.apply(_translation(message.transform)), rotation), + ) + + +def pose_from_transform(message: TransformStamped) -> PoseStamped: + """Copy the child frame's pose in the parent frame into a generated PoseStamped.""" + translation = message.transform.translation + return PoseStamped( + header=deepcopy(message.header), + pose=Pose( + position=Point(x=translation.x, y=translation.y, z=translation.z), + orientation=deepcopy(message.transform.rotation), + ), + ) + + +def transform_from_pose(message: PoseStamped, *, child_frame_id: str) -> TransformStamped: + """Copy a stamped pose into TF with an explicit child and unchanged source header.""" + if not isinstance(message, PoseStamped): + raise TypeError("Expected a generated PoseStamped") + position = message.pose.position + return TransformStamped( + header=deepcopy(message.header), + child_frame_id=child_frame_id, + transform=Transform( + translation=Vector3(x=position.x, y=position.y, z=position.z), + rotation=deepcopy(message.pose.orientation), + ), + ) + + +def pose_matrix(pose: Pose) -> NDArray[np.float64]: + """Copy a generated pose into a homogeneous matrix, rejecting invalid geometry.""" + values = np.array([pose.position.x, pose.position.y, pose.position.z], dtype=np.float64) + if not np.isfinite(values).all(): + raise ValueError("Position components must be finite") + result = np.eye(4, dtype=np.float64) + result[:3, :3] = _rotation(pose.orientation).as_matrix() + result[:3, 3] = values + return result + + +def quaternion_from_matrix(matrix: NDArray[np.float64]) -> Quaternion: + """Convert a finite proper 3-by-3 rotation matrix into a generated quaternion.""" + values = np.asarray(matrix, dtype=np.float64) + if values.shape != (3, 3) or not np.isfinite(values).all(): + raise ValueError("Rotation matrix must be finite and 3-by-3") + if not np.allclose(values.T @ values, np.eye(3), atol=1e-6, rtol=0.0) or not np.isclose( + np.linalg.det(values), 1.0, atol=1e-6, rtol=0.0 + ): + raise ValueError("Rotation matrix must be orthonormal with determinant +1") + x, y, z, w = Rotation.from_matrix(values).as_quat() + return Quaternion(x=x, y=y, z=z, w=w) + + +def pose_from_matrix(matrix: NDArray[np.float64]) -> Pose: + """Copy a finite rigid 4-by-4 transform into a generated pose.""" + values = np.asarray(matrix, dtype=np.float64) + if values.shape != (4, 4) or not np.isfinite(values).all(): + raise ValueError("Pose matrix must be finite and 4-by-4") + if not np.allclose(values[3], [0.0, 0.0, 0.0, 1.0], atol=1e-9, rtol=0.0): + raise ValueError("Pose matrix must have homogeneous last row [0, 0, 0, 1]") + return Pose( + position=Point(x=values[0, 3], y=values[1, 3], z=values[2, 3]), + orientation=quaternion_from_matrix(values[:3, :3]), + ) + + +def transform_from_matrix(matrix: NDArray[np.float64]) -> Transform: + """Copy a finite rigid homogeneous matrix into a generated transform.""" + pose = pose_from_matrix(matrix) + return Transform( + translation=Vector3(x=pose.position.x, y=pose.position.y, z=pose.position.z), + rotation=deepcopy(pose.orientation), + ) + + +def quaternion_euler(rotation: Quaternion) -> tuple[float, float, float]: + """Return fixed-axis XYZ Euler angles in radians for a generated quaternion.""" + roll, pitch, yaw = _rotation(rotation).as_euler("xyz") + return float(roll), float(pitch), float(yaw) + + +def translate_pose_local(pose: Pose, offset: Vector3) -> Pose: + """Copy a pose displaced by an offset expressed in its local axes.""" + matrix = pose_matrix(pose) + values = np.array([offset.x, offset.y, offset.z], dtype=np.float64) + if not np.isfinite(values).all(): + raise ValueError("Offset components must be finite") + delta = matrix[:3, :3] @ values + return Pose( + position=Point( + x=pose.position.x + delta[0], y=pose.position.y + delta[1], z=pose.position.z + delta[2] + ), + orientation=deepcopy(pose.orientation), + ) + + +def point_distance(first: Point, second: Point) -> float: + """Euclidean distance between generated position values.""" + return math.dist((first.x, first.y, first.z), (second.x, second.y, second.z)) + + +def quaternion_angle(first: Quaternion, second: Quaternion) -> float: + """Return the shortest angular separation in radians, normalizing inputs.""" + return float((_rotation(first).inv() * _rotation(second)).magnitude()) + + +def wrench_from_array(values: Sequence[float] | NDArray[np.float64]) -> Wrench: + """Copy force in N and torque in Nm from [fx, fy, fz, tx, ty, tz].""" + array = np.asarray(values, dtype=np.float64) + if array.shape != (6,): + raise ValueError("Expected 6 elements [fx, fy, fz, tx, ty, tz]") + return Wrench( + force=Vector3(x=array[0], y=array[1], z=array[2]), + torque=Vector3(x=array[3], y=array[4], z=array[5]), + ) + + +def wrench_array(message: Wrench) -> NDArray[np.float64]: + """Return an independent numeric force/torque array for arithmetic.""" + return np.array( + [ + message.force.x, + message.force.y, + message.force.z, + message.torque.x, + message.torque.y, + message.torque.z, + ], + dtype=np.float64, + ) + + +def vector_array(value: Vector3 | Point) -> NDArray[np.float64]: + """Copy a generated spatial vector/point into an independent XYZ array.""" + return np.array([value.x, value.y, value.z], dtype=np.float64) + + +def vector_from_array(values: Sequence[float] | NDArray[np.float64]) -> Vector3: + """Copy exactly three coordinates into a generated vector.""" + array = np.asarray(values, dtype=np.float64) + if array.shape != (3,): + raise ValueError("Expected 3 vector components") + return Vector3(x=array[0], y=array[1], z=array[2]) + + +def normalized_vector(value: Vector3) -> Vector3: + """Return a unit vector; zero length remains zero.""" + array = vector_array(value) + length = float(np.linalg.norm(array)) + return vector_from_array(array / length) if length else Vector3(x=0.0, y=0.0, z=0.0) + + +def quaternion_array(value: Quaternion) -> NDArray[np.float64]: + """Copy XYZW components without normalizing their declared values.""" + return np.array([value.x, value.y, value.z, value.w], dtype=np.float64) + + +def quaternion_from_array(values: Sequence[float] | NDArray[np.float64]) -> Quaternion: + """Copy exactly four XYZW components into a generated quaternion.""" + array = np.asarray(values, dtype=np.float64) + if array.shape != (4,): + raise ValueError("Quaternion requires exactly 4 components") + return Quaternion(x=float(array[0]), y=float(array[1]), z=float(array[2]), w=float(array[3])) + + +def quaternion_product(first: Quaternion, second: Quaternion) -> Quaternion: + """Hamilton product, retaining non-unit magnitudes and multiplication order.""" + if not isinstance(first, Quaternion) or not isinstance(second, Quaternion): + raise TypeError("Cannot multiply Quaternion with a non-quaternion value") + a, b = quaternion_array(first), quaternion_array(second) + xyz = a[3] * b[:3] + b[3] * a[:3] + np.cross(a[:3], b[:3]) + return quaternion_from_array(np.append(xyz, a[3] * b[3] - np.dot(a[:3], b[:3]))) + + +def quaternion_conjugate(value: Quaternion) -> Quaternion: + """Negate the imaginary components, preserving the scalar component.""" + return Quaternion(x=-value.x, y=-value.y, z=-value.z, w=value.w) + + +def quaternion_inverse(value: Quaternion) -> Quaternion: + """Algebraic inverse, including non-unit quaternions.""" + array = quaternion_array(value) + squared = float(np.dot(array, array)) + if squared == 0: + raise ZeroDivisionError("Cannot invert zero quaternion") + return quaternion_from_array(quaternion_array(quaternion_conjugate(value)) / squared) + + +def normalized_quaternion(value: Quaternion) -> Quaternion: + """Return unit components without changing sign; reject zero norm.""" + array = quaternion_array(value) + length = float(np.linalg.norm(array)) + if length == 0: + raise ZeroDivisionError("Cannot normalize zero quaternion") + return quaternion_from_array(array / length) + + +def rotate_vector(rotation: Quaternion, value: Vector3) -> Vector3: + """Rotate an XYZ vector using a normalized spatial rotation.""" + return vector_from_array(_rotation(rotation).apply(vector_array(value))) + + +def point_from_array(values: Sequence[float] | NDArray[np.float64]) -> Point: + """Copy exactly three Cartesian coordinates into a generated point.""" + vector = vector_from_array(values) + return Point(x=vector.x, y=vector.y, z=vector.z) + + +def compose_poses(first: Pose, second: Pose) -> Pose: + """Compose parent/local poses without mutating inputs or normalizing stored quaternions.""" + offset = _rotation(first.orientation).apply(vector_array(second.position)) + return Pose( + position=point_from_array(vector_array(first.position) + offset), + orientation=quaternion_product(first.orientation, second.orientation), + ) + + +def compose_transform_values(first: Transform, second: Transform) -> Transform: + """Compose unframed spatial transforms in parent/local order.""" + if not isinstance(first, Transform) or not isinstance(second, Transform): + raise TypeError("Expected generated Transform values") + return Transform( + translation=vector_from_array( + vector_array(first.translation) + + _rotation(first.rotation).apply(vector_array(second.translation)) + ), + rotation=quaternion_product(first.rotation, second.rotation), + ) + + +def transform_pose_local(pose: Pose, transform: Transform) -> Pose: + """Apply an offset expressed in a pose's local axes.""" + return compose_poses( + pose, + Pose( + position=point_from_array(vector_array(transform.translation)), + orientation=transform.rotation, + ), + ) + + +def relative_pose_transform(first: Pose, second: Pose) -> Transform: + """Find the local transform taking the first pose to the second.""" + rotation = quaternion_inverse(first.orientation) + return Transform( + translation=vector_from_array( + _rotation(rotation).apply(vector_array(second.position) - vector_array(first.position)) + ), + rotation=quaternion_product(rotation, second.orientation), + ) diff --git a/dimos/message_utils/test_geometry.py b/dimos/message_utils/test_geometry.py new file mode 100644 index 00000000000..a26d0a07982 --- /dev/null +++ b/dimos/message_utils/test_geometry.py @@ -0,0 +1,283 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Copyright 2026 Dimensional Inc. +# SPDX-License-Identifier: Apache-2.0 + +import math + +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +import numpy as np +import pytest + +from dimos.message_utils.geometry import ( + compose_transforms, + inverse_transform, + pose_from_matrix, + pose_from_transform, + pose_matrix, + quaternion_euler, + quaternion_from_euler, + transform_from_odometry, + transform_from_pose, + transform_matrix, + translate_pose_local, + yaw, +) +from dimos.messages.builtin_interfaces.msg import Time +from dimos.messages.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + PoseWithCovariance, + Quaternion, + Transform, + TransformStamped, + Twist, + TwistWithCovariance, + Vector3, +) +from dimos.messages.nav_msgs.msg import Odometry +from dimos.messages.std_msgs.msg import Header + + +@pytest.mark.parametrize("seconds,nanoseconds", [(0, 0), (-1, 500000000), (1700000000, 123456789)]) +def test_odometry_transform_preserves_frames_stamp_and_pose(seconds, nanoseconds): + message = Odometry( + child_frame_id="robot/lidar", + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + pose=PoseWithCovariance( + pose=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + covariance=np.zeros(36, dtype=np.float64), + ), + twist=TwistWithCovariance( + twist=Twist(linear=Vector3(x=0.0, y=0.0, z=0.0), angular=Vector3(x=0.0, y=0.0, z=0.0)), + covariance=np.zeros(36, dtype=np.float64), + ), + ) + message.header.frame_id = "robot/odom" + message.header.stamp = Time(sec=seconds, nanosec=nanoseconds) + message.pose.pose.position = Point(x=1, y=2, z=3) + message.pose.pose.orientation = Quaternion(z=0.6, w=0.8, x=0.0, y=0.0) + + transformed = transform_from_odometry(message) + + assert transformed.header.frame_id == "robot/odom" + assert transformed.child_frame_id == "robot/lidar" + assert transformed.header.stamp.sec == seconds + assert transformed.header.stamp.nanosec == nanoseconds + assert ( + transformed.transform.translation.x, + transformed.transform.translation.y, + transformed.transform.translation.z, + ) == (1, 2, 3) + assert transformed.transform.rotation == message.pose.pose.orientation + transformed.header.frame_id = "changed" + transformed.transform.rotation.w = 0 + assert message.header.frame_id == "robot/odom" + assert message.pose.pose.orientation.w == 0.8 + + +def _edge(parent, child, xyz, rotation, stamp=Time(sec=1700000000, nanosec=123456789)): + return TransformStamped( + header=Header(frame_id=parent, stamp=stamp), + child_frame_id=child, + transform=Transform(translation=Vector3(x=xyz[0], y=xyz[1], z=xyz[2]), rotation=rotation), + ) + + +def test_transform_composition_rotates_translation_and_preserves_exact_source_stamp(): + first = _edge("world", "arm", (1, -1, 0), quaternion_from_euler(0, 0, math.pi / 2)) + second = _edge( + "arm", "tool", (2, 3, 4), Quaternion(w=1, x=0.0, y=0.0, z=0.0), Time(sec=7, nanosec=0) + ) + + result = compose_transforms(first, second) + + assert result.header == first.header + assert result.child_frame_id == "tool" + translation = result.transform.translation + assert (translation.x, translation.y, translation.z) == pytest.approx((-2, 1, 4)) + assert yaw(result.transform.rotation) == pytest.approx(math.pi / 2) + result.header.stamp.nanosec = 0 + result.transform.translation.x = 99 + assert first.header.stamp.nanosec == 123456789 + assert first.transform.translation.x == 1 + assert second.transform.translation.x == 2 + + +def test_inverse_composes_to_identity_and_swaps_frames_without_losing_nanoseconds(): + original = _edge("map", "camera", (2, 3, -4), quaternion_from_euler(0.3, -0.2, 1.1)) + inverse = inverse_transform(original) + result = compose_transforms(original, inverse) + + assert inverse.header.frame_id == "camera" + assert inverse.child_frame_id == "map" + assert inverse.header.stamp == original.header.stamp + translation = result.transform.translation + assert (translation.x, translation.y, translation.z) == pytest.approx((0, 0, 0), abs=1e-12) + rotation = result.transform.rotation + assert (rotation.x, rotation.y, rotation.z, abs(rotation.w)) == pytest.approx((0, 0, 0, 1)) + + +def test_transform_composition_rejects_disconnected_frames(): + first = _edge("a", "b", (0, 0, 0), Quaternion(w=1, x=0.0, y=0.0, z=0.0)) + second = _edge("c", "d", (0, 0, 0), Quaternion(w=1, x=0.0, y=0.0, z=0.0)) + with pytest.raises(ValueError, match="Cannot compose frames"): + compose_transforms(first, second) + + +@pytest.mark.parametrize( + "rotation", + [ + Quaternion(w=0, x=0.0, y=0.0, z=0.0), + Quaternion(x=math.nan, w=1, y=0.0, z=0.0), + Quaternion(w=math.inf, x=0.0, y=0.0, z=0.0), + ], +) +def test_transform_math_rejects_invalid_rotations(rotation): + edge = _edge("a", "b", (0, 0, 0), rotation) + with pytest.raises(ValueError): + inverse_transform(edge) + with pytest.raises(ValueError): + compose_transforms(edge, _edge("b", "c", (0, 0, 0), Quaternion(w=1, x=0.0, y=0.0, z=0.0))) + + +@pytest.mark.parametrize("value", [math.nan, math.inf, -math.inf]) +def test_transform_math_rejects_nonfinite_translations_and_angles(value): + edge = _edge("a", "b", (value, 0, 0), Quaternion(w=1, x=0.0, y=0.0, z=0.0)) + with pytest.raises(ValueError, match="Translation components"): + inverse_transform(edge) + with pytest.raises(ValueError, match="Euler angles"): + quaternion_from_euler(value, 0, 0) + + +def test_pose_conversion_copies_nested_fields_and_preserves_header(): + edge = _edge("map", "robot", (1, 2, 3), Quaternion(z=0.6, w=0.8, x=0.0, y=0.0)) + pose = pose_from_transform(edge) + assert pose.header == edge.header + assert pose.pose.position == Point(x=1, y=2, z=3) + assert pose.pose.orientation == edge.transform.rotation + pose.header.frame_id = "changed" + pose.pose.orientation.w = 0 + assert edge.header.frame_id == "map" + assert edge.transform.rotation.w == 0.8 + + +def test_pose_transform_preserves_header_and_does_not_alias_input(): + source = PoseStamped( + header=Header(frame_id="odom", stamp=Time(sec=-1, nanosec=987654321)), + pose=Pose( + position=Point(x=1, y=2, z=3), orientation=Quaternion(z=0.6, w=0.8, x=0.0, y=0.0) + ), + ) + result = transform_from_pose(source, child_frame_id="robot/base") + decoded = cdr_decode(cdr_encode(result), TransformStamped) + assert decoded.header == source.header + assert decoded.child_frame_id == "robot/base" + assert decoded.transform.translation == Vector3(x=1, y=2, z=3) + assert decoded.transform.rotation == source.pose.orientation + result.header.frame_id = "changed" + result.transform.rotation.w = 0 + assert source.header.frame_id == "odom" + assert source.pose.orientation.w == 0.8 + + +def test_transform_matrix_maps_rotation_and_translation_without_aliasing(): + value = Transform( + translation=Vector3(x=2.0, y=-1.0, z=3.0), + rotation=quaternion_from_euler(0.0, 0.0, math.pi / 2), + ) + matrix = transform_matrix(value) + np.testing.assert_allclose(matrix @ [1.0, 0.0, 0.0, 1.0], [2.0, 0.0, 3.0, 1.0], atol=1e-12) + value.translation.x = 10.0 + assert matrix[0, 3] == 2.0 + + +def test_pose_matrix_roundtrip_preserves_rigid_motion_without_aliasing(): + pose = Pose( + position=Point(x=2, y=-1, z=3), + orientation=quaternion_from_euler(0, 0, math.pi / 2), + ) + matrix = pose_matrix(pose) + np.testing.assert_allclose(matrix @ [1, 0, 0, 1], [2, 0, 3, 1], atol=1e-12) + decoded = cdr_decode(cdr_encode(pose_from_matrix(matrix)), Pose) + np.testing.assert_allclose(pose_matrix(decoded), matrix, atol=1e-12) + matrix[0, 3] = 99 + decoded.position.y = 42 + assert pose.position == Point(x=2, y=-1, z=3) + assert decoded.position.x == 2 + + +@pytest.mark.parametrize( + "matrix", + [ + np.eye(3), + np.full((4, 4), math.nan), + np.diag([1, 1, 1, 2]), + np.diag([2, 1, 1, 1]), + np.diag([-1, 1, 1, 1]), + ], +) +def test_pose_from_matrix_rejects_nonrigid_or_invalid_matrices(matrix): + with pytest.raises(ValueError): + pose_from_matrix(matrix) + + +@pytest.mark.parametrize( + "pose", + [ + Pose( + position=Point(x=math.inf, y=0.0, z=0.0), + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + ), + Pose(orientation=Quaternion(w=0, x=0.0, y=0.0, z=0.0), position=Point(x=0.0, y=0.0, z=0.0)), + Pose( + orientation=Quaternion(w=math.nan, x=0.0, y=0.0, z=0.0), + position=Point(x=0.0, y=0.0, z=0.0), + ), + ], +) +def test_pose_matrix_rejects_invalid_geometry(pose): + with pytest.raises(ValueError): + pose_matrix(pose) + + +def test_local_offset_uses_pose_axes_and_copies_orientation(): + pose = Pose( + position=Point(x=1, y=2, z=3), + orientation=quaternion_from_euler(0, 0, math.pi / 2), + ) + translated = translate_pose_local(pose, Vector3(x=2, z=-1, y=0.0)) + position = translated.position + assert (position.x, position.y, position.z) == pytest.approx((1, 4, 2)) + assert quaternion_euler(translated.orientation) == pytest.approx((0, 0, math.pi / 2)) + translated.orientation.w = 0 + assert pose.orientation.w == pytest.approx(math.sqrt(0.5)) + assert pose.position == Point(x=1, y=2, z=3) + + +@pytest.mark.parametrize("value", [math.nan, math.inf, -math.inf]) +def test_local_offset_rejects_nonfinite_displacements(value): + with pytest.raises(ValueError, match="Offset components"): + translate_pose_local( + Pose( + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + position=Point(x=0.0, y=0.0, z=0.0), + ), + Vector3(x=value, y=0.0, z=0.0), + ) From dada376037db3b74760ebacbb190dd365e33c608 Mon Sep 17 00:00:00 2001 From: cc Date: Fri, 23 Oct 2026 02:11:20 -0700 Subject: [PATCH 03/10] fix(pointcloud): preserve records without serialization copies --- dimos/message_utils/pointcloud.py | 301 +++++++++++- dimos/message_utils/test_pointcloud.py | 462 ++++++++++++++++++ dimos/message_utils/test_pointcloud_colors.py | 85 ++++ 3 files changed, 846 insertions(+), 2 deletions(-) create mode 100644 dimos/message_utils/test_pointcloud.py create mode 100644 dimos/message_utils/test_pointcloud_colors.py diff --git a/dimos/message_utils/pointcloud.py b/dimos/message_utils/pointcloud.py index 1e630828eab..fed75827913 100644 --- a/dimos/message_utils/pointcloud.py +++ b/dimos/message_utils/pointcloud.py @@ -12,14 +12,24 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Read generated point-cloud layouts and extract positions and colors.""" +"""NumPy conversions for generated PointCloud2 messages.""" +from copy import deepcopy +from dataclasses import replace +import math from typing import Any import numpy as np from numpy.typing import NDArray +import open3d as o3d # type: ignore[import-untyped] -from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.message_utils.camera_info import intrinsic_matrix +from dimos.message_utils.geometry import transform_matrix +from dimos.message_utils.image import image_to_rgb, image_view +from dimos.message_utils.time import to_nanoseconds +from dimos.messages.geometry_msgs.msg import TransformStamped +from dimos.messages.sensor_msgs.msg import CameraInfo, Image, PointCloud2, PointField +from dimos.messages.std_msgs.msg import Header _FIELD_KINDS = {1: "i1", 2: "u1", 3: "i2", 4: "u2", 5: "i4", 6: "u4", 7: "f4", 8: "f8"} @@ -72,6 +82,96 @@ def pointcloud_xyz(message: PointCloud2) -> NDArray[np.float64]: return np.stack([points[name].ravel() for name in ("x", "y", "z")], axis=-1).astype(np.float64) +def pointcloud_from_xyz( + points: NDArray[Any], *, header: Header, stamps: NDArray[Any] | None = None +) -> PointCloud2: + """Copy Nx3 or HxWx3 coordinates into a tightly packed little-endian XYZ cloud. + + XYZ fields are float32. NaNs/infinities remain missing-point markers and set + is_dense=False; finite coordinates outside float32 range are rejected. + """ + values = np.asarray(points) + if values.ndim not in (2, 3) or values.shape[-1] != 3: + raise ValueError("XYZ coordinates must have shape (N, 3) or (H, W, 3)") + if values.dtype.kind not in "fiu": + raise ValueError("XYZ coordinates must be numeric") + finite = np.isfinite(values) + if np.any(np.abs(values[finite].astype(np.float64)) > np.finfo(np.float32).max): + raise ValueError("XYZ coordinates exceed float32 range") + packed = np.ascontiguousarray(values, dtype=" NDArray[np.float64] | None: + """Read per-point capture times; absence falls back to the header at the caller.""" + fields = [field for field in message.fields if field.name == "stamps"] + if not fields: + return None + if len(fields) != 1 or fields[0].datatype != PointField.FLOAT64 or fields[0].count != 1: + raise ValueError("Per-point stamps must be a scalar FLOAT64 field") + return np.asarray(pointcloud_view(message)["stamps"].reshape(-1), dtype=np.float64) + + +def select_points(message: PointCloud2, keep: NDArray[np.bool_]) -> PointCloud2: + """Copy selected point records in row order, retaining every field and point byte. + + The result is unorganized with no row padding. Point padding, endianness, + field metadata, source header, and the conservative density flag survive. + """ + pointcloud_view(message) # Validate the full declared layout before selecting bytes. + if keep.dtype != np.bool_ or keep.shape != (message.height * message.width,): + raise ValueError("point selection must be a flat boolean mask matching the point count") + records: NDArray[np.uint8] = np.ndarray( + (message.height, message.width, message.point_step), + dtype=np.uint8, + buffer=message.data, + strides=(message.row_step, message.point_step, 1), + ) + selected = records.reshape(message.height * message.width, message.point_step)[keep] + width = len(selected) + return PointCloud2( + header=deepcopy(message.header), + height=1, + width=width, + fields=deepcopy(message.fields), + is_bigendian=message.is_bigendian, + point_step=message.point_step, + row_step=width * message.point_step, + data=np.asarray(np.frombuffer(selected.tobytes(), dtype=np.uint8).copy(), dtype=np.uint8), + is_dense=message.is_dense, + ) + + def pointcloud_rgb(message: PointCloud2) -> NDArray[np.uint8] | None: """Copy RGB bytes from the standard packed rgb/rgba field, if present. @@ -88,3 +188,200 @@ def pointcloud_rgb(message: PointCloud2) -> NDArray[np.uint8] | None: raise ValueError("packed point cloud color must be a scalar FLOAT32 or UINT32 field") packed = field.view(np.dtype(">u4" if message.is_bigendian else "> shift) & 255 for shift in (16, 8, 0)], axis=-1).astype(np.uint8) + + +def concatenate_clouds(first: PointCloud2, second: PointCloud2) -> PointCloud2: + """Join compatible point records, preserving fields, endian, and point padding. + + Row padding is removed. The first frame and newest exact source stamp survive. + Empty inputs are identities; differing nonempty frames/layouts are rejected. + """ + pointcloud_view(first) + pointcloud_view(second) + if not first.width * first.height: + return deepcopy(second) + if not second.width * second.height: + return deepcopy(first) + if first.header.frame_id != second.header.frame_id: + raise ValueError("cannot concatenate point clouds in different frames") + if ( + first.fields != second.fields + or first.point_step != second.point_step + or first.is_bigendian != second.is_bigendian + ): + raise ValueError("cannot concatenate point clouds with different record layouts") + a = select_points(first, np.ones(first.width * first.height, dtype=np.bool_)) + b = select_points(second, np.ones(second.width * second.height, dtype=np.bool_)) + stamp = max((first.header.stamp, second.header.stamp), key=to_nanoseconds) + return PointCloud2( + header=Header(frame_id=first.header.frame_id, stamp=stamp), + height=1, + width=a.width + b.width, + fields=deepcopy(first.fields), + is_bigendian=first.is_bigendian, + point_step=first.point_step, + row_step=a.row_step + b.row_step, + data=np.asarray(np.concatenate((a.data, b.data)), dtype=np.uint8), + is_dense=first.is_dense and second.is_dense, + ) + + +def transform_cloud(message: PointCloud2, transform: TransformStamped) -> PointCloud2: + """Apply parent←child to XYZ while retaining the complete organized record layout. + + Copy storage before changing coordinates. Keep all other fields, point/row + padding, endian and the cloud's exact source stamp. Frames must match; + integer XYZ is rejected because transformed coordinates would be truncated. + """ + if message.header.frame_id != transform.child_frame_id: + raise ValueError("point cloud frame does not match transform child frame") + view = pointcloud_view(message) + for name in ("x", "y", "z"): + if name not in (view.dtype.names or ()) or view.dtype[name].kind != "f": + raise ValueError("transform requires scalar floating-point XYZ fields") + matrix = transform_matrix(transform.transform) + xyz = pointcloud_xyz(message) + transformed = xyz @ matrix[:3, :3].T + matrix[:3, 3] + for index, name in enumerate(("x", "y", "z")): + values = transformed[:, index] + if np.any(np.abs(values[np.isfinite(values)]) > np.finfo(view.dtype[name]).max): + raise ValueError(f"transformed {name} coordinates exceed field range") + payload = bytearray(bytes(message.data)) + writable = np.ndarray( + (message.height, message.width), + dtype=view.dtype, + buffer=payload, + strides=(message.row_step, message.point_step), + ) + for index, name in enumerate(("x", "y", "z")): + writable[name] = transformed[:, index].reshape(message.height, message.width) + output = replace(message, header=deepcopy(message.header), fields=deepcopy(message.fields)) + output.data = np.frombuffer(payload, dtype=np.uint8).copy() + output.header = Header(stamp=deepcopy(message.header.stamp), frame_id=transform.header.frame_id) + return output + + +def pointcloud_from_xyz_rgb( + points: NDArray[Any], colors: NDArray[np.uint8], *, header: Header +) -> PointCloud2: + """Copy XYZ and RGB8 into a declared XYZ FLOAT32 / packed RGB UINT32 layout.""" + cloud = pointcloud_from_xyz(points, header=header) + rgb = np.asarray(colors) + if rgb.dtype != np.uint8 or rgb.shape != np.asarray(points).shape: + raise ValueError("RGB must be uint8 with the same (..., 3) shape as XYZ") + records = np.empty((cloud.height, cloud.width), dtype=[("xyz", " PointCloud2: + """Project a rectified RGB/depth pair into the depth optical frame. + + Integer depth is multiplied by depth_scale to obtain meters; 32FC1 already + represents meters. Invalid/nonpositive/too-distant samples are removed. + Scale K when calibration resolution differs, retaining the depth header. + """ + if ( + not math.isfinite(depth_scale) + or depth_scale <= 0 + or not math.isfinite(depth_trunc) + or depth_trunc <= 0 + ): + raise ValueError("depth scale and truncation must be finite and positive") + if (color.width, color.height) != (depth.width, depth.height): + raise ValueError("color and depth dimensions do not match") + if depth.encoding not in ("16UC1", "mono16", "32FC1"): + raise ValueError("RGBD requires uint16 or float32 depth") + if calibration.width <= 0 or calibration.height <= 0: + raise ValueError("calibration dimensions must be positive") + matrix = intrinsic_matrix(calibration).copy() + matrix[0, :] *= depth.width / calibration.width + matrix[1, :] *= depth.height / calibration.height + distances = image_view(depth).astype(np.float32) + if depth.encoding != "32FC1": + distances *= depth_scale + distances[~np.isfinite(distances) | (distances <= 0)] = 0 + rgbd = o3d.geometry.RGBDImage.create_from_color_and_depth( + o3d.geometry.Image(np.ascontiguousarray(image_to_rgb(color))), + o3d.geometry.Image(np.ascontiguousarray(distances)), + depth_scale=1.0, + depth_trunc=depth_trunc, + convert_rgb_to_intensity=False, + ) + intrinsics = o3d.camera.PinholeCameraIntrinsic( + depth.width, + depth.height, + matrix[0, 0], + matrix[1, 1], + matrix[0, 2], + matrix[1, 2], + ) + cloud = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd, intrinsics) + colors = np.rint(np.asarray(cloud.colors) * 255).astype(np.uint8) + return pointcloud_from_xyz_rgb(np.asarray(cloud.points), colors, header=depth.header) + + +def voxel_downsample_cloud(message: PointCloud2, voxel_size: float = 0.025) -> PointCloud2: + """Downsample XYZ/RGB clouds with the existing Open3D tensor backend. + + Non-positive sizes and clouds below twenty points retain the existing fast + path. Unsupported extra fields are rejected rather than silently dropped. + Source headers are preserved; generated messages own no numerical methods. + """ + if not math.isfinite(voxel_size): + raise ValueError("voxel_size must be finite") + if voxel_size <= 0 or message.width * message.height < 20: + return message + names = {field.name for field in message.fields} + if names not in ({"x", "y", "z"}, {"x", "y", "z", "rgb"}): + raise ValueError("voxel downsampling supports only XYZ and optional RGB fields") + points = pointcloud_xyz(message).astype(np.float32) + if not np.isfinite(points).all(): + raise ValueError("voxel downsampling requires finite XYZ positions") + colors = pointcloud_rgb(message) + + device = o3d.core.Device("CPU:0") + cloud = o3d.t.geometry.PointCloud(o3d.core.Tensor(points, o3d.core.Dtype.Float32, device)) + if colors is not None: + cloud.point["colors"] = o3d.core.Tensor( + colors.astype(np.float32) / 255.0, o3d.core.Dtype.Float32, device + ) + reduced = cloud.voxel_down_sample(voxel_size) + xyz = reduced.point["positions"].numpy() + if colors is None: + return pointcloud_from_xyz(xyz, header=message.header) + rgb = np.clip(reduced.point["colors"].numpy() * 255.0, 0, 255).astype(np.uint8) + return pointcloud_from_xyz_rgb(xyz, rgb, header=message.header) + + +def pointcloud_to_open3d(message: PointCloud2) -> Any: + """Copy generated XYZ/RGB values into a native geometry algorithm input.""" + + result = o3d.geometry.PointCloud() + result.points = o3d.utility.Vector3dVector(pointcloud_xyz(message)) + colors = pointcloud_rgb(message) + if colors is not None: + result.colors = o3d.utility.Vector3dVector(colors.astype(np.float64) / 255.0) + return result + + +def cloud_bounds_intersect(first: PointCloud2, second: PointCloud2) -> bool: + """Test closed axis-aligned bounds; empty or entirely nonfinite clouds do not intersect.""" + a, b = pointcloud_xyz(first), pointcloud_xyz(second) + a, b = a[np.isfinite(a).all(axis=1)], b[np.isfinite(b).all(axis=1)] + if len(a) == 0 or len(b) == 0: + return False + return bool(np.all(a.min(axis=0) <= b.max(axis=0)) and np.all(a.max(axis=0) >= b.min(axis=0))) diff --git a/dimos/message_utils/test_pointcloud.py b/dimos/message_utils/test_pointcloud.py new file mode 100644 index 00000000000..48411b23d56 --- /dev/null +++ b/dimos/message_utils/test_pointcloud.py @@ -0,0 +1,462 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Copyright 2026 Dimensional Inc. +# SPDX-License-Identifier: Apache-2.0 + +import gc +import struct + +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +import numpy as np +import pytest + +from dimos.message_utils.camera_info import camera_info_from_intrinsics +from dimos.message_utils.image import image_from_array +from dimos.message_utils.pointcloud import ( + concatenate_clouds, + pointcloud_from_rgbd, + pointcloud_from_xyz, + pointcloud_from_xyz_rgb, + pointcloud_rgb, + pointcloud_stamps, + pointcloud_view, + pointcloud_xyz, + select_points, + transform_cloud, + voxel_downsample_cloud, +) +from dimos.messages.builtin_interfaces.msg import Time +from dimos.messages.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos.messages.sensor_msgs.msg import PointCloud2, PointField +from dimos.messages.std_msgs.msg import Header + + +@pytest.fixture(params=[False, True], ids=["little-endian", "big-endian"]) +def padded_cloud(request): + # Two organized rows, two points per row; fields are deliberately reordered. + endian = ">" if request.param else "<" + payload = bytearray(128) + for row in range(2): + for column in range(2): + value = row * 2 + column + 1 + struct.pack_into( + endian + "df2Hf", + payload, + row * 64 + column * 24, + value * 3.0, + value * 1.0, + value, + value + 10, + value * 2.0, + ) + return PointCloud2( + height=2, + width=2, + point_step=24, + row_step=64, + is_bigendian=request.param, + data=np.frombuffer(bytes(payload), dtype=np.uint8), + fields=[ + PointField(name="z", offset=0, datatype=8, count=1), + PointField(name="x", offset=8, datatype=7, count=1), + PointField(name="tags", offset=12, datatype=4, count=2), + PointField(name="y", offset=16, datatype=7, count=1), + ], + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + is_dense=False, + ) + + +def test_field_counts_padding_and_endianness(padded_cloud): + view = pointcloud_view(padded_cloud) + assert view["tags"].tolist() == [[[1, 11], [2, 12]], [[3, 13], [4, 14]]] + assert pointcloud_xyz(padded_cloud).tolist() == [[1, 2, 3], [2, 4, 6], [3, 6, 9], [4, 8, 12]] + assert not view.flags.writeable + with pytest.raises(ValueError, match="WRITEABLE"): + view.setflags(write=True) + copied = view.copy() + copied["x"][0, 0] = 99 + assert view["x"][0, 0] == 1 + + +def test_view_retains_storage_after_message_is_deleted(): + message = PointCloud2( + height=1, + width=1, + point_step=4, + row_step=4, + fields=[PointField(name="x", datatype=7, count=1, offset=0)], + data=np.frombuffer(struct.pack("u2"), encoding="16UC1", header=header) + # Add nonzero row padding to big-endian depth. + payload = bytes(depth.data) + depth.data = payload[:6] + b"xx" + payload[6:] + b"yy" + depth.step = 8 + calibration = camera_info_from_intrinsics(4, 8, 2, 2, 6, 4, header=header) + cloud = cdr_decode( + cdr_encode( + pointcloud_from_rgbd(color, depth, calibration, depth_scale=0.001, depth_trunc=4) + ), + PointCloud2, + ) + np.testing.assert_allclose( + pointcloud_xyz(cloud), [[-0.5, -0.25, 1], [1, -0.5, 2], [-1.5, 0, 3]] + ) + np.testing.assert_array_equal(pointcloud_rgb(cloud), rgb.reshape(-1, 3)[[0, 2, 3]]) + assert cloud.header == header + + +@pytest.mark.parametrize( + "scale,trunc", [(0, 1), (-1, 1), (float("nan"), 1), (1, 0), (1, float("inf"))] +) +def test_rgbd_rejects_invalid_units(scale, trunc): + header = Header(stamp=Time(sec=0, nanosec=0), frame_id="") + color = image_from_array(np.zeros((1, 1, 3), dtype=np.uint8), encoding="rgb8", header=header) + depth = image_from_array(np.ones((1, 1), dtype=np.uint16), encoding="16UC1", header=header) + calibration = camera_info_from_intrinsics(1, 1, 0, 0, 1, 1, header=header) + with pytest.raises(ValueError, match="finite and positive"): + pointcloud_from_rgbd(color, depth, calibration, depth_scale=scale, depth_trunc=trunc) + + +def test_transform_rejects_coordinate_overflow(): + cloud = pointcloud_from_xyz( + np.array([[3e38, 0, 0]]), header=Header(frame_id="child", stamp=Time(sec=0, nanosec=0)) + ) + transform = TransformStamped( + header=Header(frame_id="parent", stamp=Time(sec=0, nanosec=0)), + child_frame_id="child", + transform=Transform( + translation=Vector3(x=3e38, y=0.0, z=0.0), rotation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), + ) + with pytest.raises(ValueError, match="exceed field range"): + transform_cloud(cloud, transform) + assert np.isfinite(pointcloud_xyz(cloud)).all() + + +def test_voxel_downsample_retains_existing_small_cloud_fast_path(): + message = pointcloud_from_xyz( + np.zeros((19, 3)), header=Header(frame_id="lidar", stamp=Time(sec=0, nanosec=0)) + ) + assert voxel_downsample_cloud(message, 0.5) is message + assert voxel_downsample_cloud(message, 0.0) is message + with pytest.raises(ValueError, match="finite"): + voxel_downsample_cloud(message, float("nan")) + + +def test_voxel_downsample_native_centroids_colors_and_exact_header(): + pytest.importorskip("open3d", reason="native tensor downsampling requires Open3D") + points = np.repeat([[0.1, 0.1, 0.1], [0.3, 0.3, 0.3]], 20, axis=0) + colors = np.repeat([[0, 0, 0], [255, 255, 255]], 20, axis=0).astype(np.uint8) + header = Header(frame_id="optical", stamp=Time(sec=5, nanosec=123456789)) + source = pointcloud_from_xyz_rgb(points, colors, header=header) + result = voxel_downsample_cloud(source, 1.0) + assert result.header == header + assert result.width == 1 + np.testing.assert_allclose(pointcloud_xyz(result), [[0.2, 0.2, 0.2]], atol=1e-6) + np.testing.assert_array_equal(pointcloud_rgb(result), [[127, 127, 127]]) + assert source.width == 40 + + +@pytest.mark.parametrize("shape", [(3, 3), (2, 3, 3)]) +def test_per_point_stamps_survive_cdr_and_point_selection(shape): + points = np.arange(np.prod(shape), dtype=np.float32).reshape(shape) + stamps = np.arange(np.prod(shape[:-1]), dtype=np.float64).reshape(shape[:-1]) + 17.25 + header = Header(stamp=Time(sec=19, nanosec=123), frame_id="map") + cloud = pointcloud_from_xyz(points, header=header, stamps=stamps) + restored = cdr_decode(cdr_encode(cloud), PointCloud2) + np.testing.assert_array_equal(pointcloud_xyz(restored), points.reshape(-1, 3)) + np.testing.assert_array_equal(pointcloud_stamps(restored), stamps.reshape(-1)) + keep = np.arange(stamps.size) % 2 == 0 + selected = select_points(restored, keep) + np.testing.assert_array_equal(pointcloud_stamps(selected), stamps.reshape(-1)[keep]) + assert selected.header == header + assert pointcloud_stamps(pointcloud_from_xyz(points, header=header)) is None + + +@pytest.mark.parametrize("stamps", [np.array([1.0]), np.array([1.0, np.nan])]) +def test_per_point_stamps_reject_wrong_count_or_nonfinite_values(stamps): + with pytest.raises(ValueError, match="Per-point stamps"): + pointcloud_from_xyz( + np.zeros((2, 3)), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + stamps=stamps, + ) diff --git a/dimos/message_utils/test_pointcloud_colors.py b/dimos/message_utils/test_pointcloud_colors.py new file mode 100644 index 00000000000..4043a764772 --- /dev/null +++ b/dimos/message_utils/test_pointcloud_colors.py @@ -0,0 +1,85 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +import numpy as np +import pytest + +from dimos.message_utils.pointcloud import pointcloud_rgb +from dimos.messages.builtin_interfaces.msg import Time +from dimos.messages.sensor_msgs.msg import PointCloud2, PointField +from dimos.messages.std_msgs.msg import Header + + +@pytest.mark.parametrize("bigendian", [False, True]) +@pytest.mark.parametrize("name", ["rgb", "rgba"]) +@pytest.mark.parametrize("datatype", [PointField.FLOAT32, PointField.UINT32]) +def test_packed_colors_preserve_bits_and_row_padding(bigendian, name, datatype): + order = ">" if bigendian else "<" + data = bytearray(48) + values = np.ndarray((2, 2), dtype=order + "u4", buffer=data, offset=4, strides=(24, 8)) + values[:] = [[0xFF123456, 0x00ABCDEF], [0x00FF0000, 0xFF00FF00]] + cloud = PointCloud2( + height=2, + width=2, + point_step=8, + row_step=24, + is_bigendian=bigendian, + data=np.frombuffer(bytes(data), dtype=np.uint8), + fields=[PointField(name=name, offset=4, count=1, datatype=datatype)], + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + is_dense=False, + ) + decoded = cdr_decode(cdr_encode(cloud), PointCloud2) + np.testing.assert_array_equal( + pointcloud_rgb(decoded), [[0x12, 0x34, 0x56], [0xAB, 0xCD, 0xEF], [255, 0, 0], [0, 255, 0]] + ) + + +def test_missing_colors(): + assert ( + pointcloud_rgb( + PointCloud2( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + height=0, + width=0, + fields=[], + is_bigendian=False, + point_step=0, + row_step=0, + data=np.array([], dtype=np.uint8), + is_dense=False, + ) + ) + is None + ) + + +@pytest.mark.parametrize( + "datatype,count", [(PointField.FLOAT64, 1), (PointField.INT32, 1), (PointField.UINT32, 2)] +) +def test_invalid_color_layout(datatype, count): + cloud = PointCloud2( + height=1, + width=1, + point_step=8, + row_step=8, + data=np.frombuffer(bytes(8), dtype=np.uint8), + fields=[PointField(name="rgb", offset=0, count=count, datatype=datatype)], + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + is_bigendian=False, + is_dense=False, + ) + with pytest.raises(ValueError, match="packed point cloud color"): + pointcloud_rgb(cloud) From 65879496be4524a1e8e693a53fd3c30b240f4751 Mon Sep 17 00:00:00 2001 From: cc Date: Fri, 23 Oct 2026 02:15:45 -0700 Subject: [PATCH 04/10] ci(cdr): run this adapter layer tests before SDK checks --- .github/workflows/cdr-native.yml | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/.github/workflows/cdr-native.yml b/.github/workflows/cdr-native.yml index 4b36e8845e8..08a5e34f6ed 100644 --- a/.github/workflows/cdr-native.yml +++ b/.github/workflows/cdr-native.yml @@ -31,6 +31,13 @@ jobs: sudo apt-get install -y cmake pkg-config clang libclang-dev portaudio19-dev libturbojpeg ffmpeg - name: Install Python project and test dependencies run: uv sync --frozen --group tests + - name: Test generated value adapters + run: | + uv run --no-sync pytest --noconftest --import-mode=importlib -o addopts='' -q \ + dimos/message_utils/test_image_views.py \ + dimos/message_utils/test_image_compression.py \ + dimos/message_utils/test_camera_info.py \ + dimos/message_utils/test_time.py - name: Test native SDK run: | cargo test --locked --manifest-path native/rust/dimos-module/Cargo.toml From f4231128132d5ff25e15f4f6399626b7c266e28f Mon Sep 17 00:00:00 2001 From: cc Date: Fri, 23 Oct 2026 02:15:46 -0700 Subject: [PATCH 05/10] ci(cdr): run this adapter layer tests before SDK checks --- .github/workflows/cdr-native.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/cdr-native.yml b/.github/workflows/cdr-native.yml index 08a5e34f6ed..86472495100 100644 --- a/.github/workflows/cdr-native.yml +++ b/.github/workflows/cdr-native.yml @@ -37,7 +37,8 @@ jobs: dimos/message_utils/test_image_views.py \ dimos/message_utils/test_image_compression.py \ dimos/message_utils/test_camera_info.py \ - dimos/message_utils/test_time.py + dimos/message_utils/test_time.py \ + dimos/message_utils/test_geometry.py - name: Test native SDK run: | cargo test --locked --manifest-path native/rust/dimos-module/Cargo.toml From b13692563191ef0484dd71ff3e69ae2dbdd11ca3 Mon Sep 17 00:00:00 2001 From: cc Date: Fri, 23 Oct 2026 02:15:46 -0700 Subject: [PATCH 06/10] ci(cdr): run this adapter layer tests before SDK checks --- .github/workflows/cdr-native.yml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/workflows/cdr-native.yml b/.github/workflows/cdr-native.yml index 86472495100..21b5b1eed64 100644 --- a/.github/workflows/cdr-native.yml +++ b/.github/workflows/cdr-native.yml @@ -38,7 +38,9 @@ jobs: dimos/message_utils/test_image_compression.py \ dimos/message_utils/test_camera_info.py \ dimos/message_utils/test_time.py \ - dimos/message_utils/test_geometry.py + dimos/message_utils/test_geometry.py \ + dimos/message_utils/test_pointcloud.py \ + dimos/message_utils/test_pointcloud_colors.py - name: Test native SDK run: | cargo test --locked --manifest-path native/rust/dimos-module/Cargo.toml From 950d753be0ccbd65b52d2597943afa1ceeba7c95 Mon Sep 17 00:00:00 2001 From: cc Date: Mon, 26 Oct 2026 11:34:06 -0700 Subject: [PATCH 07/10] refactor: namespace sensor helpers and clarify ownership --- dimos/message_utils/camera_info.py | 21 +++-- dimos/message_utils/image.py | 79 +++++++++------- dimos/message_utils/test_camera_info.py | 28 +++--- dimos/message_utils/test_image_compression.py | 92 +++++++++---------- dimos/message_utils/test_image_views.py | 89 +++++++++++------- dimos/message_utils/time.py | 7 -- dimos/visualization/rerun/cdr.py | 12 +-- 7 files changed, 172 insertions(+), 156 deletions(-) diff --git a/dimos/message_utils/camera_info.py b/dimos/message_utils/camera_info.py index c8040990805..eb82210ba7d 100644 --- a/dimos/message_utils/camera_info.py +++ b/dimos/message_utils/camera_info.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Calibration construction and array access for generated ROS CameraInfo values.""" +"""Calibration operations for generated ROS CameraInfo wire messages. + +Import with ``from dimos.message_utils import camera_info``. Constructors return +CameraInfo messages; ``intrinsic_matrix`` returns an independent writable array. +``with_stamp`` deep-copies the calibration so prior publishes remain unchanged. +""" from copy import deepcopy import math @@ -37,7 +42,7 @@ def _dimensions(width: int, height: int) -> None: raise ValueError("Calibration width and height must be positive integers") -def camera_info_from_intrinsics( +def from_intrinsics( fx: float, fy: float, cx: float, cy: float, width: int, height: int, *, header: Header ) -> CameraInfo: """Construct an undistorted pinhole calibration with an explicit source header.""" @@ -59,7 +64,7 @@ def camera_info_from_intrinsics( ) -def camera_info_from_fov( +def from_fov( fov_degrees: float, width: int, height: int, @@ -75,9 +80,7 @@ def camera_info_from_fov( raise ValueError("Field-of-view axis must be horizontal or vertical") extent = width if axis == "horizontal" else height focal = extent / (2 * math.tan(math.radians(fov_degrees) / 2)) - return camera_info_from_intrinsics( - focal, focal, width / 2, height / 2, width, height, header=header - ) + return from_intrinsics(focal, focal, width / 2, height / 2, width, height, header=header) def _matrix(data: dict[str, Any], name: str, rows: int, columns: int | None) -> list[float]: @@ -100,7 +103,7 @@ def _matrix(data: dict[str, Any], name: str, rows: int, columns: int | None) -> return converted -def camera_info_from_yaml(path: str | Path, *, header: Header) -> CameraInfo: +def from_yaml(path: str | Path, *, header: Header) -> CameraInfo: """Read a ROS calibration YAML file, validating declared matrix dimensions.""" with Path(path).open() as source: data = yaml.safe_load(source) @@ -128,7 +131,7 @@ def intrinsic_matrix(message: CameraInfo) -> NDArray[np.float64]: return np.array(message.k, dtype=np.float64, copy=True).reshape(3, 3) -def camera_info_with_stamp(message: CameraInfo, stamp: Time) -> CameraInfo: +def with_stamp(message: CameraInfo, stamp: Time) -> CameraInfo: """Copy a calibration with an exact source stamp, without mutating prior publishes.""" result = deepcopy(message) result.header.stamp = deepcopy(stamp) @@ -162,7 +165,7 @@ def __getattr__(self, name: str) -> CameraInfo: for key in (name, snake): path = self._calibration_dir / f"{key}.yaml" if path.is_file(): - value = camera_info_from_yaml( + value = from_yaml( path, header=Header(frame_id=self._frame_id, stamp=Time(sec=0, nanosec=0)) ) self._cache[name] = value diff --git a/dimos/message_utils/image.py b/dimos/message_utils/image.py index 16548a5f175..c26a8923402 100644 --- a/dimos/message_utils/image.py +++ b/dimos/message_utils/image.py @@ -12,7 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Array views and compression for generated ROS image messages.""" +"""Image operations for generated ROS wire messages. + +Import this module with ``from dimos.message_utils import image`` and call +``image.from_array(pixels, encoding="rgb8")`` or ``image.view(message)``. +``view`` borrows the pixel buffer; ``to_rgb`` and ``to_bgr`` return independent +writable arrays. ``to_jpeg`` returns bytes and ``to_compressed`` returns a +CompressedImage message. Generated message classes remain wire data. +""" from __future__ import annotations @@ -44,8 +51,8 @@ } -def image_from_array(pixels: NDArray[Any], *, encoding: str, header: Header | None = None) -> Image: - """Copy an array into a generated image with an explicit pixel encoding.""" +def from_array(pixels: NDArray[Any], *, encoding: str, header: Header | None = None) -> Image: + """Copy pixels into an Image with an explicit encoding, reusing the supplied header.""" if encoding not in _FORMATS: raise ValueError(f"unsupported image encoding {encoding!r}") kind, channels = _FORMATS[encoding] @@ -68,8 +75,13 @@ def image_from_array(pixels: NDArray[Any], *, encoding: str, header: Header | No ) -def image_view(msg: Image) -> NDArray[Any]: - """Borrow a read-only array, retaining the message and respecting row padding/endian.""" +def view(msg: Image) -> NDArray[Any]: + """Borrow a read-only array that respects row padding and byte order. + + The array keeps the data buffer alive after the message is released. It is + not a snapshot: writes through the message's buffer remain visible. Use + ``view(message).copy()`` for independent writable pixels. + """ if msg.encoding not in _FORMATS: raise ValueError(f"unsupported image encoding {msg.encoding!r}") kind, channels = _FORMATS[msg.encoding] @@ -88,12 +100,12 @@ def image_view(msg: Image) -> NDArray[Any]: return view -def image_to_jpeg(msg: Image, quality: int = 75) -> bytes: - """Encode an 8-bit gray or color image using OpenCV's JPEG encoder.""" +def to_jpeg(msg: Image, quality: int = 75) -> bytes: + """Return JPEG bytes for an 8-bit gray or color image, without message metadata.""" if isinstance(quality, bool) or not 0 <= quality <= 100: raise ValueError("JPEG quality must be in 0..100") - pixels = image_view(msg) + pixels = view(msg) conversions = { "rgb8": cv2.COLOR_RGB2BGR, "rgba8": cv2.COLOR_RGBA2BGR, @@ -109,7 +121,7 @@ def image_to_jpeg(msg: Image, quality: int = 75) -> bytes: return bytes(encoded) -def compressed_image_from_image( +def to_compressed( message: Image, quality: int = 75, *, @@ -117,7 +129,7 @@ def compressed_image_from_image( max_width: int | None = None, effort: int | None = None, ) -> CompressedImage: - """Compress generated pixels, retaining the exact header. + """Return a CompressedImage with encoded pixels and the source header. PNG preserves uint16 depth. JXL preserves uint16/float32 depth losslessly; it requires imagecodecs and is an explicit codec extension for consumers. @@ -125,12 +137,12 @@ def compressed_image_from_image( if not isinstance(message, Image): raise TypeError("compression expects a generated Image") if max_width is not None: - message, _ = image_resize_to_fit(message, max_width, max_width) + message, _ = resize_to_fit(message, max_width, max_width) if format == "jpeg": output_encoding = "mono8" if message.encoding == "mono8" else "bgr8" - data = image_to_jpeg(message, quality=quality) + data = to_jpeg(message, quality=quality) elif format == "png": - pixels = image_view(message) + pixels = view(message) if pixels.dtype.kind != "u" or pixels.dtype.itemsize not in (1, 2): raise ValueError("PNG cannot encode floating-point depth") if message.encoding in ("rgb8", "rgba8"): @@ -143,7 +155,7 @@ def compressed_image_from_image( data = bytes(encoded) output_encoding = {"rgb8": "bgr8", "rgba8": "bgra8"}.get(message.encoding, message.encoding) elif format == "jxl": - pixels = image_view(message) + pixels = view(message) if pixels.dtype.kind not in ("u", "f") or ( pixels.dtype.kind, pixels.dtype.itemsize, @@ -168,7 +180,7 @@ def compressed_image_from_image( ) -def image_from_compressed(message: CompressedImage) -> Image: +def from_compressed(message: CompressedImage) -> Image: """Decode JPEG/PNG or the JXL extension into explicitly encoded generated pixels.""" format_name = message.format.lower() @@ -188,7 +200,7 @@ def image_from_compressed(message: CompressedImage) -> Image: raise ValueError("unsupported decoded JXL layout") if encoding == "mono16" and format_name.startswith("16uc1;"): encoding = "16UC1" - return image_from_array(jxl_pixels, encoding=encoding, header=message.header) + return from_array(jxl_pixels, encoding=encoding, header=message.header) if not any(name in format_name for name in ("jpeg", "jpg", "png")): raise ValueError(f"unsupported compressed image format {message.format!r}") pixels = cv2.imdecode(np.frombuffer(bytes(message.data), dtype=np.uint8), cv2.IMREAD_UNCHANGED) @@ -204,26 +216,26 @@ def image_from_compressed(message: CompressedImage) -> Image: raise ValueError("unsupported decoded image layout") if encoding == "mono16" and format_name.startswith("16uc1;"): encoding = "16UC1" - return image_from_array(pixels, encoding=encoding, header=message.header) + return from_array(pixels, encoding=encoding, header=message.header) -def image_brightness(message: Image) -> float: +def brightness(message: Image) -> float: """Sample mean pixel intensity in [0, 1] for 8/16-bit image encodings. Respect row padding and endian. Floating-point depth has no normalized intensity scale and is rejected rather than treated as a color image. """ - pixels = image_view(message) + pixels = view(message) if pixels.dtype.kind != "u": raise ValueError("brightness requires unsigned integer pixels") stride = max(1, max(message.height, message.width) // 256) return float(pixels[::stride, ::stride].mean() / np.iinfo(pixels.dtype).max) -def image_sharpness(message: Image) -> float: +def sharpness(message: Image) -> float: """Laplacian variance of an 8-bit visual image, downsampled to 160 pixels wide.""" - pixels = image_view(message) + pixels = view(message) codes = { "rgb8": cv2.COLOR_RGB2GRAY, "bgr8": cv2.COLOR_BGR2GRAY, @@ -247,10 +259,10 @@ def image_sharpness(message: Image) -> float: return float(cv2.Laplacian(gray, cv2.CV_64F).var()) -def image_to_bgr(message: Image) -> NDArray[np.uint8]: +def to_bgr(message: Image) -> NDArray[np.uint8]: """Return an independent BGR8 array for drawing or OpenCV color operations.""" - pixels = image_view(message) + pixels = view(message) if message.encoding == "bgr8": return pixels.copy() codes = { @@ -264,14 +276,14 @@ def image_to_bgr(message: Image) -> NDArray[np.uint8]: return np.asarray(cv2.cvtColor(pixels, codes[message.encoding]), dtype=np.uint8) -def image_to_rgb(message: Image) -> NDArray[np.uint8]: +def to_rgb(message: Image) -> NDArray[np.uint8]: """Copy a display image as RGB8; 16-bit grayscale uses its high byte. Alpha is discarded. Floating-point depth needs an explicit visualization scale and is rejected rather than assigning pixel values to meters. """ - pixels = image_view(message) + pixels = view(message) if message.encoding == "rgb8": return pixels.copy() if message.encoding in ("mono16", "16UC1"): @@ -289,8 +301,13 @@ def image_to_rgb(message: Image) -> NDArray[np.uint8]: return np.asarray(cv2.cvtColor(pixels, codes[message.encoding]), dtype=np.uint8) -def image_resize_to_fit(message: Image, max_width: int, max_height: int) -> tuple[Image, float]: - """Downscale a generated image while preserving encoding, aspect and exact header.""" +def resize_to_fit(message: Image, max_width: int, max_height: int) -> tuple[Image, float]: + """Return an image fitting the bounds and the requested downscale factor. + + If it already fits, return the original message and 1.0. Otherwise allocate + independent pixels, preserve the encoding and reuse the source header. + Integer output dimensions may round down from the scaled aspect ratio. + """ if min(max_width, max_height, message.width, message.height) <= 0: raise ValueError("Image and target dimensions must be positive") @@ -298,14 +315,14 @@ def image_resize_to_fit(message: Image, max_width: int, max_height: int) -> tupl return message, 1.0 scale = min(max_width / message.width, max_height / message.height) pixels = cv2.resize( - image_view(message), + view(message), (max(1, int(message.width * scale)), max(1, int(message.height * scale))), interpolation=cv2.INTER_LINEAR, ) - return image_from_array(pixels, encoding=message.encoding, header=message.header), scale + return from_array(pixels, encoding=message.encoding, header=message.header), scale -def image_from_file(path: str | Path, *, header: Header | None = None) -> Image: +def from_file(path: str | Path, *, header: Header | None = None) -> Image: """Decode an image file as generated RGB8 with an explicit optional source header.""" with PILImage.open(path) as image: - return image_from_array(np.asarray(image.convert("RGB")), encoding="rgb8", header=header) + return from_array(np.asarray(image.convert("RGB")), encoding="rgb8", header=header) diff --git a/dimos/message_utils/test_camera_info.py b/dimos/message_utils/test_camera_info.py index bcef99230c3..65dbfb0595d 100644 --- a/dimos/message_utils/test_camera_info.py +++ b/dimos/message_utils/test_camera_info.py @@ -22,13 +22,7 @@ import pytest import yaml -from dimos.message_utils.camera_info import ( - camera_info_from_fov, - camera_info_from_intrinsics, - camera_info_from_yaml, - camera_info_with_stamp, - intrinsic_matrix, -) +from dimos.message_utils import camera_info from dimos.messages.builtin_interfaces.msg import Time from dimos.messages.sensor_msgs.msg import CameraInfo from dimos.messages.std_msgs.msg import Header @@ -39,7 +33,7 @@ def test_real_go2_calibration_survives_cdr_with_exact_header(): header = Header(frame_id="camera_optical", stamp=Time(sec=1700000000, nanosec=123456789)) source = yaml.safe_load(_CALIBRATION.read_text()) - message = camera_info_from_yaml(_CALIBRATION, header=header) + message = camera_info.from_yaml(_CALIBRATION, header=header) decoded = cdr_decode(cdr_encode(message), CameraInfo) assert decoded.header == header assert (decoded.width, decoded.height) == (1280, 720) @@ -51,7 +45,7 @@ def test_real_go2_calibration_survives_cdr_with_exact_header(): ("p", "projection_matrix"), ]: np.testing.assert_array_equal(getattr(decoded, field), source[key]["data"]) - matrix = intrinsic_matrix(decoded) + matrix = camera_info.intrinsic_matrix(decoded) assert matrix.shape == (3, 3) matrix[0, 0] = 0 assert decoded.k[0] == source["camera_matrix"]["data"][0] @@ -61,7 +55,7 @@ def test_real_go2_calibration_survives_cdr_with_exact_header(): @pytest.mark.parametrize("axis,focal", [("vertical", 240), ("horizontal", 320)]) def test_fov_converts_selected_axis_and_builds_projection(axis, focal): - message = camera_info_from_fov( + message = camera_info.from_fov( 90, 640, 480, header=Header(frame_id="optical", stamp=Time(sec=0, nanosec=0)), axis=axis ) np.testing.assert_allclose(message.k, [focal, 0, 320, 0, focal, 240, 0, 0, 1]) @@ -73,7 +67,7 @@ def test_fov_converts_selected_axis_and_builds_projection(axis, focal): @pytest.mark.parametrize("fov", [0, -1, 180, 181, math.nan, math.inf]) def test_fov_rejects_nonphysical_angles(fov): with pytest.raises(ValueError, match="Field of view"): - camera_info_from_fov( + camera_info.from_fov( fov, 640, 480, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="") ) @@ -81,7 +75,7 @@ def test_fov_rejects_nonphysical_angles(fov): @pytest.mark.parametrize("width,height", [(0, 480), (640, -1), (True, 480), (1.5, 480)]) def test_intrinsics_reject_invalid_dimensions(width, height): with pytest.raises(ValueError, match="positive integers"): - camera_info_from_intrinsics( + camera_info.from_intrinsics( 300, 300, 320, @@ -95,7 +89,7 @@ def test_intrinsics_reject_invalid_dimensions(width, height): @pytest.mark.parametrize("fx", [0, -1, math.nan, math.inf]) def test_intrinsics_reject_invalid_focal_length(fx): with pytest.raises(ValueError, match="positive focal"): - camera_info_from_intrinsics( + camera_info.from_intrinsics( fx, 300, 320, 240, 640, 480, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="") ) @@ -113,11 +107,11 @@ def test_yaml_rejects_invalid_camera_matrix(tmp_path, mutation): path = tmp_path / "invalid.yaml" path.write_text(yaml.safe_dump(data)) with pytest.raises(ValueError, match="camera_matrix"): - camera_info_from_yaml(path, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")) + camera_info.from_yaml(path, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")) def test_stamp_copy_preserves_calibration_and_does_not_mutate_prior_publish(): - original = camera_info_from_intrinsics( + original = camera_info.from_intrinsics( 500, 510, 320, @@ -129,7 +123,7 @@ def test_stamp_copy_preserves_calibration_and_does_not_mutate_prior_publish(): original.binning_x = 2 original.roi.x_offset = 12 stamp = Time(sec=2, nanosec=123456789) - copied = camera_info_with_stamp(original, stamp) + copied = camera_info.with_stamp(original, stamp) assert copied.header.stamp == stamp assert copied.header.frame_id == original.header.frame_id assert list(copied.k) == list(original.k) @@ -146,4 +140,4 @@ def test_calibration_yaml_rejects_non_mapping_input(tmp_path): source = tmp_path / "calibration.yaml" source.write_text("- width\n- height\n") with pytest.raises(TypeError, match="Calibration YAML must contain a mapping"): - camera_info_from_yaml(source, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")) + camera_info.from_yaml(source, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")) diff --git a/dimos/message_utils/test_image_compression.py b/dimos/message_utils/test_image_compression.py index 0514d681424..2cf600fa416 100644 --- a/dimos/message_utils/test_image_compression.py +++ b/dimos/message_utils/test_image_compression.py @@ -17,14 +17,7 @@ import numpy as np import pytest -from dimos.message_utils.image import ( - compressed_image_from_image, - image_from_array, - image_from_compressed, - image_to_bgr, - image_to_rgb, - image_view, -) +from dimos.message_utils import image from dimos.message_utils.time import time_from_seconds, to_nanoseconds from dimos.messages.builtin_interfaces.msg import Time from dimos.messages.sensor_msgs.msg import CompressedImage @@ -37,43 +30,44 @@ def rgb_image(): gradient = np.linspace(0, 255, 640, dtype=np.uint8) data = np.broadcast_to(gradient, (480, 640)).copy() data = np.stack([data, data // 2, rng.randint(0, 50, (480, 640), dtype=np.uint8)], axis=-1) - return image_from_array( + return image.from_array( data, encoding="rgb8", header=Header(frame_id="cam", stamp=time_from_seconds(1234.5678)) ) def test_jpeg_roundtrip_preserves_meta_and_pixels(rgb_image) -> None: - ci = compressed_image_from_image(rgb_image, quality=90) + ci = image.to_compressed(rgb_image, quality=90) + assert isinstance(ci, CompressedImage) assert "; jpeg compressed " in ci.format assert ci.header.frame_id == "cam" assert to_nanoseconds(ci.header.stamp) == to_nanoseconds(rgb_image.header.stamp) - assert 0 < len(ci.data) < image_view(rgb_image).nbytes // 4 - img = image_from_compressed(ci) + assert 0 < len(ci.data) < image.view(rgb_image).nbytes // 4 + img = image.from_compressed(ci) assert img.encoding in ("rgb8", "bgr8") assert img.header.frame_id == "cam" assert to_nanoseconds(img.header.stamp) == to_nanoseconds(rgb_image.header.stamp) - assert image_view(img).shape == image_view(rgb_image).shape - diff = np.abs(image_to_rgb(img).astype(int) - image_view(rgb_image).astype(int)).mean() + assert image.view(img).shape == image.view(rgb_image).shape + diff = np.abs(image.to_rgb(img).astype(int) - image.view(rgb_image).astype(int)).mean() assert diff < 5, f"JPEG q90 mean pixel error too high: {diff}" def test_png_roundtrip_is_lossless_bgr(rgb_image) -> None: - bgr = image_from_array(image_to_bgr(rgb_image), encoding="bgr8", header=rgb_image.header) - ci = compressed_image_from_image(bgr, format="png") - img = image_from_compressed(ci) + bgr = image.from_array(image.to_bgr(rgb_image), encoding="bgr8", header=rgb_image.header) + ci = image.to_compressed(bgr, format="png") + img = image.from_compressed(ci) assert img.encoding == "bgr8" - assert np.array_equal(image_view(img), image_view(bgr)) + assert np.array_equal(image.view(img), image.view(bgr)) assert to_nanoseconds(img.header.stamp) == to_nanoseconds(bgr.header.stamp) def test_png_roundtrip_is_lossless_gray16() -> None: data = np.arange(100 * 80, dtype=np.uint16).reshape(100, 80) - src = image_from_array( + src = image.from_array( data, encoding="mono16", header=Header(frame_id="d", stamp=time_from_seconds(1.0)) ) - img = image_from_compressed(compressed_image_from_image(src, format="png")) + img = image.from_compressed(image.to_compressed(src, format="png")) assert img.encoding == "mono16" - assert np.array_equal(image_view(img), data) + assert np.array_equal(image.view(img), data) def test_png_decode_alpha_is_bgra() -> None: @@ -81,7 +75,7 @@ def test_png_decode_alpha_is_bgra() -> None: arr[..., 3] = 128 ok, buf = cv2.imencode(".png", arr) assert ok - img = image_from_compressed( + img = image.from_compressed( CompressedImage( data=np.frombuffer(buf.tobytes(), dtype=np.uint8), format="png", @@ -89,11 +83,11 @@ def test_png_decode_alpha_is_bgra() -> None: ) ) assert img.encoding == "bgra8" - assert np.array_equal(image_view(img), arr) + assert np.array_equal(image.view(img), arr) def test_lcm_roundtrip(rgb_image) -> None: - ci = compressed_image_from_image(rgb_image) + ci = image.to_compressed(rgb_image) wire = cdr_encode(ci) out = cdr_decode(wire, CompressedImage) np.testing.assert_array_equal(out.data, ci.data) @@ -103,85 +97,85 @@ def test_lcm_roundtrip(rgb_image) -> None: def test_jpeg_rejects_depth_formats() -> None: - depth = image_from_array( + depth = image.from_array( np.zeros((10, 10), dtype=np.uint16), encoding="mono16", header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) with pytest.raises(ValueError, match="JPEG requires 8-bit"): - compressed_image_from_image(depth) + image.to_compressed(depth) def test_png_rejects_float_depth() -> None: - depth = image_from_array( + depth = image.from_array( np.zeros((10, 10), dtype=np.float32), encoding="32FC1", header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) with pytest.raises(ValueError, match="PNG cannot encode"): - compressed_image_from_image(depth, format="png") + image.to_compressed(depth, format="png") def test_from_image_rejects_non_image() -> None: with pytest.raises(TypeError): - compressed_image_from_image(b"not an image") + image.to_compressed(b"not an image") def test_max_width_resizes(rgb_image) -> None: - img = image_from_compressed(compressed_image_from_image(rgb_image, max_width=320)) + img = image.from_compressed(image.to_compressed(rgb_image, max_width=320)) assert img.width <= 320 assert img.height <= 320 def test_jxl_roundtrip_is_lossless_gray16() -> None: data = np.arange(100 * 80, dtype=np.uint16).reshape(100, 80) - src = image_from_array( + src = image.from_array( data, encoding="mono16", header=Header(frame_id="d", stamp=time_from_seconds(1.0)) ) - ci = compressed_image_from_image(src, format="jxl") + ci = image.to_compressed(src, format="jxl") assert "; jxl compressed " in ci.format - img = image_from_compressed(ci) + img = image.from_compressed(ci) assert img.encoding == "mono16" - assert np.array_equal(image_view(img), data) + assert np.array_equal(image.view(img), data) def test_jxl_roundtrip_is_lossless_float_depth() -> None: data = np.linspace(0.1, 10.0, 100 * 80, dtype=np.float32).reshape(100, 80) - src = image_from_array( + src = image.from_array( data, encoding="32FC1", header=Header(frame_id="d", stamp=time_from_seconds(2.0)) ) - ci = compressed_image_from_image(src, format="jxl") - img = image_from_compressed(ci) + ci = image.to_compressed(src, format="jxl") + img = image.from_compressed(ci) assert img.encoding == "32FC1" - assert np.array_equal(image_view(img), data) + assert np.array_equal(image.view(img), data) assert to_nanoseconds(img.header.stamp) == to_nanoseconds(src.header.stamp) assert img.header.frame_id == "d" def test_jxl_lossy_rgb_honors_quality(rgb_image) -> None: - ci = compressed_image_from_image(rgb_image, format="jxl", quality=90) - assert 0 < len(ci.data) < image_view(rgb_image).nbytes // 4 - img = image_from_compressed(ci) + ci = image.to_compressed(rgb_image, format="jxl", quality=90) + assert 0 < len(ci.data) < image.view(rgb_image).nbytes // 4 + img = image.from_compressed(ci) assert img.encoding in ("rgb8", "bgr8") - diff = np.abs(image_to_rgb(img).astype(int) - image_view(rgb_image).astype(int)).mean() + diff = np.abs(image.to_rgb(img).astype(int) - image.view(rgb_image).astype(int)).mean() assert diff < 5, f"JXL q90 mean pixel error too high: {diff}" def test_jxl_effort_trades_cpu_for_size(rgb_image) -> None: - fast = compressed_image_from_image(rgb_image, format="jxl", effort=1) - thorough = compressed_image_from_image(rgb_image, format="jxl", effort=7) + fast = image.to_compressed(rgb_image, format="jxl", effort=1) + thorough = image.to_compressed(rgb_image, format="jxl", effort=7) assert len(thorough.data) <= len(fast.data) def test_jxl_rejects_float64() -> None: - depth = image_from_array(np.zeros((10, 10), dtype=np.float64), encoding="64FC1") + depth = image.from_array(np.zeros((10, 10), dtype=np.float64), encoding="64FC1") with pytest.raises(ValueError, match="JXL cannot encode"): - compressed_image_from_image(depth, format="jxl") + image.to_compressed(depth, format="jxl") @pytest.mark.parametrize("codec", ["png", "jxl"]) def test_compressed_big_endian_depth(codec: str) -> None: pixels = np.arange(80, dtype=">u2").reshape(10, 8) * 257 - source = image_from_array(pixels.astype(">u2"), encoding="16UC1") - decoded = image_from_compressed(compressed_image_from_image(source, format=codec)) - np.testing.assert_array_equal(image_view(decoded), pixels) + source = image.from_array(pixels.astype(">u2"), encoding="16UC1") + decoded = image.from_compressed(image.to_compressed(source, format=codec)) + np.testing.assert_array_equal(image.view(decoded), pixels) diff --git a/dimos/message_utils/test_image_views.py b/dimos/message_utils/test_image_views.py index a4d0680bcaf..38a7ed025d4 100644 --- a/dimos/message_utils/test_image_views.py +++ b/dimos/message_utils/test_image_views.py @@ -22,15 +22,7 @@ from PIL import Image as PILImage import pytest -from dimos.message_utils.image import ( - image_brightness, - image_from_array, - image_from_file, - image_sharpness, - image_to_jpeg, - image_to_rgb, - image_view, -) +from dimos.message_utils import image from dimos.message_utils.time import time_from_nanoseconds from dimos.messages.builtin_interfaces.msg import Time from dimos.messages.sensor_msgs.msg import Image @@ -47,9 +39,13 @@ def test_padded_image_view_retains_owner_and_is_readonly() -> None: header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), is_bigendian=0, ) - pixels = image_view(msg) + pixels = image.view(msg) np.testing.assert_array_equal(pixels, [[1, 2], [3, 4]]) assert pixels.strides == (4, 1) + assert np.shares_memory(pixels, msg.data) + msg.data[0] = 5 + assert pixels[0, 0] == 5 + msg.data[0] = 1 with pytest.raises(ValueError): pixels[0, 0] = 9 del msg @@ -70,16 +66,16 @@ def test_big_endian_depth_view() -> None: data=np.array([0x01, 0x02, 0x03, 0x04], dtype=np.uint8), header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) - np.testing.assert_array_equal(image_view(msg), [[258, 772]]) + np.testing.assert_array_equal(image.view(msg), [[258, 772]]) def test_image_array_copy_preserves_endian_and_handles_strides() -> None: pixels = np.arange(24, dtype=">u2").reshape(4, 6)[::2, ::2] - msg = image_from_array(pixels, encoding="16UC1") + msg = image.from_array(pixels, encoding="16UC1") assert (msg.width, msg.height, msg.step, msg.is_bigendian) == (3, 2, 6, 1) expected = pixels.copy() pixels[:] = 0 - np.testing.assert_array_equal(image_view(msg), expected) + np.testing.assert_array_equal(image.view(msg), expected) @pytest.mark.parametrize( @@ -88,13 +84,13 @@ def test_image_array_copy_preserves_endian_and_handles_strides() -> None: ) def test_image_array_requires_matching_encoding(pixels: np.ndarray, encoding: str) -> None: with pytest.raises(ValueError, match="does not match"): - image_from_array(pixels, encoding=encoding) + image.from_array(pixels, encoding=encoding) @pytest.mark.parametrize("step,data", [(1, [1]), (2, [1]), (2, [1, 2, 3])]) def test_invalid_image_layout(step: int, data: list[int]) -> None: with pytest.raises(ValueError, match="dimensions"): - image_view( + image.view( Image( width=2, height=1, @@ -117,7 +113,8 @@ def test_jpeg_preserves_rgb_color() -> None: header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), is_bigendian=0, ) - jpeg = image_to_jpeg(msg, quality=95) + jpeg = image.to_jpeg(msg, quality=95) + assert isinstance(jpeg, bytes) bgr = cv2.imdecode(np.frombuffer(jpeg, np.uint8), cv2.IMREAD_COLOR) assert bgr.shape == (16, 16, 3) assert bgr[8, 8, 2] > 250 and bgr[8, 8, 0] < 5 @@ -129,17 +126,17 @@ def test_sharpness_prefers_edges_to_blurred_pixels(encoding): channels = 1 if encoding == "mono8" else (4 if "a" in encoding else 3) pixels = mono if channels == 1 else np.repeat(mono[:, :, None], channels, axis=2) blurred = cv2.GaussianBlur(pixels, (5, 5), 0) - sharp = image_from_array(pixels, encoding=encoding) - soft = image_from_array(blurred, encoding=encoding) - assert image_sharpness(sharp) > image_sharpness(soft) + sharp = image.from_array(pixels, encoding=encoding) + soft = image.from_array(blurred, encoding=encoding) + assert image.sharpness(sharp) > image.sharpness(soft) def test_sharpness_rejects_depth_and_empty_images(): - depth = image_from_array(np.zeros((2, 2), dtype=np.float32), encoding="32FC1") + depth = image.from_array(np.zeros((2, 2), dtype=np.float32), encoding="32FC1") with pytest.raises(ValueError, match="8-bit visual"): - image_sharpness(depth) + image.sharpness(depth) with pytest.raises(ValueError, match="nonempty"): - image_sharpness( + image.sharpness( Image( encoding="mono8", header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), @@ -162,12 +159,12 @@ def test_sharpness_rejects_depth_and_empty_images(): ], ) def test_rgb_model_input_is_color_correct_and_independent(encoding: str, pixel: list[int]) -> None: - source = image_from_array(np.array([[pixel]], dtype=np.uint8), encoding=encoding) + source = image.from_array(np.array([[pixel]], dtype=np.uint8), encoding=encoding) restored = cdr_decode(cdr_encode(source), Image) - rgb = image_to_rgb(restored) + rgb = image.to_rgb(restored) np.testing.assert_array_equal(rgb, [[[255, 0, 17]]]) rgb[0, 0] = 0 - np.testing.assert_array_equal(image_view(restored), [[pixel]]) + np.testing.assert_array_equal(image.view(restored), [[pixel]]) def test_rgb_model_input_respects_big_endian_padded_grayscale() -> None: @@ -181,14 +178,14 @@ def test_rgb_model_input_respects_big_endian_padded_grayscale() -> None: header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) np.testing.assert_array_equal( - image_to_rgb(msg), [[[18, 18, 18], [255, 255, 255]], [[1, 1, 1], [0, 0, 0]]] + image.to_rgb(msg), [[[18, 18, 18], [255, 255, 255]], [[1, 1, 1], [0, 0, 0]]] ) def test_rgb_model_input_requires_explicit_float_depth_scale() -> None: - msg = image_from_array(np.array([[2.5]], dtype=np.float32), encoding="32FC1") + msg = image.from_array(np.array([[2.5]], dtype=np.float32), encoding="32FC1") with pytest.raises(ValueError, match="RGB8"): - image_to_rgb(msg) + image.to_rgb(msg) def test_brightness_ignores_padding_and_normalizes_big_endian_u16() -> None: @@ -201,10 +198,10 @@ def test_brightness_ignores_padding_and_normalizes_big_endian_u16() -> None: data=np.array([0, 0, 255, 255, 255, 255], dtype=np.uint8), header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), ) - assert image_brightness(message) == pytest.approx(0.5) - depth = image_from_array(np.array([[1.0]], dtype=np.float32), encoding="32FC1") + assert image.brightness(message) == pytest.approx(0.5) + depth = image.from_array(np.array([[1.0]], dtype=np.float32), encoding="32FC1") with pytest.raises(ValueError, match="unsigned integer"): - image_brightness(depth) + image.brightness(depth) @pytest.mark.parametrize("mode", ["L", "RGBA"]) @@ -214,10 +211,10 @@ def test_image_file_decodes_rgb_and_preserves_source_header(tmp_path, mode): source = PILImage.fromarray(pixels).convert(mode) source.save(path) header = Header(stamp=time_from_nanoseconds(1234567890), frame_id="camera") - message = image_from_file(path, header=header) + message = image.from_file(path, header=header) assert message.encoding == "rgb8" assert cdr_encode(message.header) == cdr_encode(header) - np.testing.assert_array_equal(image_view(message), [[[7, 7, 7], [19, 19, 19]]]) + np.testing.assert_array_equal(image.view(message), [[[7, 7, 7], [19, 19, 19]]]) @pytest.mark.parametrize( @@ -226,8 +223,30 @@ def test_image_file_decodes_rgb_and_preserves_source_header(tmp_path, mode): ) def test_remaining_raw_encodings_keep_signed_float_and_endian_values(encoding, dtype, shape): pixels = (np.arange(np.prod(shape)).reshape(shape) - 5).astype(dtype) - source = image_from_array(pixels, encoding=encoding) + source = image.from_array(pixels, encoding=encoding) decoded = cdr_decode(cdr_encode(source), Image) assert decoded.encoding == encoding assert decoded.is_bigendian == 1 - np.testing.assert_array_equal(image_view(decoded), pixels) + np.testing.assert_array_equal(image.view(decoded), pixels) + + +@pytest.mark.parametrize("encoding", ["rgb8", "bgr8"]) +def test_bgr_array_is_writable_and_independent(encoding: str) -> None: + message = image.from_array(np.full((2, 3, 3), 17, dtype=np.uint8), encoding=encoding) + pixels = image.to_bgr(message) + pixels[:] = 0 + assert np.all(image.view(message) == 17) + + +def test_resize_reuses_fitting_message_and_copies_downscaled_pixels() -> None: + message = image.from_array(np.full((4, 8, 3), 17, dtype=np.uint8), encoding="rgb8") + unchanged, scale = image.resize_to_fit(message, 8, 4) + assert unchanged is message + assert scale == 1.0 + resized, scale = image.resize_to_fit(message, 4, 4) + assert resized is not message + assert (resized.width, resized.height, scale) == (4, 2, 0.5) + assert resized.encoding == message.encoding + assert resized.header == message.header + resized.data[:] = 0 + assert np.all(image.view(message) == 17) diff --git a/dimos/message_utils/time.py b/dimos/message_utils/time.py index 3d9197cb58a..81ef26b173e 100644 --- a/dimos/message_utils/time.py +++ b/dimos/message_utils/time.py @@ -16,7 +16,6 @@ import math import time -from typing import Any from dimos.messages.builtin_interfaces.msg import Duration, Time from dimos.messages.std_msgs.msg import Header @@ -69,9 +68,3 @@ def duration_from_seconds(value: float) -> Duration: def header_now(frame_id: str = "") -> Header: """Construct a wall-clock header; generated defaults themselves remain zero.""" return Header(stamp=time_from_nanoseconds(time.time_ns()), frame_id=frame_id) - - -def message_header(message: Any) -> Header | None: - """Return source metadata from a generated stamped value.""" - header = getattr(message, "header", None) - return header if isinstance(header, Header) else None diff --git a/dimos/visualization/rerun/cdr.py b/dimos/visualization/rerun/cdr.py index 920de372993..aaa2e623407 100644 --- a/dimos/visualization/rerun/cdr.py +++ b/dimos/visualization/rerun/cdr.py @@ -22,11 +22,7 @@ import numpy as np -from dimos.message_utils.image import ( - image_from_compressed, - image_to_rgb, - image_view, -) +from dimos.message_utils import image from dimos.message_utils.pointcloud import pointcloud_rgb, pointcloud_xyz from dimos.messages.dimos_msgs.msg import EntityMarkers from dimos.messages.geometry_msgs.msg import PointStamped, PoseStamped @@ -70,14 +66,14 @@ def image_archetype(message: Image | CompressedImage) -> rr.Image | rr.DepthImag if isinstance(message, CompressedImage): name = message.format.lower() if "jxl" in name or ("png" in name and name.startswith("16uc1;")): - return image_archetype(image_from_compressed(message)) + return image_archetype(image.from_compressed(message)) if "jpeg" in name or "jpg" in name: return rr.EncodedImage(contents=bytes(message.data), media_type="image/jpeg") if "png" in name: return rr.EncodedImage(contents=bytes(message.data), media_type="image/png") raise ValueError(f"unsupported compressed image format {message.format!r}") if message.encoding in ("16UC1", "32FC1", "mono16"): - pixels = image_view(message) + pixels = image.view(message) pixels = np.ascontiguousarray(pixels, dtype=pixels.dtype.newbyteorder("=")) if message.encoding != "mono16": return rr.DepthImage(pixels, meter=1000.0 if message.encoding == "16UC1" else 1.0) @@ -89,7 +85,7 @@ def image_archetype(message: Image | CompressedImage) -> rr.Image | rr.DepthImag color_model="L", ) return rr.Image( - bytes=image_to_rgb(message).tobytes(), + bytes=image.to_rgb(message).tobytes(), width=message.width, height=message.height, datatype=np.uint8, From 95139fae4846faae61098dcde37ce6d3249bfb74 Mon Sep 17 00:00:00 2001 From: cc Date: Mon, 26 Oct 2026 11:43:57 -0700 Subject: [PATCH 08/10] refactor(messages): include occupancy and coordinate updates in utility layer --- .github/workflows/cdr-native.yml | 3 +- dimos/message_utils/occupancy.py | 172 +++++++++++++++++++++ dimos/message_utils/pointcloud.py | 22 ++- dimos/message_utils/test_occupancy.py | 202 +++++++++++++++++++++++++ dimos/message_utils/test_pointcloud.py | 26 ++++ scripts/check_message_utils_api.py | 2 +- 6 files changed, 421 insertions(+), 6 deletions(-) create mode 100644 dimos/message_utils/occupancy.py create mode 100644 dimos/message_utils/test_occupancy.py diff --git a/.github/workflows/cdr-native.yml b/.github/workflows/cdr-native.yml index 220d091efe3..dbe5af20194 100644 --- a/.github/workflows/cdr-native.yml +++ b/.github/workflows/cdr-native.yml @@ -41,7 +41,8 @@ jobs: dimos/message_utils/test_time.py \ dimos/message_utils/test_geometry.py \ dimos/message_utils/test_pointcloud.py \ - dimos/message_utils/test_pointcloud_colors.py + dimos/message_utils/test_pointcloud_colors.py \ + dimos/message_utils/test_occupancy.py - name: Test native SDK run: | cargo test --locked --manifest-path native/rust/dimos-module/Cargo.toml diff --git a/dimos/message_utils/occupancy.py b/dimos/message_utils/occupancy.py new file mode 100644 index 00000000000..eaf703bf241 --- /dev/null +++ b/dimos/message_utils/occupancy.py @@ -0,0 +1,172 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Occupancy operations on generated grid messages. + +Import with ``from dimos.message_utils import occupancy``. ``view`` borrows +read-only cells, and ``from_array`` copies cells and metadata into a message. +``block_max_reduce`` returns its input when the grid is smaller than a block. +""" + +from copy import deepcopy +from enum import IntEnum +import math +from pathlib import Path +from typing import Any + +import numpy as np +from numpy.typing import NDArray +from PIL import Image + +from dimos.message_utils import geometry +from dimos.messages.geometry_msgs.msg import Point, Pose, Quaternion +from dimos.messages.nav_msgs.msg import MapMetaData, OccupancyGrid +from dimos.messages.std_msgs.msg import Header + + +class CostValues(IntEnum): + """Standard cost values for occupancy grid cells. + + These values follow the ROS nav_msgs/OccupancyGrid convention: + - 0: Free space + - 1-99: Occupied space with varying cost levels + - 100: Lethal obstacle (definitely occupied) + - -1: Unknown space + """ + + UNKNOWN = -1 # Unknown space + FREE = 0 # Free space + OCCUPIED = 100 # Occupied/lethal space + + +def view(msg: OccupancyGrid) -> NDArray[np.int8]: + """Borrow read-only row-major cells; source buffer writes remain visible.""" + if len(msg.data) != msg.info.height * msg.info.width: + raise ValueError("occupancy data length does not match dimensions") + view = msg.data.view().reshape(msg.info.height, msg.info.width) + view.setflags(write=False) + return view + + +def block_max_reduce(cells: NDArray[np.int8], factor: int) -> NDArray[np.int8]: + """Coarsen an occupancy grid by taking the max of factor x factor blocks. + + Block maximum (not mean) so coarsening never erases an obstacle; a block + is unknown (-1) only when every cell in it is unknown. Trailing remainder + rows/cols are trimmed, keeping row 0/col 0 - the origin corner - so the + grid origin stays valid. Grids thinner than `factor` pass through. + """ + h, w = cells.shape[:2] + new_h, new_w = h // factor, w // factor + if new_h == 0 or new_w == 0: + return cells + trimmed = cells[: new_h * factor, : new_w * factor] + blocks = trimmed.reshape(new_h, factor, new_w, factor) + # Sink unknown below every known value for the max, then map it back. + as_int = blocks.astype(np.int16) + known = np.where(as_int < 0, -1000, as_int) + reduced = known.max(axis=(1, 3)) + reduced[reduced == -1000] = -1 + result: NDArray[np.int8] = reduced.astype(np.int8) + return result + + +def extent(message: OccupancyGrid) -> tuple[float, float]: + """Return physical width and height, rejecting an invalid grid resolution.""" + info = message.info + if not math.isfinite(info.resolution) or info.resolution <= 0: + raise ValueError("OccupancyGrid resolution must be finite and positive") + return info.width * info.resolution, info.height * info.resolution + + +def world_to_grid(message: OccupancyGrid, point: Point) -> tuple[float, float]: + """Convert a world point to continuous grid coordinates, including origin rotation.""" + extent(message) + matrix = geometry.pose_matrix(message.info.origin) + coordinates = np.array([point.x, point.y, point.z]) + if not np.isfinite(coordinates).all(): + raise ValueError("grid coordinates must be finite") + local = matrix[:3, :3].T @ (coordinates - matrix[:3, 3]) + if not np.isfinite(local).all(): + raise ValueError("grid coordinates must be finite") + return float(local[0] / message.info.resolution), float(local[1] / message.info.resolution) + + +def grid_to_world(message: OccupancyGrid, coordinates: tuple[float, float]) -> Point: + """Convert continuous grid coordinates to the grid plane in world coordinates.""" + extent(message) + x, y = coordinates + if not math.isfinite(x) or not math.isfinite(y): + raise ValueError("grid coordinates must be finite") + result = geometry.pose_matrix(message.info.origin) @ np.array( + [x * message.info.resolution, y * message.info.resolution, 0.0, 1.0] + ) + return Point(x=float(result[0]), y=float(result[1]), z=float(result[2])) + + +def from_file(path: Path, *, header: Header, resolution: float = 0.05) -> OccupancyGrid: + """Load a static .npy or grayscale PNG grid with explicit frame/time metadata. + + Cells must be ROS occupancy values (-1 unknown, 0..100 known); PNG 255 is + unknown. NumPy object/pickle loading is disabled. Row order is unchanged. + """ + if not math.isfinite(resolution) or resolution <= 0: + raise ValueError("OccupancyGrid resolution must be finite and positive") + cells: NDArray[Any] + if path.suffix == ".npy": + cells = np.load(path, allow_pickle=False) + elif path.suffix == ".png": + with Image.open(path) as image: + cells = np.asarray(image.convert("L")).astype(np.int16) + cells[cells == 255] = -1 + else: + raise ValueError(f"Unsupported occupancy file format: {path.suffix}") + return from_array(cells, header=header, resolution=resolution) + + +def from_array( + cells: NDArray[Any], *, header: Header, resolution: float, origin: Pose | None = None +) -> OccupancyGrid: + """Copy a 2D ROS occupancy array with explicit map metadata.""" + if not math.isfinite(resolution) or resolution <= 0: + raise ValueError("OccupancyGrid resolution must be finite and positive") + if cells.ndim != 2 or cells.dtype.kind not in "iu": + raise ValueError("Occupancy cells must be a 2D integer array") + if np.any((cells < -1) | (cells > 100)): + raise ValueError("Occupancy cells must be -1 or 0..100") + return OccupancyGrid( + header=deepcopy(header), + info=MapMetaData( + map_load_time=deepcopy(header.stamp), + resolution=resolution, + height=cells.shape[0], + width=cells.shape[1], + origin=deepcopy(origin) + if origin is not None + else Pose( + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + position=Point(x=0.0, y=0.0, z=0.0), + ), + ), + data=np.asarray(cells.astype(np.int8).reshape(-1), dtype=np.int8), + ) + + +def cell_value(message: OccupancyGrid, point: Point) -> int: + """Read a world cell, retaining the planner's unknown result outside the map.""" + x, y = world_to_grid(message, point) + column, row = int(x), int(y) + if not (0 <= column < message.info.width and 0 <= row < message.info.height): + return CostValues.UNKNOWN + return int(view(message)[row, column]) diff --git a/dimos/message_utils/pointcloud.py b/dimos/message_utils/pointcloud.py index f118918a152..2b7010bc97d 100644 --- a/dimos/message_utils/pointcloud.py +++ b/dimos/message_utils/pointcloud.py @@ -239,13 +239,27 @@ def transform(message: PointCloud2, transform: TransformStamped) -> PointCloud2: """ if message.header.frame_id != transform.child_frame_id: raise ValueError("point cloud frame does not match transform child frame") + matrix = geometry.transform_matrix(transform.transform) + coordinates = xyz(message) + transformed = coordinates @ matrix[:3, :3].T + matrix[:3, 3] + output = with_xyz(message, transformed) + output.header.frame_id = transform.header.frame_id + return output + + +def with_xyz(message: PointCloud2, coordinates: NDArray[Any]) -> PointCloud2: + """Replace XYZ on a copy, preserving all other fields, padding and exact Header. + + Accept one coordinate triple per point in row order. This supports per-point + motion correction without losing capture times, intensity or organized layout. + """ records = view(message) + transformed = np.asarray(coordinates) + if transformed.shape != (message.height * message.width, 3): + raise ValueError("replacement coordinates must have shape (height * width, 3)") for name in ("x", "y", "z"): if name not in (records.dtype.names or ()) or records.dtype[name].kind != "f": raise ValueError("transform requires scalar floating-point XYZ fields") - matrix = geometry.transform_matrix(transform.transform) - coordinates = xyz(message) - transformed = coordinates @ matrix[:3, :3].T + matrix[:3, 3] for index, name in enumerate(("x", "y", "z")): values = transformed[:, index] if np.any(np.abs(values[np.isfinite(values)]) > np.finfo(records.dtype[name]).max): @@ -261,7 +275,7 @@ def transform(message: PointCloud2, transform: TransformStamped) -> PointCloud2: writable[name] = transformed[:, index].reshape(message.height, message.width) output = replace(message, header=deepcopy(message.header), fields=deepcopy(message.fields)) output.data = np.frombuffer(payload, dtype=np.uint8).copy() - output.header = Header(stamp=deepcopy(message.header.stamp), frame_id=transform.header.frame_id) + output.is_dense = message.is_dense and bool(np.isfinite(transformed).all()) return output diff --git a/dimos/message_utils/test_occupancy.py b/dimos/message_utils/test_occupancy.py new file mode 100644 index 00000000000..50c94b4ab9b --- /dev/null +++ b/dimos/message_utils/test_occupancy.py @@ -0,0 +1,202 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math + +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +import numpy as np +from PIL import Image +import pytest + +from dimos.message_utils import geometry, occupancy +from dimos.messages.builtin_interfaces.msg import Time +from dimos.messages.geometry_msgs.msg import Point, Pose, Quaternion +from dimos.messages.nav_msgs.msg import MapMetaData, OccupancyGrid +from dimos.messages.std_msgs.msg import Header + + +def test_occupancy_cells_keep_ros_row_order_and_do_not_alias_message(): + message = OccupancyGrid( + info=MapMetaData( + width=3, + height=2, + resolution=0.1, + map_load_time=Time(sec=0, nanosec=0), + origin=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.array([-1, 0, 100, 25, 50, 75], dtype=np.int8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + decoded = cdr_decode(cdr_encode(message), OccupancyGrid) + values = occupancy.view(decoded) + np.testing.assert_array_equal(values, [[-1, 0, 100], [25, 50, 75]]) + with pytest.raises(ValueError, match="read-only"): + values[0, 0] = 0 + copied = values.copy() + copied[0, 0] = 0 + assert decoded.data[0] == -1 + + +def test_occupancy_rejects_dimensions_that_do_not_match_payload(): + message = OccupancyGrid( + info=MapMetaData( + width=2, + height=2, + resolution=1.0, + map_load_time=Time(sec=0, nanosec=0), + origin=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.array([0], dtype=np.int8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + with pytest.raises(ValueError, match="data length"): + occupancy.view(message) + + +@pytest.mark.parametrize("resolution", [0.0, -0.1, math.nan, math.inf]) +def test_occupancy_rejects_invalid_resolution_for_nonempty_grid(resolution): + message = OccupancyGrid( + info=MapMetaData( + width=1, + height=1, + resolution=resolution, + map_load_time=Time(sec=0, nanosec=0), + origin=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.array([0], dtype=np.int8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + with pytest.raises(ValueError, match="resolution"): + occupancy.extent(message) + + +def test_empty_default_occupancy_grid_has_no_cells(): + assert occupancy.view( + OccupancyGrid( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + info=MapMetaData( + map_load_time=Time(sec=0, nanosec=0), + resolution=0.0, + width=0, + height=0, + origin=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.array([], dtype=np.int8), + ) + ).shape == (0, 0) + + +@pytest.mark.parametrize("suffix", [".npy", ".png"]) +def test_static_occupancy_loader_preserves_cells_header_and_identity(tmp_path, suffix): + cells = np.array([[-1, 0, 100], [25, 50, 75]], dtype=np.int8) + path = tmp_path / ("grid" + suffix) + if suffix == ".npy": + np.save(path, cells) + else: + Image.fromarray(cells.astype(np.uint8)).save(path) + header = Header(frame_id="map", stamp=Time(sec=123, nanosec=456)) + output = cdr_decode( + cdr_encode(occupancy.from_file(path, header=header, resolution=0.2)), + OccupancyGrid, + ) + np.testing.assert_array_equal(occupancy.view(output), cells) + assert output.header == header + assert output.info.map_load_time == header.stamp + assert output.info.origin.orientation.w == 1 + assert output.info.resolution == float(np.float32(0.2)) + + +@pytest.mark.parametrize( + "cells", + [np.zeros((2, 2), dtype=float), np.zeros(3, dtype=int), np.array([[101]]), np.array([[-2]])], +) +def test_static_occupancy_loader_rejects_invalid_cells(tmp_path, cells): + path = tmp_path / "grid.npy" + np.save(path, cells) + with pytest.raises(ValueError): + occupancy.from_file(path, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")) + + +@pytest.mark.parametrize("ns", [0, -1, 1700000000123456789]) +def test_array_constructor_copies_data_and_metadata_without_rounding_source_time(ns): + seconds, nanoseconds = divmod(ns, 1_000_000_000) + cells = np.array([[-1, 0], [100, 50]], dtype=np.int8) + header = Header(frame_id="map", stamp=Time(sec=seconds, nanosec=nanoseconds)) + origin = Pose( + position=Point(x=10, y=20, z=3), + orientation=geometry.quaternion_from_euler(0, 0, math.pi / 2), + ) + + message = occupancy.from_array(cells, header=header, resolution=0.5, origin=origin) + header.stamp.sec = 77 + cells[0, 0] = 44 + origin.position.x = 99 + actual = cdr_decode(cdr_encode(message), OccupancyGrid) + + assert actual.header.stamp.sec == seconds + assert actual.header.stamp.nanosec == nanoseconds + assert actual.info.map_load_time == actual.header.stamp + assert actual.info.origin.position.x == 10 + np.testing.assert_array_equal(occupancy.view(actual), [[-1, 0], [100, 50]]) + + +def test_coordinate_helpers_account_for_rotated_origin_and_grid_plane(): + origin = Pose( + position=Point(x=10, y=20, z=3), + orientation=geometry.quaternion_from_euler(0, 0, math.pi / 2), + ) + message = occupancy.from_array( + np.zeros((4, 4), dtype=np.int8), + header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)), + resolution=0.5, + origin=origin, + ) + + point = occupancy.grid_to_world(message, (2, 3)) + + assert (point.x, point.y, point.z) == pytest.approx((8.5, 21, 3)) + assert occupancy.world_to_grid(message, point) == pytest.approx((2, 3)) + assert occupancy.extent(message) == (2, 2) + + +@pytest.mark.parametrize("value", [math.nan, math.inf, -math.inf]) +def test_coordinate_helpers_reject_nonfinite_coordinates(value): + message = occupancy.from_array( + np.zeros((1, 1), dtype=np.int8), + header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)), + resolution=1, + ) + with pytest.raises(ValueError, match="finite"): + occupancy.grid_to_world(message, (value, 0)) + with pytest.raises(ValueError, match="finite"): + occupancy.world_to_grid(message, Point(x=value, y=0, z=0)) + + +def test_numpy_loader_never_unpickles_objects(tmp_path): + path = tmp_path / "objects.npy" + np.save(path, np.array([[{}]], dtype=object), allow_pickle=True) + with pytest.raises(ValueError, match="allow_pickle=False"): + occupancy.from_file(path, header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0))) diff --git a/dimos/message_utils/test_pointcloud.py b/dimos/message_utils/test_pointcloud.py index 5e0e472ea1a..85767859789 100644 --- a/dimos/message_utils/test_pointcloud.py +++ b/dimos/message_utils/test_pointcloud.py @@ -448,6 +448,32 @@ def test_per_point_stamps_reject_wrong_count_or_nonfinite_values(stamps): ) +def test_per_point_coordinate_replacement_keeps_organized_records(padded_cloud): + padded_cloud.header.stamp = Time(sec=1700000000, nanosec=123456789) + before = cdr_encode(padded_cloud) + result = pointcloud.with_xyz( + padded_cloud, pointcloud.xyz(padded_cloud) + np.array([0.5, -1, 2]) + ) + np.testing.assert_array_equal( + pointcloud.view(result)["tags"], pointcloud.view(padded_cloud)["tags"] + ) + assert result.header == padded_cloud.header + assert (result.height, result.width, result.row_step) == (2, 2, 64) + assert result.fields == padded_cloud.fields and result.is_bigendian == padded_cloud.is_bigendian + np.testing.assert_allclose( + pointcloud.xyz(result), [[1.5, 1, 5], [2.5, 3, 8], [3.5, 5, 11], [4.5, 7, 14]] + ) + assert bytes(result.data)[48:64] == bytes(padded_cloud.data)[48:64] + assert bytes(result.data)[112:128] == bytes(padded_cloud.data)[112:128] + assert cdr_encode(padded_cloud) == before + + +@pytest.mark.parametrize("coordinates", [np.zeros((3, 3)), np.zeros((2, 2, 3))]) +def test_per_point_coordinate_replacement_rejects_wrong_count(padded_cloud, coordinates): + with pytest.raises(ValueError, match="shape"): + pointcloud.with_xyz(padded_cloud, coordinates) + + def test_xyz_constructor_owns_float32_input_and_stamp_arrays() -> None: points = np.array([[1, 2, 3]], dtype=np.float32) header = Header(stamp=Time(sec=0, nanosec=0), frame_id="camera") diff --git a/scripts/check_message_utils_api.py b/scripts/check_message_utils_api.py index a48f047c173..798996736ed 100644 --- a/scripts/check_message_utils_api.py +++ b/scripts/check_message_utils_api.py @@ -21,7 +21,7 @@ import ast from pathlib import Path -DOMAINS = {"image", "camera_info", "geometry", "pointcloud", "time"} +DOMAINS = {"image", "camera_info", "geometry", "pointcloud", "time", "occupancy"} ROOT = Path(__file__).resolve().parents[1] From 3b97c5a763acfb142121402f0e409565756595c4 Mon Sep 17 00:00:00 2001 From: cc Date: Mon, 26 Oct 2026 11:48:59 -0700 Subject: [PATCH 09/10] test(messages): check native Python callers for stale helper APIs --- scripts/check_message_utils_api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/check_message_utils_api.py b/scripts/check_message_utils_api.py index 798996736ed..3c79091f9f3 100644 --- a/scripts/check_message_utils_api.py +++ b/scripts/check_message_utils_api.py @@ -40,7 +40,7 @@ def check() -> list[str]: if isinstance(target, ast.Name) ) errors = [] - for folder in ("dimos", "examples", "scripts"): + for folder in ("dimos", "examples", "scripts", "native"): for path in sorted((ROOT / folder).rglob("*.py")): source = path.read_text() if "dimos.message_utils" not in source or path == Path(__file__).resolve(): From c3c2a40bf28f57e4674cbdcf66598def229068cc Mon Sep 17 00:00:00 2001 From: cc Date: Mon, 26 Oct 2026 12:30:39 -0700 Subject: [PATCH 10/10] chore: remove message utility API checker --- scripts/check_message_utils_api.py | 82 ------------------------------ 1 file changed, 82 deletions(-) delete mode 100644 scripts/check_message_utils_api.py diff --git a/scripts/check_message_utils_api.py b/scripts/check_message_utils_api.py deleted file mode 100644 index 3c79091f9f3..00000000000 --- a/scripts/check_message_utils_api.py +++ /dev/null @@ -1,82 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Check message helper callers use domain namespaces and current function names. - -Run ``python scripts/check_message_utils_api.py`` from any checkout in the CDR -stack after migrating callers. This check does not import optional runtimes. -""" - -import ast -from pathlib import Path - -DOMAINS = {"image", "camera_info", "geometry", "pointcloud", "time", "occupancy"} -ROOT = Path(__file__).resolve().parents[1] - - -def check() -> list[str]: - exports: dict[str, set[str]] = {} - for domain in DOMAINS: - module = ast.parse((ROOT / "dimos/message_utils" / f"{domain}.py").read_text()) - exports[domain] = { - node.name for node in module.body if isinstance(node, (ast.FunctionDef, ast.ClassDef)) - } - exports[domain].update( - target.id - for node in module.body - if isinstance(node, ast.Assign) - for target in node.targets - if isinstance(target, ast.Name) - ) - errors = [] - for folder in ("dimos", "examples", "scripts", "native"): - for path in sorted((ROOT / folder).rglob("*.py")): - source = path.read_text() - if "dimos.message_utils" not in source or path == Path(__file__).resolve(): - continue - tree = ast.parse(source) - aliases: dict[str, str] = {} - for node in ast.walk(tree): - if not isinstance(node, ast.ImportFrom) or node.module is None: - continue - if node.module == "dimos.message_utils": - aliases.update( - (name.asname or name.name, name.name) - for name in node.names - if name.name in DOMAINS - ) - elif node.module.startswith("dimos.message_utils."): - domain = node.module.rsplit(".", 1)[-1] - for name in node.names: - if domain in DOMAINS and not name.name[0].isupper(): - errors.append( - f"{path.relative_to(ROOT)}:{node.lineno}: " - f"import the {domain} module instead of {name.name}" - ) - for node in ast.walk(tree): - if isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name): - target_domain = aliases.get(node.value.id) - if target_domain is not None and node.attr not in exports[target_domain]: - errors.append( - f"{path.relative_to(ROOT)}:{node.lineno}: " - f"unknown helper {target_domain}.{node.attr}" - ) - return errors - - -if __name__ == "__main__": - failures = check() - for failure in failures: - print(failure) - raise SystemExit(bool(failures))