diff --git a/dimos/message_utils/camera_info.py b/dimos/message_utils/camera_info.py new file mode 100644 index 00000000000..eb82210ba7d --- /dev/null +++ b/dimos/message_utils/camera_info.py @@ -0,0 +1,181 @@ +# 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 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 +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 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 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 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 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 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 = 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/geometry.py b/dimos/message_utils/geometry.py new file mode 100644 index 00000000000..c8d1878b817 --- /dev/null +++ b/dimos/message_utils/geometry.py @@ -0,0 +1,382 @@ +# 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 wire values. + +Import with ``from dimos.message_utils import geometry``. Result names state +which value is constructed, such as ``geometry.quaternion_from_euler`` or +``geometry.transform_from_pose``. Array conversions return independent writable +storage; operations return new values without mutating their inputs. +""" + +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/image.py b/dimos/message_utils/image.py index 08cfb6296bf..c26a8923402 100644 --- a/dimos/message_utils/image.py +++ b/dimos/message_utils/image.py @@ -12,13 +12,25 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Read and convert generated image buffers.""" +"""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 + +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,10 +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] @@ -65,8 +75,13 @@ def _image_from_array( ) -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] @@ -85,9 +100,88 @@ def image_view(msg: Image) -> NDArray[Any]: return view -def image_from_compressed(message: CompressedImage) -> Image: +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 = 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 to_compressed( + message: Image, + quality: int = 75, + *, + format: str = "jpeg", + max_width: int | None = None, + effort: int | None = None, +) -> CompressedImage: + """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. + """ + if not isinstance(message, Image): + raise TypeError("compression expects a generated Image") + if max_width is not None: + message, _ = resize_to_fit(message, max_width, max_width) + if format == "jpeg": + output_encoding = "mono8" if message.encoding == "mono8" else "bgr8" + data = to_jpeg(message, quality=quality) + elif format == "png": + 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"): + 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 = 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 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 +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) @@ -122,18 +216,74 @@ 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 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 = 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 sharpness(message: Image) -> float: + """Laplacian variance of an 8-bit visual image, downsampled to 160 pixels wide.""" + pixels = 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 to_bgr(message: Image) -> NDArray[np.uint8]: + """Return an independent BGR8 array for drawing or OpenCV color operations.""" -def image_to_rgb(message: Image) -> NDArray[np.uint8]: + pixels = 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 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. """ - import cv2 - pixels = image_view(message) + pixels = view(message) if message.encoding == "rgb8": return pixels.copy() if message.encoding in ("mono16", "16UC1"): @@ -149,3 +299,30 @@ 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 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") + 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( + view(message), + (max(1, int(message.width * scale)), max(1, int(message.height * scale))), + interpolation=cv2.INTER_LINEAR, + ) + return from_array(pixels, encoding=message.encoding, header=message.header), scale + + +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 from_array(np.asarray(image.convert("RGB")), encoding="rgb8", header=header) 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 1e630828eab..2b7010bc97d 100644 --- a/dimos/message_utils/pointcloud.py +++ b/dimos/message_utils/pointcloud.py @@ -12,23 +12,37 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Read generated point-cloud layouts and extract positions and colors.""" +"""Point cloud operations for generated PointCloud2 wire messages. +Import with ``from dimos.message_utils import pointcloud``. ``view`` borrows +read-only structured records; ``xyz``, ``rgb`` and ``stamps`` return independent +writable arrays. Constructors and record operations return messages with copied +storage. ``voxel_downsample`` may return its input unchanged. +""" + +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 import camera_info, geometry, image, time as message_time +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"} -def pointcloud_view(message: PointCloud2) -> NDArray[Any]: +def view(message: PointCloud2) -> NDArray[Any]: """Borrow read-only structured (height, width) points, including row/field padding. Fields retain their declared counts and endianness. The returned array keeps - the message's storage alive. Use ``view.copy()`` for independent mutable data. + the message's storage alive. Writes through the source buffer remain visible. + Use ``view(message).copy()`` for independent mutable data. """ if message.width and not message.point_step: raise ValueError("nonempty point cloud requires a positive point_step") @@ -63,22 +77,112 @@ def pointcloud_view(message: PointCloud2) -> NDArray[Any]: return view -def pointcloud_xyz(message: PointCloud2) -> NDArray[np.float64]: +def xyz(message: PointCloud2) -> NDArray[np.float64]: """Copy scalar x/y/z fields into an N-by-3 float64 array, preserving point order.""" - points = pointcloud_view(message) + points = view(message) for name in ("x", "y", "z"): if name not in (points.dtype.names or ()) or points[name].ndim != 2: raise ValueError(f"point cloud requires a scalar {name!r} field") return np.stack([points[name].ravel() for name in ("x", "y", "z")], axis=-1).astype(np.float64) -def pointcloud_rgb(message: PointCloud2) -> NDArray[np.uint8] | None: +def 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.array(values, dtype=" NDArray[np.float64] | None: + """Copy per-point capture times, or return None when the field is absent.""" + 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.array(view(message)["stamps"].reshape(-1), dtype=np.float64, copy=True) + + +def select(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. + """ + 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 rgb(message: PointCloud2) -> NDArray[np.uint8] | None: """Copy RGB bytes from the standard packed rgb/rgba field, if present. FLOAT32 fields carry color bits, not numeric floating-point color values. Both representations preserve the declared byte order and row padding. """ - points = pointcloud_view(message) + points = view(message) names = points.dtype.names or () name = next((name for name in ("rgb", "rgba") if name in names), None) if name is None: @@ -88,3 +192,212 @@ 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(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. + """ + view(first) + 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(first, np.ones(first.width * first.height, dtype=np.bool_)) + b = select(second, np.ones(second.width * second.height, dtype=np.bool_)) + stamp = max((first.header.stamp, second.header.stamp), key=message_time.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(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") + 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") + 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): + raise ValueError(f"transformed {name} coordinates exceed field range") + payload = bytearray(bytes(message.data)) + writable = np.ndarray( + (message.height, message.width), + dtype=records.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.is_dense = message.is_dense and bool(np.isfinite(transformed).all()) + return output + + +def 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 = 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 = camera_info.intrinsic_matrix(calibration) + 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 from_xyz_rgb(np.asarray(cloud.points), colors, header=depth.header) + + +def voxel_downsample(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 return the original + message without copying. 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 = xyz(message).astype(np.float32) + if not np.isfinite(points).all(): + raise ValueError("voxel downsampling requires finite XYZ positions") + colors = 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) + reduced_points = reduced.point["positions"].numpy() + if colors is None: + return from_xyz(reduced_points, header=message.header) + reduced_colors = np.clip(reduced.point["colors"].numpy() * 255.0, 0, 255).astype(np.uint8) + return from_xyz_rgb(reduced_points, reduced_colors, header=message.header) + + +def 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(xyz(message)) + colors = rgb(message) + if colors is not None: + result.colors = o3d.utility.Vector3dVector(colors.astype(np.float64) / 255.0) + return result + + +def bounds_intersect(first: PointCloud2, second: PointCloud2) -> bool: + """Test closed axis-aligned bounds; empty or entirely nonfinite clouds do not intersect.""" + a, b = xyz(first), 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_camera_info.py b/dimos/message_utils/test_camera_info.py new file mode 100644 index 00000000000..65dbfb0595d --- /dev/null +++ b/dimos/message_utils/test_camera_info.py @@ -0,0 +1,143 @@ +# 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 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 + +_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 = camera_info.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_geometry.py b/dimos/message_utils/test_geometry.py new file mode 100644 index 00000000000..189378c15a2 --- /dev/null +++ b/dimos/message_utils/test_geometry.py @@ -0,0 +1,272 @@ +# 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 import geometry +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 = geometry.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), geometry.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 = geometry.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 geometry.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), geometry.quaternion_from_euler(0.3, -0.2, 1.1)) + inverse = geometry.inverse_transform(original) + result = geometry.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"): + geometry.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): + geometry.inverse_transform(edge) + with pytest.raises(ValueError): + geometry.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"): + geometry.inverse_transform(edge) + with pytest.raises(ValueError, match="Euler angles"): + geometry.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 = geometry.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 = geometry.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=geometry.quaternion_from_euler(0.0, 0.0, math.pi / 2), + ) + matrix = geometry.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=geometry.quaternion_from_euler(0, 0, math.pi / 2), + ) + matrix = geometry.pose_matrix(pose) + np.testing.assert_allclose(matrix @ [1, 0, 0, 1], [2, 0, 3, 1], atol=1e-12) + decoded = cdr_decode(cdr_encode(geometry.pose_from_matrix(matrix)), Pose) + np.testing.assert_allclose(geometry.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): + geometry.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): + geometry.pose_matrix(pose) + + +def test_local_offset_uses_pose_axes_and_copies_orientation(): + pose = Pose( + position=Point(x=1, y=2, z=3), + orientation=geometry.quaternion_from_euler(0, 0, math.pi / 2), + ) + translated = geometry.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 geometry.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"): + geometry.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), + ) diff --git a/dimos/message_utils/test_image_compression.py b/dimos/message_utils/test_image_compression.py new file mode 100644 index 00000000000..c84407e86bf --- /dev/null +++ b/dimos/message_utils/test_image_compression.py @@ -0,0 +1,202 @@ +# 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 import image, time as message_time +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=message_time.time_from_seconds(1234.5678)), + ) + + +def test_jpeg_roundtrip_preserves_meta_and_pixels(rgb_image) -> None: + 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 message_time.to_nanoseconds(ci.header.stamp) == message_time.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 message_time.to_nanoseconds(img.header.stamp) == message_time.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 = 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 message_time.to_nanoseconds(img.header.stamp) == message_time.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=message_time.time_from_seconds(1.0)), + ) + img = image.from_compressed(image.to_compressed(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 = image.to_compressed(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( + message_time.to_nanoseconds(out.header.stamp) + - message_time.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"): + image.to_compressed(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"): + image.to_compressed(depth, format="png") + + +def test_from_image_rejects_non_image() -> None: + with pytest.raises(TypeError): + image.to_compressed(b"not an image") + + +def test_max_width_resizes(rgb_image) -> None: + 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( + data, + encoding="mono16", + header=Header(frame_id="d", stamp=message_time.time_from_seconds(1.0)), + ) + ci = image.to_compressed(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=message_time.time_from_seconds(2.0)), + ) + 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 message_time.to_nanoseconds(img.header.stamp) == message_time.to_nanoseconds( + src.header.stamp + ) + assert img.header.frame_id == "d" + + +def test_jxl_lossy_rgb_honors_quality(rgb_image) -> None: + 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() + assert diff < 5, f"JXL q90 mean pixel error too high: {diff}" + + +def test_jxl_effort_trades_cpu_for_size(rgb_image) -> None: + 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") + with pytest.raises(ValueError, match="JXL cannot encode"): + 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(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 new file mode 100644 index 00000000000..2dbb2d607ae --- /dev/null +++ b/dimos/message_utils/test_image_views.py @@ -0,0 +1,251 @@ +# 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 import image, time as message_time +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) + 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 + 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) + 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 + + +@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=message_time.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) + + +@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/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 new file mode 100644 index 00000000000..85767859789 --- /dev/null +++ b/dimos/message_utils/test_pointcloud.py @@ -0,0 +1,488 @@ +# 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 import camera_info, image, pointcloud +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"): + pointcloud.transform(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 pointcloud.voxel_downsample(message, 0.5) is message + assert pointcloud.voxel_downsample(message, 0.0) is message + with pytest.raises(ValueError, match="finite"): + pointcloud.voxel_downsample(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 = pointcloud.voxel_downsample(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 = pointcloud.select(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, + ) + + +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") + message = pointcloud.from_xyz(points, header=header) + points[:] = 0 + np.testing.assert_array_equal(pointcloud.xyz(message), [[1, 2, 3]]) + stamped = pointcloud.from_xyz(points, header=header, stamps=np.array([1.25])) + times = pointcloud.stamps(stamped) + assert times is not None + assert times.flags.writeable + times[:] = 0 + np.testing.assert_array_equal(pointcloud.stamps(stamped), [1.25]) diff --git a/dimos/message_utils/test_pointcloud_colors.py b/dimos/message_utils/test_pointcloud_colors.py new file mode 100644 index 00000000000..07e53e72ee9 --- /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 import pointcloud +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) diff --git a/dimos/message_utils/test_time.py b/dimos/message_utils/test_time.py new file mode 100644 index 00000000000..f739bec3e43 --- /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 import time as message_time + + +@pytest.mark.parametrize( + "factory", [message_time.time_from_nanoseconds, message_time.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 message_time.to_nanoseconds(stamp) == value + assert message_time.to_nanoseconds(cdr_decode(cdr_encode(stamp), stamp.__msgtype__)) == value + + +@pytest.mark.parametrize( + "factory", [message_time.time_from_seconds, message_time.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 message_time.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"): + message_time.time_from_seconds(value) + + +def test_invalid_ros_fraction_is_rejected(): + stamp = message_time.time_from_nanoseconds(0) + stamp.nanosec = 1_000_000_000 + + with pytest.raises(ValueError, match="nanosec"): + message_time.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 = message_time.header_now("camera") + + assert message_time.to_nanoseconds(header.stamp) == 1_700_000_000_123_456_789 + assert header.frame_id == "camera" + + +@pytest.mark.parametrize( + "factory", [message_time.time_from_nanoseconds, message_time.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 message_time.to_nanoseconds(factory(value)) == value + + +@pytest.mark.parametrize( + "factory", [message_time.time_from_nanoseconds, message_time.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..81ef26b173e --- /dev/null +++ b/dimos/message_utils/time.py @@ -0,0 +1,70 @@ +# 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 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) diff --git a/dimos/visualization/rerun/cdr.py b/dimos/visualization/rerun/cdr.py index 920de372993..43a2597e4d9 100644 --- a/dimos/visualization/rerun/cdr.py +++ b/dimos/visualization/rerun/cdr.py @@ -22,12 +22,7 @@ import numpy as np -from dimos.message_utils.image import ( - image_from_compressed, - image_to_rgb, - image_view, -) -from dimos.message_utils.pointcloud import pointcloud_rgb, pointcloud_xyz +from dimos.message_utils import image, pointcloud from dimos.messages.dimos_msgs.msg import EntityMarkers from dimos.messages.geometry_msgs.msg import PointStamped, PoseStamped from dimos.messages.nav_msgs.msg import Odometry, Path @@ -70,14 +65,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 +84,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, @@ -114,7 +109,7 @@ def cloud_archetype( from matplotlib import colormaps import rerun as rr - points, colors = pointcloud_xyz(message), pointcloud_rgb(message) + points, colors = pointcloud.xyz(message), pointcloud.rgb(message) keep = np.isfinite(points).all(axis=1) if bottom_cutoff is not None: keep &= points[:, 2] >= bottom_cutoff