diff --git a/.github/workflows/message-codegen.yml b/.github/workflows/message-codegen.yml index 0890e1e007a..39f8f31277f 100644 --- a/.github/workflows/message-codegen.yml +++ b/.github/workflows/message-codegen.yml @@ -27,7 +27,8 @@ on: - dimos/perception/localize/** - dimos/perception/detection/project.py - dimos/models/embedding/** - - dimos/perception/experimental/moduleDB.py + - dimos/perception/experimental/** + - dimos/models/segmentation/** - dimos/models/vl/base.py - dimos/models/vl/qwen.py - dimos/models/vl/test_base.py @@ -190,3 +191,9 @@ jobs: run: | .native-review/bin/python -I -c 'import dimos; import pytest; import sys; raise SystemExit(pytest.main(sys.argv[1:]))' --noconftest --import-mode=importlib -o addopts='' -q "$PWD/dimos/perception/localize/test_rig.py" "$PWD/dimos/perception/detection/type/detection3d/test_generated_projection.py" "$PWD/dimos/memory/store/test_cdr_recording.py" "$PWD/dimos/memory/codecs/test_codecs.py" MYPYPATH="$PWD/build/message-source/python" .native-review/bin/python -m mypy --follow-imports=silent dimos/perception/localize/association.py dimos/perception/localize/dandetect.py dimos/perception/localize/module.py dimos/perception/localize/rig.py dimos/perception/localize/support_plane.py dimos/perception/localize/types.py dimos/perception/localize/localize.py dimos/perception/localize/tools/tool_localize.py dimos/perception/localize/demo_blueprints/xarm_feed.py dimos/perception/detection/project.py dimos/memory/codecs/jpeg.py dimos/memory/codecs/base.py + - name: Verify generated experimental objects, tracking and segmentation boundaries + if: runner.os == 'Linux' + run: | + .native-review/bin/python -m pip install 'hydra-core==1.3.2' + .native-review/bin/python -I -c 'import dimos; import pytest; import sys; raise SystemExit(pytest.main(sys.argv[1:]))' --noconftest --import-mode=importlib -o addopts='' -o 'markers=self_hosted: optional heavyweight integration' --strict-markers -m 'not self_hosted' -q "$PWD/dimos/perception/experimental/test_object_generated.py" "$PWD/dimos/perception/experimental/test_object_scene_registration_temporal.py" "$PWD/dimos/perception/experimental/test_generated_tracker.py" "$PWD/dimos/perception/experimental/test_spatial_memory.py" "$PWD/dimos/models/segmentation/test_generated_edge_tam_images.py" "$PWD/dimos/models/segmentation/test_yoloe.py" + MYPYPATH="$PWD/build/message-source/python" .native-review/bin/python -m mypy --follow-imports=silent dimos/perception/experimental/object.py dimos/perception/experimental/objectDB.py dimos/perception/experimental/object_scene_registration.py dimos/perception/experimental/object_scene_registration_spec.py dimos/perception/experimental/object_tracker.py dimos/perception/experimental/object_tracker_2d.py dimos/perception/experimental/object_tracker_3d.py dimos/perception/experimental/spatial_perception.py dimos/models/segmentation/edge_tam.py dimos/models/segmentation/yoloe.py diff --git a/dimos/models/segmentation/edge_tam.py b/dimos/models/segmentation/edge_tam.py index 667b05b4047..b65977283c0 100644 --- a/dimos/models/segmentation/edge_tam.py +++ b/dimos/models/segmentation/edge_tam.py @@ -27,8 +27,9 @@ from PIL import Image as PILImage import torch +from dimos.messages.sensor_msgs.msg import Image from dimos.models.base import default_torch_device -from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.image import image_to_bgr, image_to_jpeg from dimos.perception.detection.detectors.base import Detector from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D @@ -172,7 +173,7 @@ def propose_all( ) self._mask_generators[points_per_side] = generator - rgb = cv2.cvtColor(image.to_opencv(), cv2.COLOR_BGR2RGB) + rgb = cv2.cvtColor(image_to_bgr(image), cv2.COLOR_BGR2RGB) amp = ( torch.autocast("cuda", dtype=torch.bfloat16) if self._predictor.device.type == "cuda" @@ -203,7 +204,7 @@ def segment( return ImageDetections2D(detections.image, []) image = detections.image - rgb = cv2.cvtColor(image.to_opencv(), cv2.COLOR_BGR2RGB) + rgb = cv2.cvtColor(image_to_bgr(image), cv2.COLOR_BGR2RGB) # bfloat16 autocast is CUDA-only; MPS runs in float32 amp = ( @@ -254,7 +255,7 @@ def _prepare_frame(self, image: Image) -> torch.Tensor: """Prepare frame for SAM2 (resize, normalize, convert to tensor).""" import cv2 - cv_image = image.to_opencv() + cv_image = image_to_bgr(image) rgb_image = cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB) pil_image = PILImage.fromarray(rgb_image) @@ -407,7 +408,7 @@ def stop(self) -> None: def _temp_dir_context(image: Image) -> Generator[str, None, None]: path = tempfile.mkdtemp() - image.save(f"{path}/00000.jpg") + Path(path, "00000.jpg").write_bytes(image_to_jpeg(image)) try: yield path diff --git a/dimos/models/segmentation/test_generated_edge_tam_images.py b/dimos/models/segmentation/test_generated_edge_tam_images.py new file mode 100644 index 00000000000..ac079411bb5 --- /dev/null +++ b/dimos/models/segmentation/test_generated_edge_tam_images.py @@ -0,0 +1,50 @@ +# 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 pathlib import Path +from types import SimpleNamespace + +import cv2 +import numpy as np +import torch + +from dimos.models.segmentation import edge_tam +from dimos.msgs.image import image_from_array +from dimos.perception.detection.type.detection2d.seg import Detection2DSeg + + +def test_generated_tracker_frame_and_initial_jpeg(monkeypatch): + monkeypatch.setattr( + edge_tam, "_build_model", lambda: SimpleNamespace(image_size=8, device="cpu") + ) + processor = edge_tam.EdgeTAMProcessor() + pixels = np.zeros((12, 16, 3), dtype=np.uint8) + pixels[:, :, 0] = 255 + image = image_from_array(pixels, encoding="rgb8") + frame = processor._prepare_frame(image) + assert frame.shape == (3, 8, 8) + torch.testing.assert_close( + frame[:, 0, 0], torch.tensor([(1 - 0.485) / 0.229, -0.456 / 0.224, -0.406 / 0.225]) + ) + with edge_tam._temp_dir_context(image) as directory: + path = Path(directory) + decoded = cv2.imread(str(path / "00000.jpg")) + assert decoded.shape == (12, 16, 3) + assert decoded[0, 0, 2] > 240 + assert not path.exists() + mask = torch.zeros((12, 16)) + mask[2:8, 3:10] = 1 + detection = Detection2DSeg.from_sam2_result(mask, 7, image) + assert detection.bbox == (3, 2, 9, 7) + assert detection.to_ros_detection2d().header == image.header diff --git a/dimos/models/segmentation/test_yoloe.py b/dimos/models/segmentation/test_yoloe.py index 5848f1aeca5..5d2efaffa50 100644 --- a/dimos/models/segmentation/test_yoloe.py +++ b/dimos/models/segmentation/test_yoloe.py @@ -19,19 +19,21 @@ import numpy as np from ultralytics.models.yolo.yoloe import YOLOEVPSegPredictor +from dimos.messages.sensor_msgs.msg import Image +from dimos.messages.std_msgs.msg import Header from dimos.models.segmentation.yoloe import YoloeBoxSegmenter -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.image import image_from_array, image_to_bgr +from dimos.msgs.time import time_from_seconds, to_seconds from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.perception.detection.type.detection2d.seg import Detection2DSeg def _image() -> Image: - return Image( - data=np.zeros((32, 48, 3), dtype=np.uint8), - format=ImageFormat.BGR, - frame_id="camera", - ts=12.5, + return image_from_array( + np.zeros((32, 48, 3), dtype=np.uint8), + encoding="bgr8", + header=Header(frame_id="camera", stamp=time_from_seconds(12.5)), ) @@ -42,7 +44,7 @@ def _box(image: Image, *, bbox: tuple[float, float, float, float]) -> Detection2 class_id=7, confidence=0.73, name="white mug", - ts=image.ts, + ts=to_seconds(image.header.stamp), image=image, ) @@ -60,7 +62,7 @@ def _mask( class_id=prompt_id, confidence=0.11, name=f"prompt-{prompt_id}", - ts=image.ts, + ts=to_seconds(image.header.stamp), image=image, mask=np.full((32, 48), value, dtype=np.uint8), ) @@ -137,7 +139,7 @@ def test_visual_box_prompts_use_segmentation_predictor(mocker: Any) -> None: get_data.assert_called_once_with("models_yoloe") model_factory.assert_called_once_with(Path("/models/yoloe-11s-seg.pt")) kwargs = dict(model.predict.call_args.kwargs) - np.testing.assert_array_equal(kwargs.pop("source"), image.to_opencv()) + np.testing.assert_array_equal(kwargs.pop("source"), image_to_bgr(image)) visual_prompts = kwargs.pop("visual_prompts") np.testing.assert_array_equal( visual_prompts["bboxes"], np.asarray([detection.bbox], dtype=np.float64) diff --git a/dimos/models/segmentation/yoloe.py b/dimos/models/segmentation/yoloe.py index a13d0a9d009..fd075227d27 100644 --- a/dimos/models/segmentation/yoloe.py +++ b/dimos/models/segmentation/yoloe.py @@ -23,7 +23,8 @@ from numpy.typing import NDArray from ultralytics.models.yolo.yoloe import YOLOEVPSegPredictor -from dimos.msgs.sensor_msgs.Image import Image +from dimos.messages.sensor_msgs.msg import Image +from dimos.msgs.image import image_to_bgr from dimos.perception.detection.detectors.yoloe import Yoloe2DDetector, YoloePromptMode from dimos.perception.detection.type.detection2d.bbox import Bbox, Detection2DBBox from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D @@ -58,7 +59,7 @@ def predict_image(self, image: Image) -> ImageDetections2D: raise RuntimeError("YOLO-E segmentation requires visual box prompts") with self._lock: results = self.model.predict( - source=image.to_opencv(), + source=image_to_bgr(image), device=self.device, conf=self.conf, iou=0.6, diff --git a/dimos/perception/experimental/object.py b/dimos/perception/experimental/object.py index 76f4cc425c1..6d06f01ca0a 100644 --- a/dimos/perception/experimental/object.py +++ b/dimos/perception/experimental/object.py @@ -19,25 +19,43 @@ from typing import TYPE_CHECKING, Any import uuid -from dimos_lcm.geometry_msgs import Pose -from dimos_lcm.vision_msgs import BoundingBox3D, ObjectHypothesis, ObjectHypothesisWithPose import numpy as np -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.std_msgs.Header import Header -from dimos.msgs.vision_msgs.Detection3D import Detection3D as ROSDetection3D -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray +from dimos.messages.builtin_interfaces.msg import Time +from dimos.messages.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + PoseWithCovariance, + Quaternion, + TransformStamped, + Vector3, +) +from dimos.messages.sensor_msgs.msg import CameraInfo, Image, PointCloud2 +from dimos.messages.std_msgs.msg import Header +from dimos.messages.vision_msgs.msg import ( + BoundingBox3D, + Detection3D as ROSDetection3D, + Detection3DArray, + ObjectHypothesis, + ObjectHypothesisWithPose, +) +from dimos.msgs.geometry import quaternion_from_matrix, transform_matrix +from dimos.msgs.image import image_to_rgb, image_view +from dimos.msgs.pointcloud import ( + concatenate_clouds, + pointcloud_from_xyz, + pointcloud_from_xyz_rgb, + pointcloud_rgb, + pointcloud_to_open3d, + pointcloud_xyz, + voxel_downsample_cloud, +) +from dimos.msgs.time import time_from_seconds from dimos.perception.detection.type.detection2d.seg import Detection2DSeg from dimos.perception.detection.type.detection3d.base import Detection3D if TYPE_CHECKING: - from dimos_lcm.sensor_msgs import CameraInfo - from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D @@ -54,7 +72,7 @@ class Object(Detection3D): size: Vector3 pose: PoseStamped pointcloud: PointCloud2 - camera_transform: Transform | None = None + camera_transform: TransformStamped | None = None mask: np.ndarray[Any, np.dtype[np.uint8]] | None = None detections_count: int = 1 visual_embedding: np.ndarray[Any, np.dtype[np.float32]] | None = None @@ -68,8 +86,8 @@ def __post_init__(self) -> None: def set_center(self, center: Vector3) -> None: """Update the canonical center and its pose representation together.""" - self.center = Vector3(center) - self.pose.position = self.center + self.center = Vector3(x=center.x, y=center.y, z=center.z) + self.pose.pose.position = Point(x=center.x, y=center.y, z=center.z) def update_object(self, other: Object, *, accumulate_pointcloud: bool = True) -> None: """Update this object with data from another detection. @@ -89,7 +107,7 @@ def update_object(self, other: Object, *, accumulate_pointcloud: bool = True) -> and other.camera_transform is not None and self.camera_transform is not None ): - self.pointcloud = self.pointcloud + other.pointcloud + self.pointcloud = concatenate_clouds(self.pointcloud, other.pointcloud) else: self.pointcloud = other.pointcloud @@ -115,20 +133,20 @@ def update_object(self, other: Object, *, accumulate_pointcloud: bool = True) -> def get_oriented_bounding_box(self) -> Any: """Get oriented bounding box of the pointcloud.""" - return self.pointcloud.oriented_bounding_box + return pointcloud_to_open3d(self.pointcloud).get_oriented_bounding_box() def _detection3d_bbox_components(self) -> tuple[Vector3, Quaternion, Vector3]: """Return canonical geometry without refitting the point cloud.""" center = self.center size = self.size - orientation = self.pose.orientation + orientation = self.pose.pose.orientation return ( - Vector3(center.x, center.y, center.z), - Quaternion(orientation.x, orientation.y, orientation.z, orientation.w), + Vector3(x=center.x, y=center.y, z=center.z), + Quaternion(x=orientation.x, y=orientation.y, z=orientation.z, w=orientation.w), Vector3( - max(float(size.x), 1e-3), - max(float(size.y), 1e-3), - max(float(size.z), 1e-3), + x=max(float(size.x), 1e-3), + y=max(float(size.y), 1e-3), + z=max(float(size.z), 1e-3), ), ) @@ -142,21 +160,38 @@ def to_detection3d_msg(self) -> ROSDetection3D: """Convert to ROS Detection3D message.""" center, orientation, size = self._detection3d_bbox_components() - msg = ROSDetection3D() - msg.header = Header(self.ts, self.frame_id) + msg = ROSDetection3D( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + results=[], + bbox=BoundingBox3D( + center=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), + ), + size=Vector3(x=0.0, y=0.0, z=0.0), + ), + id="", + ) + msg.header = self.pose.header msg.id = self.object_id msg.results = [ ObjectHypothesisWithPose( hypothesis=ObjectHypothesis( class_id=self.name, score=self.confidence, - ) + ), + 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), + ), ) ] - msg.results_length = len(msg.results) msg.bbox = BoundingBox3D( center=Pose( - position=center, + position=Point(x=center.x, y=center.y, z=center.z), orientation=orientation, ), size=size, @@ -179,6 +214,7 @@ def agent_encode(self) -> dict[str, Any]: def to_dict(self) -> dict[str, Any]: """Convert object to dictionary with all relevant data.""" + colors = pointcloud_rgb(self.pointcloud) return { "object_id": self.object_id, "track_id": self.track_id, @@ -188,8 +224,11 @@ def to_dict(self) -> dict[str, Any]: "identity_basis": self.identity_basis, "last_seen_ts": self.last_seen_ts, "mask": self.mask, - "pointcloud": self.pointcloud.as_numpy(), - "image": self.image.as_numpy() if self.image else None, + "pointcloud": ( + pointcloud_xyz(self.pointcloud), + None if colors is None else colors.astype(np.float64) / 255.0, + ), + "image": image_view(self.image) if self.image else None, } @classmethod @@ -199,7 +238,7 @@ def from_2d_to_list( color_image: Image, depth_image: Image, camera_info: CameraInfo, - camera_transform: Transform | None = None, + camera_transform: TransformStamped | None = None, depth_scale: float = 1.0, depth_trunc: float = 10.0, statistical_nb_neighbors: int = 10, @@ -241,16 +280,13 @@ def from_2d_to_list( import cv2 import open3d as o3d # type: ignore[import-untyped] - color_cv = color_image.to_opencv() - if color_cv.ndim == 3 and color_cv.shape[2] == 3: - color_cv = cv2.cvtColor(color_cv, cv2.COLOR_BGR2RGB) - - depth_cv = depth_image.to_opencv() + color_cv = image_to_rgb(color_image) + depth_cv = image_view(depth_image) h, w = depth_cv.shape[:2] # Build Open3D camera intrinsics - fx, fy = camera_info.K[0], camera_info.K[4] - cx, cy = camera_info.K[2], camera_info.K[5] + fx, fy = camera_info.k[0], camera_info.k[4] + cx, cy = camera_info.k[2], camera_info.k[5] intrinsic_o3d = o3d.camera.PinholeCameraIntrinsic(w, h, fx, fy, cx, cy) objects: list[Object] = [] @@ -289,56 +325,56 @@ def from_2d_to_list( ) pcd = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd, intrinsic_o3d) - pc0 = PointCloud2( - pcd, - frame_id=depth_image.frame_id, - ts=depth_image.ts, - ).voxel_downsample(voxel_downsample) - - pcd_filtered, _ = pc0.pointcloud.remove_statistical_outlier( + initial = pointcloud_from_xyz_rgb( + np.asarray(pcd.points), + (np.asarray(pcd.colors) * 255).astype(np.uint8), + header=depth_image.header, + ) + pc0 = voxel_downsample_cloud(initial, voxel_downsample) + pcd_filtered, _ = pointcloud_to_open3d(pc0).remove_statistical_outlier( nb_neighbors=statistical_nb_neighbors, std_ratio=statistical_std_ratio, ) - if len(pcd_filtered.points) < 10: continue - - pc = PointCloud2( - pcd_filtered, - frame_id=depth_image.frame_id, - ts=depth_image.ts, - ) - - # Transform pointcloud to world frame if camera_transform is provided + header = depth_image.header if camera_transform is not None: - pc = pc.transform(camera_transform) - frame_id = camera_transform.frame_id - else: - frame_id = depth_image.frame_id + if camera_transform.child_frame_id != header.frame_id: + raise ValueError("Camera transform source frame does not match depth image") + pcd_filtered.transform(transform_matrix(camera_transform.transform)) + header = Header( + stamp=depth_image.header.stamp, frame_id=camera_transform.header.frame_id + ) + pc = pointcloud_from_xyz_rgb( + np.asarray(pcd_filtered.points), + (np.asarray(pcd_filtered.colors) * 255).astype(np.uint8), + header=header, + ) + frame_id = header.frame_id # Compute bounding box: AABB for stable upright obstacles, OBB for tighter fit if use_aabb: - aabb = pc.pointcloud.get_axis_aligned_bounding_box() + aabb = pcd_filtered.get_axis_aligned_bounding_box() aabb_center = (aabb.min_bound + aabb.max_bound) / 2.0 aabb_extent = aabb.max_bound - aabb.min_bound - center = Vector3(aabb_center[0], aabb_center[1], aabb_center[2]) + center = Vector3(x=aabb_center[0], y=aabb_center[1], z=aabb_center[2]) sx, sy, sz = float(aabb_extent[0]), float(aabb_extent[1]), float(aabb_extent[2]) - orientation = Quaternion(0.0, 0.0, 0.0, 1.0) + orientation = Quaternion(w=1.0, x=0.0, y=0.0, z=0.0) else: - obb = pc.pointcloud.get_oriented_bounding_box() - center = Vector3(obb.center[0], obb.center[1], obb.center[2]) + obb = pcd_filtered.get_oriented_bounding_box() + center = Vector3(x=obb.center[0], y=obb.center[1], z=obb.center[2]) sx, sy, sz = float(obb.extent[0]), float(obb.extent[1]), float(obb.extent[2]) - orientation = Quaternion.from_rotation_matrix(obb.R) + orientation = quaternion_from_matrix(np.asarray(obb.R)) if max_obstacle_width > 0: sx = min(sx, max_obstacle_width) sy = min(sy, max_obstacle_width) - size = Vector3(sx, sy, sz) + size = Vector3(x=sx, y=sy, z=sz) pose = PoseStamped( - ts=det.ts, - frame_id=frame_id, - position=center, - orientation=orientation, + header=Header(stamp=det.image.header.stamp, frame_id=frame_id), + pose=Pose( + position=Point(x=center.x, y=center.y, z=center.z), orientation=orientation + ), ) # Skip objects too far from origin (background detections) @@ -380,52 +416,38 @@ def aggregate_pointclouds(objects: list[Object]) -> PointCloud2: Returns: Combined PointCloud2 with all points colored by object (empty if no points). """ - import open3d as o3d # type: ignore[import-untyped] - - if not objects: - return PointCloud2(pointcloud=o3d.geometry.PointCloud(), frame_id="", ts=0.0) - + header = ( + Header(stamp=Time(sec=0, nanosec=0), frame_id="") + if not objects + else Header(frame_id=objects[0].frame_id, stamp=objects[0].pointcloud.header.stamp) + ) all_points = [] all_colors = [] - for obj in objects: - points, colors = obj.pointcloud.as_numpy() + if obj.frame_id != header.frame_id: + raise ValueError("Cannot aggregate object clouds in different frames") + points = pointcloud_xyz(obj.pointcloud) if len(points) == 0: continue - + colors = pointcloud_rgb(obj.pointcloud) try: seed = int(obj.object_id, 16) except (ValueError, TypeError): seed = abs(hash(obj.object_id)) rng = np.random.default_rng(abs(seed)) track_color = rng.integers(50, 255, 3) / 255.0 - - if colors is not None: - blended = np.clip(0.6 * colors + 0.4 * track_color, 0.0, 1.0) - else: - blended = np.tile(track_color, (len(points), 1)) - + blended = ( + np.clip(0.6 * colors.astype(np.float64) / 255.0 + 0.4 * track_color, 0.0, 1.0) + if colors is not None + else np.tile(track_color, (len(points), 1)) + ) all_points.append(points) all_colors.append(blended) - if not all_points: - return PointCloud2( - pointcloud=o3d.geometry.PointCloud(), frame_id=objects[0].frame_id, ts=objects[0].ts - ) - - combined_points = np.vstack(all_points) - combined_colors = np.vstack(all_colors) - - pc = PointCloud2.from_numpy( - combined_points, - frame_id=objects[0].frame_id, - timestamp=objects[0].ts, + return pointcloud_from_xyz(np.empty((0, 3), dtype=np.float32), header=header) + return pointcloud_from_xyz_rgb( + np.vstack(all_points), (np.vstack(all_colors) * 255).astype(np.uint8), header=header ) - pcd = pc.pointcloud - pcd.colors = o3d.utility.Vector3dVector(combined_colors) - pc.pointcloud = pcd - - return pc def to_detection3d_array( @@ -448,11 +470,14 @@ def to_detection3d_array( resolved_frame_id = frame_id if resolved_frame_id is None: resolved_frame_id = objects[0].frame_id if objects else "" - resolved_ts = ts - if resolved_ts is None: - resolved_ts = objects[0].ts if objects else 0.0 + stamp = ( + time_from_seconds(ts) + if ts is not None + else objects[0].pose.header.stamp + if objects + else Time(sec=0, nanosec=0) + ) return Detection3DArray( - detections_length=len(detections), - header=Header(resolved_ts, resolved_frame_id), + header=Header(stamp=stamp, frame_id=resolved_frame_id), detections=detections, ) diff --git a/dimos/perception/experimental/objectDB.py b/dimos/perception/experimental/objectDB.py index d2d8e7e1bea..e8cf858350d 100644 --- a/dimos/perception/experimental/objectDB.py +++ b/dimos/perception/experimental/objectDB.py @@ -14,15 +14,19 @@ from __future__ import annotations +import math import threading import time from typing import TYPE_CHECKING, Any -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +import numpy as np + +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.msgs.time import to_nanoseconds from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: - from dimos.msgs.geometry_msgs.Vector3 import Vector3 + from dimos.messages.geometry_msgs.msg import Vector3 from dimos.perception.experimental.object import Object logger = setup_logger() @@ -176,20 +180,19 @@ def find_nearest( if not candidates: return None - return min(candidates, key=lambda obj: position.distance(obj.center)) + return min( + candidates, + key=lambda obj: math.dist( + (position.x, position.y, position.z), (obj.center.x, obj.center.y, obj.center.z) + ), + ) def clear(self) -> None: """Clear all objects from the database.""" - import open3d as o3d # type: ignore[import-untyped] - with self._lock: - # Drop Open3D pointcloud references before clearing to reduce shutdown warnings. + # Release large generated point buffers, retaining metadata on held objects. for obj in list(self._pending_objects.values()) + list(self._objects.values()): - obj.pointcloud = PointCloud2( - pointcloud=o3d.geometry.PointCloud(), - frame_id=obj.pointcloud.frame_id, - ts=obj.pointcloud.ts, - ) + obj.pointcloud = pointcloud_from_xyz(np.empty((0, 3)), header=obj.pointcloud.header) self._pending_objects.clear() self._objects.clear() self._track_id_map.clear() @@ -220,8 +223,6 @@ def _match(self, obj: Object, now: float) -> tuple[Object | None, str | None]: return None, None def _insert_pending(self, obj: Object, now: float) -> Object: - if not obj.ts: - obj.ts = now obj.last_seen_ts = now self._pending_objects[obj.object_id] = obj if obj.track_id >= 0: @@ -234,11 +235,13 @@ def _update_existing(self, existing: Object, obj: Object, now: float) -> bool: self._track_id_map[obj.track_id] = existing.object_id # Multiple prompts or repeated scans may produce the same object from # one camera frame. Only distinct source observations advance memory. - if existing.ts == obj.ts: + if (existing.image.header.frame_id, to_nanoseconds(existing.image.header.stamp)) == ( + obj.image.header.frame_id, + to_nanoseconds(obj.image.header.stamp), + ): return False existing.update_object(obj) - existing.ts = obj.ts or now existing.last_seen_ts = now return True @@ -283,13 +286,22 @@ def _match_by_distance(self, obj: Object) -> Object | None: candidates = [ o for o in all_objects - if o.center is not None and obj.center.distance(o.center) < self._distance_threshold + if o.center is not None + and math.dist( + (obj.center.x, obj.center.y, obj.center.z), (o.center.x, o.center.y, o.center.z) + ) + < self._distance_threshold ] if not candidates: return None - return min(candidates, key=lambda o: obj.center.distance(o.center)) + return min( + candidates, + key=lambda o: math.dist( + (obj.center.x, obj.center.y, obj.center.z), (o.center.x, o.center.y, o.center.z) + ), + ) def _prune_stale_pending(self, now: float) -> None: if self._pending_ttl_s <= 0: diff --git a/dimos/perception/experimental/object_scene_registration.py b/dimos/perception/experimental/object_scene_registration.py index 5e43cb490d0..4397fc79a4c 100644 --- a/dimos/perception/experimental/object_scene_registration.py +++ b/dimos/perception/experimental/object_scene_registration.py @@ -21,21 +21,21 @@ import numpy as np from numpy.typing import NDArray +from reactivex import operators as ops from dimos.agents.annotation import skill from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out +from dimos.messages.geometry_msgs.msg import TransformStamped +from dimos.messages.sensor_msgs.msg import CameraInfo, Image, PointCloud2 +from dimos.messages.tf2_msgs.msg import TFMessage +from dimos.messages.vision_msgs.msg import Detection2DArray, Detection3DArray from dimos.models.segmentation.edge_tam import BoxPromptImageSegmenter from dimos.models.segmentation.yoloe import YoloeBoxSegmenter -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.std_msgs.Header import Header -from dimos.msgs.tf2_msgs.TFMessage import TFMessage -from dimos.msgs.vision_msgs.Detection2DArray import Detection2DArray -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray +from dimos.msgs.image import image_from_array, image_view +from dimos.msgs.pointcloud import pointcloud_from_xyz, transform_cloud +from dimos.msgs.time import to_seconds from dimos.perception.detection.detectors.yoloe import Yoloe2DDetector, YoloePromptMode from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.perception.experimental.object import ( @@ -45,13 +45,23 @@ to_detection3d_array, ) from dimos.perception.experimental.objectDB import ObjectDB -from dimos.types.timestamped import align_timestamped +from dimos.types.timestamped import TimestampedData, align_timestamped from dimos.utils.logging_config import setup_logger from dimos.utils.reactive import backpressure logger = setup_logger() +def _timed_image(image: Image) -> TimestampedData[Image]: + return TimestampedData(image, to_seconds(image.header.stamp)) + + +def _aligned_images( + pair: tuple[TimestampedData[Image], TimestampedData[Image]], +) -> tuple[Image, Image]: + return pair[0].value, pair[1].value + + class ObjectSceneRegistrationConfig(ModuleConfig): target_frame: str = "base_link" prompt_mode: YoloePromptMode = YoloePromptMode.LRPC @@ -89,7 +99,7 @@ class ObjectSceneRegistrationModule(Module): _processing_lock: threading.RLock # A tuple assignment/read is atomic, so depth and its transform cannot be # observed from different frames by get_full_scene_pointcloud(). - _latest_scene_snapshot: tuple[Image, Transform | None] | None = None + _latest_scene_snapshot: tuple[Image, TransformStamped | None] | None = None config: ObjectSceneRegistrationConfig @@ -155,11 +165,11 @@ def start(self) -> None: self.camera_info.subscribe(lambda msg: setattr(self, "_camera_info", msg)) aligned_frames = align_timestamped( - self.color_image.observable(), - self.depth_image.observable(), + self.color_image.observable().pipe(ops.map(_timed_image)), + self.depth_image.observable().pipe(ops.map(_timed_image)), buffer_size=2.0, match_tolerance=0.1, - ) + ).pipe(ops.map(_aligned_images)) backpressure(aligned_frames).subscribe(self._on_aligned_frames) def _create_segmenter(self) -> BoxPromptImageSegmenter | None: @@ -229,11 +239,12 @@ def scan_scene(self, text: list[str] | None = None) -> Detection3DArray: if frames is None: return to_detection3d_array([], frame_id=self._target_frame) objects = self._scan_scene_objects(frames) - return to_detection3d_array( + result = to_detection3d_array( objects, frame_id=self._target_frame, - ts=frames[0].ts, ) + result.header.stamp = frames[0].header.stamp + return result def _scan_scene_objects(self, frames: tuple[Image, Image] | None = None) -> list[DetObject]: """Process one frame and return only objects observed during this scan.""" @@ -275,7 +286,7 @@ def get_object_pointcloud_by_object_id(self, object_id: str) -> PointCloud2 | No logger.warning(f"No object found with object_id='{object_id}'") return None pc = obj.pointcloud - num_points = len(pc.pointcloud.points) if pc else 0 + num_points = pc.width * pc.height logger.info(f"Found object '{object_id}' ({obj.name}) with {num_points} points") return pc @@ -313,7 +324,7 @@ def get_full_scene_pointcloud( return None depth_image, camera_transform = scene_snapshot - depth_cv = depth_image.to_opencv() + depth_cv = image_view(depth_image) h, w = depth_cv.shape[:2] # Zero out excluded object's depth @@ -324,8 +335,8 @@ def get_full_scene_pointcloud( depth_cv[exclude_mask > 0] = 0 # Build pointcloud from depth - fx, fy = self._camera_info.K[0], self._camera_info.K[4] - cx, cy = self._camera_info.K[2], self._camera_info.K[5] + fx, fy = self._camera_info.k[0], self._camera_info.k[4] + cx, cy = self._camera_info.k[2], self._camera_info.k[5] intrinsic = o3d.camera.PinholeCameraIntrinsic(w, h, fx, fy, cx, cy) depth_o3d = o3d.geometry.Image(depth_cv.astype(np.float32)) @@ -338,14 +349,9 @@ def get_full_scene_pointcloud( pcd = pcd.voxel_down_sample(voxel_size) - pc = PointCloud2( - pcd, - frame_id=depth_image.frame_id, - ts=depth_image.ts, - ) - + pc = pointcloud_from_xyz(np.asarray(pcd.points), header=depth_image.header) if camera_transform is not None: - pc = pc.transform(camera_transform) + pc = transform_cloud(pc, camera_transform) return pc @@ -412,21 +418,19 @@ def _process_images(self, color_msg: Image, depth_msg: Image) -> list[DetObject] color_image = color_msg # Convert depth to meters (float32) - depth_cv = depth_msg.to_opencv() - if depth_msg.format == ImageFormat.DEPTH16: + depth_cv = image_view(depth_msg) + if depth_msg.encoding == "16UC1": depth_cv = depth_cv.astype(np.float32) / 1000.0 elif depth_cv.dtype != np.float32: depth_cv = depth_cv.astype(np.float32) - depth_image = Image( - data=depth_cv, format=ImageFormat.DEPTH, frame_id=depth_msg.frame_id, ts=depth_msg.ts - ) + depth_image = image_from_array(depth_cv, encoding="32FC1", header=depth_msg.header) camera_transform = None - if self._target_frame != color_image.frame_id: + if self._target_frame != color_image.header.frame_id: camera_transform = self.tfbuffer.get( self._target_frame, - color_image.frame_id, - color_image.ts, + color_image.header.frame_id, + to_seconds(color_image.header.stamp), 0.1, forward_tolerance=0.2, ) @@ -457,8 +461,7 @@ def _process_images(self, color_msg: Image, depth_msg: Image) -> list[DetObject] detections_2d = self._segmenter.segment(detections_2d) detections_2d_msg = Detection2DArray( - detections_length=len(detections_2d.detections), - header=Header(color_image.ts, color_image.frame_id or ""), + header=color_image.header, detections=[det.to_ros_detection2d() for det in detections_2d.detections], ) self.detections_2d.publish(detections_2d_msg) @@ -473,13 +476,13 @@ def _process_3d_detections( detections_2d: ImageDetections2D[Any], color_image: Image, depth_image: Image, - camera_transform: Transform | None, + camera_transform: TransformStamped | None, ) -> list[DetObject]: """Convert 2D detections to 3D and publish.""" if self._camera_info is None: return [] - if self._target_frame != color_image.frame_id and camera_transform is None: + if self._target_frame != color_image.header.frame_id and camera_transform is None: logger.warning("Failed to lookup transform from camera frame to target frame") return [] diff --git a/dimos/perception/experimental/object_scene_registration_spec.py b/dimos/perception/experimental/object_scene_registration_spec.py index 5f0cbbf8e1b..a082529e424 100644 --- a/dimos/perception/experimental/object_scene_registration_spec.py +++ b/dimos/perception/experimental/object_scene_registration_spec.py @@ -14,8 +14,8 @@ from typing import Protocol -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.messages.vision_msgs.msg import Detection3DArray from dimos.spec.utils import Spec diff --git a/dimos/perception/experimental/object_tracker.py b/dimos/perception/experimental/object_tracker.py index 2331dbfd8ca..f14260a53c0 100644 --- a/dimos/perception/experimental/object_tracker.py +++ b/dimos/perception/experimental/object_tracker.py @@ -12,41 +12,57 @@ # See the License for the specific language governing permissions and # limitations under the License. +import math import threading import time from typing import Any -# Import LCM messages -from dimos_lcm.vision_msgs import ( - Detection2D, - Detection3D, - ObjectHypothesisWithPose, -) import numpy as np from numpy.typing import NDArray +from reactivex import operators as ops from reactivex.disposable import Disposable from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.msgs.std_msgs.Header import Header -from dimos.msgs.tf2_msgs.TFMessage import TFMessage -from dimos.msgs.vision_msgs.Detection2DArray import Detection2DArray -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray -from dimos.types.timestamped import align_timestamped -from dimos.utils.logging_config import setup_logger -from dimos.utils.transform_utils import ( - euler_to_quaternion, - optical_to_robot_frame, - yaw_towards_point, +from dimos.messages.builtin_interfaces.msg import Time +from dimos.messages.geometry_msgs.msg import ( + Point, + Pose, + PoseWithCovariance, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos.messages.sensor_msgs.msg import CameraInfo, Image +from dimos.messages.std_msgs.msg import Header +from dimos.messages.tf2_msgs.msg import TFMessage + +# Generated detection values +from dimos.messages.vision_msgs.msg import ( + BoundingBox2D, + BoundingBox3D, + Detection2D, + Detection2DArray, + Detection3D, + Detection3DArray, + ObjectHypothesis, + ObjectHypothesisWithPose, + Point2D, + Pose2D, ) +from dimos.msgs.geometry import quaternion_from_euler +from dimos.msgs.image import image_from_array, image_to_rgb, image_view +from dimos.msgs.time import time_from_seconds, to_seconds +from dimos.types.timestamped import TimestampedData, align_timestamped +from dimos.utils.logging_config import setup_logger + + +def _timed_image(image: Image) -> TimestampedData[Image]: + return TimestampedData(image, to_seconds(image.header.stamp)) + logger = setup_logger() @@ -128,20 +144,20 @@ def start(self) -> None: # Subscribe to aligned rgb and depth streams def on_aligned_frames(frames_tuple) -> None: # type: ignore[no-untyped-def] - rgb_msg, depth_msg = frames_tuple + rgb_msg, depth_msg = (item.value for item in frames_tuple) with self._frame_lock: - self._latest_rgb_frame = rgb_msg.data + self._latest_rgb_frame = image_to_rgb(rgb_msg) - depth_data = depth_msg.data + depth_data = image_view(depth_msg) # Convert from millimeters to meters if depth is DEPTH16 format - if depth_msg.format == ImageFormat.DEPTH16: + if depth_msg.encoding in ("16UC1", "mono16"): depth_data = depth_data.astype(np.float32) / 1000.0 self._latest_depth_frame = depth_data # Create aligned observable for RGB and depth aligned_frames = align_timestamped( - self.color_image.observable(), - self.depth.observable(), + self.color_image.pure_observable().pipe(ops.map(_timed_image)), + self.depth.pure_observable().pipe(ops.map(_timed_image)), buffer_size=2.0, # 2 second buffer match_tolerance=0.5, # 500ms tolerance ) @@ -154,10 +170,10 @@ def on_camera_info(camera_info_msg: CameraInfo) -> None: # Extract intrinsics from camera info K matrix # K is a 3x3 matrix in row-major order: [fx, 0, cx, 0, fy, cy, 0, 0, 1] self.camera_intrinsics = [ # type: ignore[assignment] - camera_info_msg.K[0], - camera_info_msg.K[4], - camera_info_msg.K[2], - camera_info_msg.K[5], + camera_info_msg.k[0], + camera_info_msg.k[4], + camera_info_msg.k[2], + camera_info_msg.k[5], ] unsub = self.camera_info.subscribe(on_camera_info) # type: ignore[assignment] @@ -309,8 +325,12 @@ def _reset_tracking_state(self) -> None: self.tracking_frame_count = 0 # Reset frame counter # Publish empty detections to clear any visualizations - empty_2d = Detection2DArray(detections_length=0, header=Header(), detections=[]) - empty_3d = Detection3DArray(detections_length=0, header=Header(), detections=[]) + empty_2d = Detection2DArray( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), detections=[] + ) + empty_3d = Detection3DArray( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), detections=[] + ) self._latest_detection2d = empty_2d self._latest_detection3d = empty_3d self._detection_event.clear() @@ -409,9 +429,9 @@ def _process_tracking(self) -> None: return # Create detections if tracking succeeded - header = Header(self.frame_id) - detection2darray = Detection2DArray(detections_length=0, header=header, detections=[]) - detection3darray = Detection3DArray(detections_length=0, header=header, detections=[]) + header = Header(stamp=time_from_seconds(time.time()), frame_id=self.config.frame_id) + detection2darray = Detection2DArray(header=header, detections=[]) + detection3darray = Detection3DArray(header=header, detections=[]) if final_success and current_bbox_x1y1x2y2 is not None: x1, y1, x2, y2 = current_bbox_x1y1x2y2 @@ -421,13 +441,28 @@ def _process_tracking(self) -> None: height = float(y2 - y1) # Create Detection2D - detection_2d = Detection2D() + detection_2d = Detection2D( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + results=[], + bbox=BoundingBox2D( + center=Pose2D(position=Point2D(x=0.0, y=0.0), theta=0.0), size_x=0.0, size_y=0.0 + ), + id="", + ) detection_2d.id = "0" - detection_2d.results_length = 1 detection_2d.header = header # Create hypothesis - hypothesis = ObjectHypothesisWithPose() + hypothesis = ObjectHypothesisWithPose( + hypothesis=ObjectHypothesis(class_id="", score=0.0), + 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), + ), + ) hypothesis.hypothesis.class_id = "tracked_object" hypothesis.hypothesis.score = 1.0 detection_2d.results = [hypothesis] @@ -439,8 +474,9 @@ def _process_tracking(self) -> None: detection_2d.bbox.size_x = width detection_2d.bbox.size_y = height - detection2darray = Detection2DArray() - detection2darray.detections_length = 1 + detection2darray = Detection2DArray( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), detections=[] + ) detection2darray.header = header detection2darray.detections = [detection_2d] @@ -460,18 +496,12 @@ def _process_tracking(self) -> None: x_optical = (center_x - cx) * z_optical / fx y_optical = (center_y - cy) * z_optical / fy - # Create pose in optical frame - optical_pose = Pose() - optical_pose.position = Vector3(x_optical, y_optical, z_optical) - optical_pose.orientation = Quaternion(0.0, 0.0, 0.0, 1.0) # Identity for now - - # Convert to robot frame - robot_pose = optical_to_robot_frame(optical_pose) - - # Calculate orientation: object facing towards camera (origin) - yaw = yaw_towards_point(robot_pose.position) - euler = Vector3(0.0, 0.0, yaw) # Only yaw, no roll/pitch - robot_pose.orientation = euler_to_quaternion(euler) + # Optical (right, down, forward) -> robot (forward, left, up). + position = Point(x=z_optical, y=-x_optical, z=-y_optical) + robot_pose = Pose( + position=position, + orientation=quaternion_from_euler(0, 0, math.atan2(position.y, position.x)), + ) # Estimate object size in meters size_x = width * z_optical / fx @@ -479,42 +509,61 @@ def _process_tracking(self) -> None: size_z = 0.1 # Default depth size # Create Detection3D - detection_3d = Detection3D() + detection_3d = Detection3D( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + results=[], + bbox=BoundingBox3D( + center=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), + ), + size=Vector3(x=0.0, y=0.0, z=0.0), + ), + id="", + ) detection_3d.id = "0" - detection_3d.results_length = 1 detection_3d.header = header # Reuse hypothesis from 2D detection_3d.results = [hypothesis] # Create 3D bounding box with robot frame pose - detection_3d.bbox.center = Pose() + detection_3d.bbox.center = 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), + ) detection_3d.bbox.center.position = robot_pose.position detection_3d.bbox.center.orientation = robot_pose.orientation - detection_3d.bbox.size = Vector3(size_x, size_y, size_z) + detection_3d.bbox.size = Vector3(x=size_x, y=size_y, z=size_z) - detection3darray = Detection3DArray() - detection3darray.detections_length = 1 + detection3darray = Detection3DArray( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), detections=[] + ) detection3darray.header = header detection3darray.detections = [detection_3d] # Publish transform for tracked object # The optical pose is in camera optical frame, so publish it relative to the camera frame - tracked_object_tf = Transform( - translation=robot_pose.position, - rotation=robot_pose.orientation, - frame_id=self.frame_id, # Use configured camera frame + tracked_object_tf = TransformStamped( + header=header, child_frame_id="tracked_object", - ts=header.ts, + transform=Transform( + translation=Vector3( + x=robot_pose.position.x, + y=robot_pose.position.y, + z=robot_pose.position.z, + ), + rotation=robot_pose.orientation, + ), ) - self.tf.publish(TFMessage(tracked_object_tf)) + self.tf.publish(TFMessage(transforms=[tracked_object_tf])) # Store latest detections for RPC access self._latest_detection2d = detection2darray self._latest_detection3d = detection3darray # Signal that new detections are available - if detection2darray.detections_length > 0 or detection3darray.detections_length > 0: + if bool(detection2darray.detections) or bool(detection3darray.detections): self._detection_event.set() # Publish detections @@ -524,12 +573,8 @@ def _process_tracking(self) -> None: # Create and publish visualization if tracking is active if self.tracking_initialized: # Convert single detection to list for visualization - detections_3d = ( - detection3darray.detections if detection3darray.detections_length > 0 else [] - ) - detections_2d = ( - detection2darray.detections if detection2darray.detections_length > 0 else [] - ) + detections_3d = detection3darray.detections if bool(detection3darray.detections) else [] + detections_2d = detection2darray.detections if bool(detection2darray.detections) else [] if detections_3d and detections_2d: # Extract 2D bbox for visualization @@ -555,7 +600,9 @@ def _process_tracking(self) -> None: viz_image = self._draw_reid_matches(viz_image) # Convert to Image message and publish - viz_msg = Image.from_numpy(viz_image) + viz_msg = image_from_array( + viz_image, encoding="rgb8", header=self._latest_detection2d.header + ) self.tracked_overlay.publish(viz_msg) def _draw_reid_matches(self, image: NDArray[np.uint8]) -> NDArray[np.uint8]: diff --git a/dimos/perception/experimental/object_tracker_2d.py b/dimos/perception/experimental/object_tracker_2d.py index 385d0ff98d8..a83bd0f10c4 100644 --- a/dimos/perception/experimental/object_tracker_2d.py +++ b/dimos/perception/experimental/object_tracker_2d.py @@ -17,15 +17,6 @@ import time from typing import Any -# Import LCM messages -from dimos_lcm.vision_msgs import ( - BoundingBox2D, - Detection2D, - ObjectHypothesis, - ObjectHypothesisWithPose, - Point2D, - Pose2D, -) import numpy as np from numpy.typing import NDArray from reactivex.disposable import Disposable @@ -34,9 +25,22 @@ from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.msgs.std_msgs.Header import Header -from dimos.msgs.vision_msgs.Detection2DArray import Detection2DArray +from dimos.messages.geometry_msgs.msg import Point, Pose, PoseWithCovariance, Quaternion +from dimos.messages.sensor_msgs.msg import Image +from dimos.messages.std_msgs.msg import Header + +# Generated detection values +from dimos.messages.vision_msgs.msg import ( + BoundingBox2D, + Detection2D, + Detection2DArray, + ObjectHypothesis, + ObjectHypothesisWithPose, + Point2D, + Pose2D, +) +from dimos.msgs.image import image_from_array, image_to_rgb +from dimos.msgs.time import time_from_seconds from dimos.utils.logging_config import setup_logger logger = setup_logger(level=logging.INFO) @@ -91,7 +95,7 @@ def start(self) -> None: def on_frame(frame_msg: Image) -> None: arrival_time = time.perf_counter() with self._frame_lock: - self._latest_rgb_frame = frame_msg.data + self._latest_rgb_frame = image_to_rgb(frame_msg) self._frame_arrival_time = arrival_time unsub = self.color_image.subscribe(on_frame) @@ -176,7 +180,8 @@ def _reset_tracking_state(self) -> None: # Publish empty detection empty_2d = Detection2DArray( - detections_length=0, header=Header(time.time(), self.frame_id), detections=[] + header=Header(stamp=time_from_seconds(time.time()), frame_id=self.config.frame_id), + detections=[], ) self._latest_detection2d = empty_2d self.detection2darray.publish(empty_2d) @@ -260,12 +265,11 @@ def _process_tracking(self) -> None: height = float(y2 - y1) # Create 2D detection header - header = Header(time.time(), self.frame_id) + header = Header(stamp=time_from_seconds(time.time()), frame_id=self.config.frame_id) # Create Detection2D with all fields in constructors detection_2d = Detection2D( id="0", - results_length=1, header=header, bbox=BoundingBox2D( center=Pose2D(position=Point2D(x=center_x, y=center_y), theta=0.0), @@ -274,14 +278,19 @@ def _process_tracking(self) -> None: ), results=[ ObjectHypothesisWithPose( - hypothesis=ObjectHypothesis(class_id="tracked_object", score=1.0) + hypothesis=ObjectHypothesis(class_id="tracked_object", score=1.0), + 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), + ), ) ], ) - detection2darray = Detection2DArray( - detections_length=1, header=header, detections=[detection_2d] - ) + detection2darray = Detection2DArray(header=header, detections=[detection_2d]) # Store and publish self._latest_detection2d = detection2darray @@ -290,7 +299,7 @@ def _process_tracking(self) -> None: # Create visualization viz_image = self._draw_visualization(frame, current_bbox_x1y1x2y2) viz_copy = viz_image.copy() # Force copy needed to prevent frame reuse - viz_msg = Image.from_numpy(viz_copy, format=ImageFormat.RGB) + viz_msg = image_from_array(viz_copy, encoding="rgb8", header=header) self.tracked_overlay.publish(viz_msg) def _draw_visualization(self, image: NDArray[np.uint8], bbox: list[int]) -> NDArray[np.uint8]: diff --git a/dimos/perception/experimental/object_tracker_3d.py b/dimos/perception/experimental/object_tracker_3d.py index f2f5a116b10..9fef533934b 100644 --- a/dimos/perception/experimental/object_tracker_3d.py +++ b/dimos/perception/experimental/object_tracker_3d.py @@ -13,33 +13,47 @@ # limitations under the License. -# Import LCM messages -from dimos_lcm.sensor_msgs import CameraInfo -from dimos_lcm.vision_msgs import ( - Detection3D, - ObjectHypothesisWithPose, -) +import math + import numpy as np +from reactivex import operators as ops from dimos.core.core import rpc from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.msgs.std_msgs.Header import Header -from dimos.msgs.tf2_msgs.TFMessage import TFMessage -from dimos.msgs.vision_msgs.Detection2DArray import Detection2DArray -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray + +# Generated detection values +from dimos.messages.builtin_interfaces.msg import Time +from dimos.messages.geometry_msgs.msg import ( + Point, + Pose, + PoseWithCovariance, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos.messages.sensor_msgs.msg import CameraInfo, Image +from dimos.messages.std_msgs.msg import Header +from dimos.messages.tf2_msgs.msg import TFMessage +from dimos.messages.vision_msgs.msg import ( + BoundingBox3D, + Detection2DArray, + Detection3D, + Detection3DArray, + ObjectHypothesis, + ObjectHypothesisWithPose, +) +from dimos.msgs.geometry import quaternion_from_euler +from dimos.msgs.image import image_from_array, image_to_rgb, image_view +from dimos.msgs.time import to_seconds from dimos.perception.experimental.object_tracker_2d import ObjectTracker2D -from dimos.types.timestamped import align_timestamped +from dimos.types.timestamped import TimestampedData, align_timestamped from dimos.utils.logging_config import setup_logger -from dimos.utils.transform_utils import ( - euler_to_quaternion, - optical_to_robot_frame, - yaw_towards_point, -) + + +def _timed_image(image: Image) -> TimestampedData[Image]: + return TimestampedData(image, to_seconds(image.header.stamp)) + logger = setup_logger() @@ -78,20 +92,20 @@ def start(self) -> None: # Subscribe to aligned RGB and depth streams def on_aligned_frames(frames_tuple) -> None: # type: ignore[no-untyped-def] - rgb_msg, depth_msg = frames_tuple + rgb_msg, depth_msg = (item.value for item in frames_tuple) with self._frame_lock: - self._latest_rgb_frame = rgb_msg.data + self._latest_rgb_frame = image_to_rgb(rgb_msg) - depth_data = depth_msg.data + depth_data = image_view(depth_msg) # Convert from millimeters to meters if depth is DEPTH16 format - if depth_msg.format == ImageFormat.DEPTH16: + if depth_msg.encoding in ("16UC1", "mono16"): depth_data = depth_data.astype(np.float32) / 1000.0 self._latest_depth_frame = depth_data # Create aligned observable for RGB and depth aligned_frames = align_timestamped( - self.color_image.observable(), - self.depth.observable(), + self.color_image.pure_observable().pipe(ops.map(_timed_image)), + self.depth.pure_observable().pipe(ops.map(_timed_image)), buffer_size=2.0, # 2 second buffer match_tolerance=0.5, # 500ms tolerance ) @@ -103,10 +117,10 @@ def on_camera_info(camera_info_msg: CameraInfo) -> None: self._latest_camera_info = camera_info_msg # Extract intrinsics: K is [fx, 0, cx, 0, fy, cy, 0, 0, 1] self.camera_intrinsics = [ # type: ignore[assignment] - camera_info_msg.K[0], - camera_info_msg.K[4], - camera_info_msg.K[2], - camera_info_msg.K[5], + camera_info_msg.k[0], + camera_info_msg.k[4], + camera_info_msg.k[2], + camera_info_msg.k[5], ] self.camera_info.subscribe(on_camera_info) @@ -127,7 +141,7 @@ def _process_tracking(self) -> None: # Enhance with 3D if we have depth and a valid 2D detection if ( self._latest_detection2d - and self._latest_detection2d.detections_length > 0 + and bool(self._latest_detection2d.detections) and self._latest_depth_frame is not None and self.camera_intrinsics is not None ): @@ -161,12 +175,14 @@ def _process_tracking(self) -> None: if self.last_good_matches and self.last_roi_kps and self.last_roi_bbox: viz_image = self._draw_reid_overlay(viz_image) - viz_msg = Image.from_numpy(viz_image) + viz_msg = image_from_array( + viz_image, encoding="rgb8", header=self._latest_detection2d.header + ) self.tracked_overlay.publish(viz_msg) def _create_detection3d_from_2d(self, detection2d: Detection2DArray) -> Detection3DArray | None: """Create 3D detection from 2D detection using depth.""" - if detection2d.detections_length == 0: + if not detection2d.detections: return None det_2d = detection2d.detections[0] @@ -196,18 +212,12 @@ def _create_detection3d_from_2d(self, detection2d: Detection2DArray) -> Detectio x_optical = (center_x - cx) * z_optical / fx # type: ignore[has-type] y_optical = (center_y - cy) * z_optical / fy # type: ignore[has-type] - # Create pose in optical frame - optical_pose = Pose() - optical_pose.position = Vector3(x_optical, y_optical, z_optical) - optical_pose.orientation = Quaternion(0.0, 0.0, 0.0, 1.0) - - # Convert to robot frame - robot_pose = optical_to_robot_frame(optical_pose) - - # Calculate orientation: object facing towards camera - yaw = yaw_towards_point(robot_pose.position) - euler = Vector3(0.0, 0.0, yaw) - robot_pose.orientation = euler_to_quaternion(euler) + # Optical (right, down, forward) -> robot (forward, left, up). + position = Point(x=z_optical, y=-x_optical, z=-y_optical) + robot_pose = Pose( + position=position, + orientation=quaternion_from_euler(0, 0, math.atan2(position.y, position.x)), + ) # Estimate object size in meters size_x = width * z_optical / fx # type: ignore[has-type] @@ -215,38 +225,63 @@ def _create_detection3d_from_2d(self, detection2d: Detection2DArray) -> Detectio size_z = 0.1 # Default depth size # Create Detection3D - header = Header(self.frame_id) - detection_3d = Detection3D() + header = Header(stamp=detection2d.header.stamp, frame_id=self.config.frame_id) + detection_3d = Detection3D( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + results=[], + bbox=BoundingBox3D( + center=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), + ), + size=Vector3(x=0.0, y=0.0, z=0.0), + ), + id="", + ) detection_3d.id = "0" - detection_3d.results_length = 1 detection_3d.header = header # Create hypothesis - hypothesis = ObjectHypothesisWithPose() + hypothesis = ObjectHypothesisWithPose( + hypothesis=ObjectHypothesis(class_id="", score=0.0), + 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), + ), + ) hypothesis.hypothesis.class_id = "tracked_object" hypothesis.hypothesis.score = 1.0 detection_3d.results = [hypothesis] # Create 3D bounding box - detection_3d.bbox.center = Pose() + detection_3d.bbox.center = 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) + ) detection_3d.bbox.center.position = robot_pose.position detection_3d.bbox.center.orientation = robot_pose.orientation - detection_3d.bbox.size = Vector3(size_x, size_y, size_z) + detection_3d.bbox.size = Vector3(x=size_x, y=size_y, z=size_z) - detection3darray = Detection3DArray() - detection3darray.detections_length = 1 + detection3darray = Detection3DArray( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), detections=[] + ) detection3darray.header = header detection3darray.detections = [detection_3d] # Publish TF for tracked object - tracked_object_tf = Transform( - translation=robot_pose.position, - rotation=robot_pose.orientation, - frame_id=self.frame_id, + tracked_object_tf = TransformStamped( + header=header, child_frame_id="tracked_object", - ts=header.ts, + transform=Transform( + translation=Vector3( + x=robot_pose.position.x, y=robot_pose.position.y, z=robot_pose.position.z + ), + rotation=robot_pose.orientation, + ), ) - self.tf.publish(TFMessage(tracked_object_tf)) + self.tf.publish(TFMessage(transforms=[tracked_object_tf])) return detection3darray diff --git a/dimos/perception/experimental/spatial_perception.py b/dimos/perception/experimental/spatial_perception.py index 54f15cbd7b4..4953140d19b 100644 --- a/dimos/perception/experimental/spatial_perception.py +++ b/dimos/perception/experimental/spatial_perception.py @@ -19,7 +19,7 @@ from datetime import datetime import os import time -from typing import TYPE_CHECKING, Any +from typing import Any import uuid import numpy as np @@ -30,17 +30,17 @@ from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In -from dimos.msgs.sensor_msgs.Image import Image -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.messages.geometry_msgs.msg import Point +from dimos.messages.sensor_msgs.msg import Image +from dimos.messages.tf2_msgs.msg import TFMessage +from dimos.msgs.geometry import pose_from_transform, quaternion_euler +from dimos.msgs.image import image_to_bgr from dimos.perception.experimental.image_embedding import ImageEmbeddingProvider from dimos.perception.experimental.spatial_vector_db import SpatialVectorDB from dimos.perception.experimental.visual_memory import VisualMemory from dimos.types.robot_location import RobotLocation from dimos.utils.logging_config import setup_logger -if TYPE_CHECKING: - from dimos.msgs.geometry_msgs.Vector3 import Vector3 - _OUTPUT_DIR = DIMOS_PROJECT_ROOT / "assets" / "output" _MEMORY_DIR = _OUTPUT_DIR / "memory" _SPATIAL_MEMORY_DIR = _MEMORY_DIR / "spatial_memory" @@ -154,7 +154,7 @@ def __init__(self, **kwargs: Any) -> None: embedding_provider=self.embedding_provider, ) - self.last_position: Vector3 | None = None + self.last_position: Point | None = None self.last_record_time: float | None = None self.frame_count: int = 0 @@ -167,23 +167,14 @@ def __init__(self, **kwargs: Any) -> None: self._latest_video_frame: np.ndarray | None = None self._process_interval = 1 + def _receive_frame(self, image_msg: Image) -> None: + self._latest_video_frame = image_to_bgr(image_msg) + @rpc def start(self) -> None: - import cv2 - super().start() - # Subscribe to LCM streams - def set_video(image_msg: Image) -> None: - # Convert Image message to numpy array - if hasattr(image_msg, "data"): - frame = image_msg.data - frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) - self._latest_video_frame = frame - else: - logger.warning("Received image message without data attribute") - - self.register_disposable(Disposable(self.color_image.subscribe(set_video))) + self.register_disposable(Disposable(self.color_image.subscribe(self._receive_frame))) # Start periodic processing using interval self.register_disposable( @@ -211,7 +202,7 @@ def _process_frame(self) -> None: return # Create Pose object with position and orientation - current_pose = tf.to_pose() + current_pose = pose_from_transform(tf).pose # Process the frame directly try: @@ -242,16 +233,16 @@ def _process_frame(self) -> None: frame_id = f"frame_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:8]}" # Get euler angles from quaternion orientation for metadata - euler = tf.rotation.to_euler() + roll, pitch, yaw = quaternion_euler(tf.transform.rotation) # Create metadata dictionary with primitive types only metadata = { "pos_x": float(current_pose.position.x), "pos_y": float(current_pose.position.y), "pos_z": float(current_pose.position.z), - "rot_x": float(euler.x), - "rot_y": float(euler.y), - "rot_z": float(euler.z), + "rot_x": roll, + "rot_y": pitch, + "rot_z": yaw, "timestamp": current_time, "frame_id": frame_id, } @@ -271,7 +262,7 @@ def _process_frame(self) -> None: logger.info( f"Stored frame at position ({current_pose.position.x:.2f}, {current_pose.position.y:.2f}, {current_pose.position.z:.2f}), " - f"rotation ({euler.x:.2f}, {euler.y:.2f}, {euler.z:.2f}) " + f"rotation ({roll:.2f}, {pitch:.2f}, {yaw:.2f}) " f"stored {self.stored_frame_count}/{self.frame_count} frames" ) @@ -374,7 +365,7 @@ def process_combined_data(data): # type: ignore[no-untyped-def] vector_id=frame_id, image=frame, embedding=frame_embedding, metadata=metadata ) - self.last_position = position_v3 + self.last_position = Point(x=position_v3.x, y=position_v3.y, z=position_v3.z) self.last_record_time = current_time self.stored_frame_count += 1 @@ -477,11 +468,12 @@ def add_named_location( return False # Create RobotLocation object - location = RobotLocation( # type: ignore[call-arg] + translation = tf.transform.translation + location = RobotLocation( name=name, - position=tf.translation, - rotation=tf.rotation.to_euler(), - description=description or f"Location: {name}", + position=(translation.x, translation.y, translation.z), + rotation=quaternion_euler(tf.transform.rotation), + metadata={"description": description or f"Location: {name}"}, timestamp=time.time(), ) diff --git a/dimos/perception/experimental/test_generated_tracker.py b/dimos/perception/experimental/test_generated_tracker.py new file mode 100644 index 00000000000..bd915755548 --- /dev/null +++ b/dimos/perception/experimental/test_generated_tracker.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. + +from contextlib import ExitStack +from threading import Event + +from dimos_cdrgen.registry import decode, encode +import numpy as np +import pytest + +from dimos.core.cdr_transport import CdrTransport +from dimos.messages.builtin_interfaces.msg import Time +from dimos.messages.std_msgs.msg import Header +from dimos.messages.vision_msgs.msg import ( + BoundingBox2D, + Detection2D, + Detection2DArray, + Point2D, + Pose2D, +) +from dimos.msgs.image import image_view +from dimos.perception.experimental.object_tracker_2d import ObjectTracker2D +from dimos.perception.experimental.object_tracker_3d import ObjectTracker3D + + +@pytest.fixture +def tracker_factory(): + with ExitStack() as cleanup: + + def create(cls): + transports = cleanup.enter_context(ExitStack()) + module = cls(frame_id="camera") + cleanup.callback(module.stop) + for name, output in module.outputs.items(): + wire = CdrTransport("tracker_" + name, output.type, backend="lcm", url="memq://") + transports.callback(wire.stop) + output.transport = wire + return module + + yield create, cleanup + + +@pytest.mark.parametrize("succeeded", [True, False]) +def test_tracker_publishes_generated_detection_and_clears_failed_track( + tracker_factory, mocker, succeeded +): + create, cleanup = tracker_factory + module = create(ObjectTracker2D) + tracker = mocker.Mock() + tracker.update.return_value = (succeeded, (10, 12, 20, 16)) + mocker.patch.object(module, "tracker", tracker) + module.tracking_initialized = True + module._latest_rgb_frame = np.zeros((64, 64, 3), dtype=np.uint8) + detections, overlays = [], [] + detection_ready, overlay_ready = Event(), Event() + + def receive_detection(value): + detections.append(value) + detection_ready.set() + + def receive_overlay(value): + overlays.append(value) + overlay_ready.set() + + cleanup.callback(module.detection2darray.transport.subscribe(receive_detection)) + cleanup.callback(module.tracked_overlay.transport.subscribe(receive_overlay)) + module._process_tracking() + assert detection_ready.wait(2) + assert len(detections) == 1 + result = detections[0] + assert result.header.frame_id == "camera" + if not succeeded: + assert not result.detections + assert not module.tracking_initialized + assert not overlays + return + assert overlay_ready.wait(2) + (detection,) = result.detections + assert (detection.bbox.center.position.x, detection.bbox.center.position.y) == (20, 20) + assert (detection.bbox.size_x, detection.bbox.size_y) == (20, 16) + assert detection.results[0].hypothesis.class_id == "tracked_object" + assert len(overlays) == 1 + assert overlays[0].header == result.header + assert image_view(overlays[0]).shape == (64, 64, 3) + assert overlays[0].encoding == "rgb8" + + +@pytest.mark.parametrize("depth", [2.0, 0.0, float("nan")]) +def test_tracker3d_projection_preserves_depth_filter_and_source_stamp(tracker_factory, depth): + create, cleanup = tracker_factory + module = create(ObjectTracker3D) + module.camera_intrinsics = [100.0, 100.0, 32.0, 32.0] + module._latest_depth_frame = np.full((64, 64), depth, dtype=np.float32) + header = Header(stamp=Time(sec=1700000000, nanosec=123456789), frame_id="camera") + detection = Detection2DArray( + header=header, + detections=[ + Detection2D( + header=header, + bbox=BoundingBox2D( + center=Pose2D(position=Point2D(x=32, y=32), theta=0), size_x=20, size_y=16 + ), + results=[], + id="", + ) + ], + ) + received = [] + ready = Event() + + def receive(value): + received.append(value) + ready.set() + + cleanup.callback(module.tf.transport.subscribe(receive)) + output = module._create_detection3d_from_2d(detection) + if not np.isfinite(depth) or depth <= 0: + assert output is None + assert not received + return + result = decode(encode(output), type(output)) + assert result.header == header + box = result.detections[0].bbox + assert (box.center.position.x, box.center.position.y, box.center.position.z) == (2, 0, 0) + np.testing.assert_allclose([box.size.x, box.size.y, box.size.z], [0.4, 0.32, 0.1]) + assert ready.wait(2) + assert len(received) == 1 + (transform,) = received[0].transforms + assert transform.header == header + assert transform.child_frame_id == "tracked_object" + assert transform.transform.translation.x == 2 + assert transform.transform.rotation.w == 1 diff --git a/dimos/perception/experimental/test_object_generated.py b/dimos/perception/experimental/test_object_generated.py new file mode 100644 index 00000000000..53cecb53da9 --- /dev/null +++ b/dimos/perception/experimental/test_object_generated.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. + +"""Object geometry and deduplication use generated ROS values.""" + +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +import numpy as np + +from dimos.messages.builtin_interfaces.msg import Time +from dimos.messages.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos.messages.sensor_msgs.msg import Image, PointCloud2 +from dimos.messages.std_msgs.msg import Header +from dimos.messages.vision_msgs.msg import Detection3D +from dimos.msgs.camera_info import camera_info_from_intrinsics +from dimos.msgs.image import image_from_array +from dimos.msgs.pointcloud import pointcloud_from_xyz_rgb, pointcloud_rgb, pointcloud_xyz +from dimos.msgs.time import time_from_seconds +from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D +from dimos.perception.detection.type.detection2d.seg import Detection2DSeg +from dimos.perception.experimental.object import Object, aggregate_pointclouds, to_detection3d_array +from dimos.perception.experimental.objectDB import ObjectDB + + +def _object(object_id, x, timestamp, source_stamp=None): + stamp = time_from_seconds(timestamp) if source_stamp is None else source_stamp + return Object( + object_id=object_id, + name="cup", + center=Vector3(x=x, y=0.0, z=0.0), + size=Vector3(x=0.1, y=0.2, z=0.3), + pose=PoseStamped( + header=Header(frame_id="world", stamp=stamp), + 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), + ), + ), + pointcloud=PointCloud2( + header=Header(stamp=stamp, frame_id="world"), + height=0, + width=0, + fields=[], + is_bigendian=False, + point_step=0, + row_step=0, + data=np.array([], dtype=np.uint8), + is_dense=False, + ), + bbox=(0.0, 0.0, 1.0, 1.0), + track_id=-1, + class_id=0, + confidence=0.9, + ts=timestamp, + frame_id="world", + image=Image( + header=Header(stamp=stamp, frame_id="world"), + height=0, + width=0, + encoding="", + is_bigendian=0, + step=0, + data=np.array([], dtype=np.uint8), + ), + ) + + +def test_detection_cdr_preserves_source_pose_header_and_current_center(): + obj = _object("cup-1", 0.2, 1.0, source_stamp=Time(sec=1700000000, nanosec=123456789)) + obj.set_center(Vector3(x=0.4, y=0.5, z=0.6)) + decoded = cdr_decode(cdr_encode(obj.to_detection3d_msg()), Detection3D) + assert decoded.header == obj.pose.header + assert decoded.bbox.center == obj.pose.pose + assert decoded.bbox.size == obj.size + assert decoded.results[0].hypothesis.class_id == "cup" + assert decoded.results[0].hypothesis.score == 0.9 + array = to_detection3d_array([obj]) + assert array.header == obj.pose.header + decoded.bbox.center.position.x = 99 + assert obj.center.x == 0.4 + assert obj.pose.pose.position.x == 0.4 + + +def test_spatial_deduplication_and_nearest_lookup_use_generated_centers(): + database = ObjectDB(distance_threshold=0.2, min_detections_for_permanent=1) + first = _object("first", 0.0, 1.0) + nearby = _object("nearby", 0.1, 2.0) + distant = _object("distant", 2.0, 3.0) + assert database.add_objects([first]) == [first] + assert database.add_objects([nearby]) == [first] + assert first.center.x == 0.1 + assert database.add_objects([distant]) == [distant] + assert database.find_nearest(Vector3(x=1.8, y=0.0, z=0.0)) is distant + assert database.find_nearest(Vector3(x=0.2, y=0.0, z=0.0)) is first + + +def test_generated_object_cloud_aggregation_and_dictionary_preserve_values(): + obj = _object("abc", 0.2, 1.0) + stamp = Time(sec=1700000000, nanosec=987654321) + obj.pointcloud = pointcloud_from_xyz_rgb( + np.array([[1.0, 2.0, 3.0]], dtype=np.float32), + np.array([[255, 128, 0]], dtype=np.uint8), + header=Header(frame_id="world", stamp=stamp), + ) + pixels = np.zeros((2, 3, 3), dtype=np.uint8) + obj.image = image_from_array(pixels, encoding="rgb8", header=obj.pointcloud.header) + exported = obj.to_dict() + np.testing.assert_array_equal(exported["pointcloud"][0], [[1.0, 2.0, 3.0]]) + np.testing.assert_allclose(exported["pointcloud"][1], [[1.0, 128 / 255.0, 0.0]]) + np.testing.assert_array_equal(exported["image"], pixels) + merged = aggregate_pointclouds([obj]) + decoded = cdr_decode(cdr_encode(merged), PointCloud2) + assert decoded.header.stamp == stamp and decoded.header.frame_id == "world" + np.testing.assert_array_equal(pointcloud_xyz(decoded), [[1.0, 2.0, 3.0]]) + assert pointcloud_rgb(decoded) is not None + assert aggregate_pointclouds([]).width == 0 + + +def test_zero_and_adjacent_source_stamps_are_distinct_observations(): + for sec in [0, 1700000000]: + database = ObjectDB(min_detections_for_permanent=2) + stamp = Time(sec=sec, nanosec=0) + first = _object("first", 0, float(sec), source_stamp=stamp) + duplicate = _object("duplicate", 0, float(sec), source_stamp=stamp) + newer = _object("newer", 0, float(sec), source_stamp=Time(sec=sec, nanosec=1)) + database.add_objects([first, duplicate]) + assert first.detections_count == 1 + assert first.ts == float(sec) + database.add_objects([newer]) + assert first.detections_count == 2 + assert first.pose.header.stamp == newer.pose.header.stamp + assert database.get_objects() == [first] + + +def test_actual_rgbd_object_projection_keeps_color_and_depth_source_headers(): + color_header = Header(frame_id="camera", stamp=Time(sec=1700000000, nanosec=123456788)) + depth_header = Header(frame_id="camera", stamp=Time(sec=1700000000, nanosec=123456789)) + pixels = np.zeros((32, 32, 3), dtype=np.uint8) + pixels[:, :, 0] = 255 + color = image_from_array(pixels, encoding="rgb8", header=color_header) + depth = image_from_array( + np.full((32, 32), 2, dtype=np.float32), encoding="32FC1", header=depth_header + ) + mask = np.zeros((32, 32), dtype=np.uint8) + mask[8:24, 8:24] = 255 + detection = Detection2DSeg( + bbox=(8, 8, 24, 24), + track_id=7, + class_id=0, + confidence=0.9, + name="cup", + ts=1700000000.1234567, + image=color, + mask=mask, + ) + camera = camera_info_from_intrinsics(100, 100, 16, 16, 32, 32, header=color_header) + mount = TransformStamped( + header=Header(frame_id="world", stamp=Time(sec=1, nanosec=0)), + child_frame_id="camera", + transform=Transform( + translation=Vector3(x=1, y=0, z=0), rotation=Quaternion(x=0, y=0, z=0, w=1) + ), + ) + objects = Object.from_2d_to_list( + ImageDetections2D(color, [detection]), + color, + depth, + camera, + mount, + use_aabb=True, + mask_erode_pixels=0, + statistical_std_ratio=10, + ) + assert len(objects) == 1 + obj = objects[0] + assert obj.pose.header.stamp == color_header.stamp + assert obj.pointcloud.header.stamp == depth_header.stamp + assert obj.pose.header.frame_id == obj.pointcloud.header.frame_id == "world" + np.testing.assert_allclose(pointcloud_xyz(obj.pointcloud)[:, 2], 2) + np.testing.assert_array_equal(pointcloud_rgb(obj.pointcloud)[:, 0], 255) + assert 0.9 < obj.center.x < 1.1 + wire = cdr_decode(cdr_encode(obj.to_detection3d_msg()), Detection3D) + assert wire.header == obj.pose.header + assert wire.id == obj.object_id diff --git a/dimos/perception/experimental/test_object_scene_registration_temporal.py b/dimos/perception/experimental/test_object_scene_registration_temporal.py index 82f14417f9c..7446d649355 100644 --- a/dimos/perception/experimental/test_object_scene_registration_temporal.py +++ b/dimos/perception/experimental/test_object_scene_registration_temporal.py @@ -14,8 +14,7 @@ from __future__ import annotations -from collections.abc import Iterator -import sys +from contextlib import ExitStack from threading import Event, Thread from typing import Any from unittest.mock import ANY, MagicMock, call @@ -23,66 +22,78 @@ import numpy as np import pytest -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray +from dimos.core.cdr_transport import CdrTransport +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 Image +from dimos.messages.std_msgs.msg import Header +from dimos.messages.tf2_msgs.msg import TFMessage +from dimos.messages.vision_msgs.msg import Detection3DArray +from dimos.msgs.camera_info import camera_info_from_intrinsics +from dimos.msgs.image import image_from_array +from dimos.msgs.pointcloud import pointcloud_xyz +from dimos.msgs.time import time_from_seconds +from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox from dimos.perception.detection.type.detection2d.imageDetections2D import ImageDetections2D from dimos.perception.experimental.object_scene_registration import ObjectSceneRegistrationModule from dimos.perception.experimental.objectDB import ObjectDB +from dimos.perception.experimental.test_object_generated import _object -class _FakeTF: - def __init__(self, result: Any) -> None: - self.result = result - self.calls: list[tuple[Any, ...]] = [] +def _image(timestamp: float) -> Image: + return image_from_array( + np.ones((2, 2), dtype=np.float32), + encoding="32FC1", + header=Header(frame_id="camera", stamp=time_from_seconds(timestamp)), + ) - def get(self, *args: Any, **kwargs: Any) -> Any: - self.calls.append((args, kwargs)) - return self.result - def dispose(self) -> None: - pass +@pytest.fixture +def module_factory(): + with ExitStack() as cleanup: + def create(**kwargs): + module = ObjectSceneRegistrationModule(**kwargs) + cleanup.callback(module.stop) + module._camera_info = camera_info_from_intrinsics( + 1, 1, 0, 0, 2, 2, header=Header(frame_id="camera", stamp=Time(sec=0, nanosec=0)) + ) + wire = CdrTransport("object_scene_tf", TFMessage, backend="lcm", url="memq://") + cleanup.callback(wire.stop) + module.tf.transport = wire + return module -def _image(timestamp: float) -> Image: - return Image( - data=np.ones((2, 2), dtype=np.float32), - format=ImageFormat.DEPTH, - frame_id="camera", - ts=timestamp, - ) + yield create @pytest.fixture -def module() -> Iterator[ObjectSceneRegistrationModule]: - module = ObjectSceneRegistrationModule(target_frame="map") - module._camera_info = MagicMock(K=[1.0, 0.0, 0.0, 0.0, 1.0, 0.0]) - module._latest_scene_snapshot = None - yield module - module.stop() +def module(module_factory): + return module_factory(target_frame="map") def test_transform_is_captured_before_slow_detector( - monkeypatch: Any, module: ObjectSceneRegistrationModule + monkeypatch: Any, module: ObjectSceneRegistrationModule, module_factory, mocker ) -> None: now = [0.0] - transform = MagicMock(name="fresh_transform") - - class _ExpiringTF(_FakeTF): - def get(self, *args: Any, **kwargs: Any) -> Any: - super().get(*args, **kwargs) - return transform if now[0] <= 10.0 else None - - tf = _ExpiringTF(None) - module._tf = tf # type: ignore[assignment] + transform = TransformStamped( + header=Header(frame_id="map", stamp=time_from_seconds(12.5)), + child_frame_id="camera", + transform=Transform( + translation=Vector3(x=1, y=0, z=0), rotation=Quaternion(x=0, y=0, z=0, w=1) + ), + ) + lookup = mocker.patch.object( + module.tfbuffer, + "get", + side_effect=lambda *args, **kwargs: transform if now[0] <= 10 else None, + ) module._detector_backend = "owlv2" module._detector = MagicMock() module._text_prompts = ["cup"] - module.detections_2d = MagicMock() - color = Image( - data=np.zeros((2, 2, 3), dtype=np.uint8), - format=ImageFormat.BGR, - frame_id="camera", - ts=12.5, + color = image_from_array( + np.zeros((2, 2, 3), dtype=np.uint8), + encoding="bgr8", + header=Header(frame_id="camera", stamp=time_from_seconds(12.5)), ) detections = ImageDetections2D(color, []) @@ -96,14 +107,20 @@ def slow_detection(*_args: Any, **_kwargs: Any) -> ImageDetections2D: module._process_images(color, _image(12.5)) - assert tf.calls == [(("map", "camera", 12.5, 0.1), {"forward_tolerance": 0.2})] + lookup.assert_called_once_with("map", "camera", 12.5, 0.1, forward_tolerance=0.2) process_3d.assert_called_once_with(detections, color, ANY, transform) def test_failed_lookup_does_not_retry_without_time_or_replace_coherent_cache( - monkeypatch: Any, module: ObjectSceneRegistrationModule + monkeypatch: Any, module: ObjectSceneRegistrationModule, module_factory, mocker ) -> None: - old_transform = MagicMock(name="old_transform") + old_transform = TransformStamped( + header=Header(frame_id="map", stamp=time_from_seconds(1)), + child_frame_id="camera", + transform=Transform( + translation=Vector3(x=0, y=0, z=0), rotation=Quaternion(x=0, y=0, z=0, w=1) + ), + ) warning = MagicMock() monkeypatch.setattr( "dimos.perception.experimental.object_scene_registration.logger.warning", warning @@ -116,7 +133,7 @@ def test_failed_lookup_does_not_retry_without_time_or_replace_coherent_cache( old_depth = _image(1.0) ObjectSceneRegistrationModule._process_3d_detections( module, - MagicMock(spec=ImageDetections2D), + ImageDetections2D(_image(1), []), old_depth, old_depth, old_transform, @@ -124,7 +141,7 @@ def test_failed_lookup_does_not_retry_without_time_or_replace_coherent_cache( new_depth = _image(2.0) ObjectSceneRegistrationModule._process_3d_detections( module, - MagicMock(spec=ImageDetections2D), + ImageDetections2D(_image(1), []), new_depth, new_depth, None, @@ -135,53 +152,44 @@ def test_failed_lookup_does_not_retry_without_time_or_replace_coherent_cache( def test_full_scene_pointcloud_uses_one_coherent_scene_snapshot( - monkeypatch: Any, module: ObjectSceneRegistrationModule + monkeypatch: Any, module: ObjectSceneRegistrationModule, module_factory, mocker ) -> None: depth = _image(3.0) - transform = MagicMock(name="transform") - module._tf = _FakeTF(transform) # type: ignore[assignment] - module._latest_scene_snapshot = (depth, transform) - - class _PointCloud: - points = list(range(100)) - - def voxel_down_sample(self, voxel_size: float) -> _PointCloud: - return self - - pointcloud = _PointCloud() - fake_o3d = MagicMock() - fake_o3d.camera.PinholeCameraIntrinsic.return_value = MagicMock() - fake_o3d.geometry.Image.return_value = MagicMock() - fake_o3d.geometry.PointCloud.create_from_depth_image.return_value = pointcloud - # open3d is imported inside the method under test, so swap the module itself - monkeypatch.setitem(sys.modules, "open3d", fake_o3d) - - result = MagicMock() - result.transform.side_effect = lambda used_transform: ( - result if used_transform is transform else pytest.fail("mixed scene snapshot") + transform = TransformStamped( + header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)), + child_frame_id="camera", + transform=Transform( + translation=Vector3(x=1, y=0.0, z=0.0), rotation=Quaternion(w=1, x=0.0, y=0.0, z=0.0) + ), ) - monkeypatch.setattr( - "dimos.perception.experimental.object_scene_registration.PointCloud2", - lambda *_args, **_kwargs: result, + depth = image_from_array( + np.ones((20, 20), dtype=np.float32), encoding="32FC1", header=depth.header ) - - module.get_full_scene_pointcloud() - result.transform.assert_called_once_with(transform) - - -def test_owlv2_queries_configured_prompts(monkeypatch: Any) -> None: - module = ObjectSceneRegistrationModule( + module._camera_info = camera_info_from_intrinsics(100, 100, 10, 10, 20, 20, header=depth.header) + module._latest_scene_snapshot = (depth, transform) + lookup = mocker.spy(module.tfbuffer, "get") + result = module.get_full_scene_pointcloud(voxel_size=0.001) + assert result is not None + assert result.header.frame_id == "map" + assert result.header.stamp == depth.header.stamp + points = pointcloud_xyz(result) + assert len(points) == 400 + np.testing.assert_allclose(points[:, 2], 1) + assert points[:, 0].min() == pytest.approx(0.9) + assert points[:, 0].max() == pytest.approx(1.09) + lookup.assert_not_called() + + +def test_owlv2_queries_configured_prompts(monkeypatch: Any, module_factory, mocker) -> None: + module = module_factory( target_frame="camera", detector_backend="owlv2", detector_confidence=0.07 ) - module._camera_info = MagicMock() module._detector = MagicMock() module._text_prompts = ["mug"] - module.detections_2d = MagicMock() - color = Image( - data=np.zeros((2, 2, 3), dtype=np.uint8), - format=ImageFormat.BGR, - frame_id="camera", - ts=4.0, + color = image_from_array( + np.zeros((2, 2, 3), dtype=np.uint8), + encoding="bgr8", + header=Header(frame_id="camera", stamp=time_from_seconds(4.0)), ) detections = ImageDetections2D(color, []) module._detector.query_detections.return_value = detections @@ -195,8 +203,8 @@ def test_owlv2_queries_configured_prompts(monkeypatch: Any) -> None: module.stop() -def test_owlv2_yolo_constructs_box_prompt_segmenter(mocker: Any) -> None: - module = ObjectSceneRegistrationModule( +def test_owlv2_yolo_constructs_box_prompt_segmenter(mocker: Any, module_factory) -> None: + module = module_factory( detector_backend="owlv2", segmentation_backend="yolo", segmentation_confidence=0.04, @@ -216,30 +224,31 @@ def test_owlv2_yolo_constructs_box_prompt_segmenter(mocker: Any) -> None: segmenter_factory.assert_called_once_with(confidence=0.04) -def test_moondream_queries_each_configured_prompt(monkeypatch: Any) -> None: - module = ObjectSceneRegistrationModule(target_frame="camera", detector_backend="moondream") - module._camera_info = MagicMock() - module._detector = MagicMock() +def test_moondream_queries_each_configured_prompt(monkeypatch: Any, module_factory, mocker) -> None: + module = module_factory(target_frame="camera", detector_backend="moondream") + detector = MagicMock() + module._detector = detector module._text_prompts = ["cup", "bottle"] - module.detections_2d = MagicMock() - color = Image( - data=np.zeros((2, 2, 3), dtype=np.uint8), - format=ImageFormat.BGR, - frame_id="camera", - ts=4.0, + color = image_from_array( + np.zeros((2, 2, 3), dtype=np.uint8), + encoding="bgr8", + header=Header(frame_id="camera", stamp=time_from_seconds(4.0)), ) - cup = MagicMock(track_id=0, class_id=-1) - bottle = MagicMock(track_id=0, class_id=-1) - module._detector.query_detections.side_effect = [ + cup = Detection2DBBox((0, 0, 1, 1), 0, -1, 0.9, "cup", 4.0, color) + bottle = Detection2DBBox((1, 1, 2, 2), 0, -1, 0.8, "bottle", 4.0, color) + detector.query_detections.side_effect = [ ImageDetections2D(color, [cup]), ImageDetections2D(color, [bottle]), ] process_3d = MagicMock() monkeypatch.setattr(module, "_process_3d_detections", process_3d) - module._process_images(color, _image(4.0)) + try: + module._process_images(color, _image(4.0)) + finally: + module.stop() - assert module._detector.query_detections.call_args_list == [ + assert detector.query_detections.call_args_list == [ call(color, "cup"), call(color, "bottle"), ] @@ -247,17 +256,14 @@ def test_moondream_queries_each_configured_prompt(monkeypatch: Any) -> None: assert combined.detections == [cup, bottle] assert (cup.track_id, cup.class_id) == (-1, 0) assert (bottle.track_id, bottle.class_id) == (-1, 1) - module.stop() -def test_edgetam_refines_detector_output(monkeypatch: Any) -> None: - module = ObjectSceneRegistrationModule(target_frame="camera", segmentation_backend="edgetam") - module._camera_info = MagicMock() - color = Image( - data=np.zeros((2, 2, 3), dtype=np.uint8), - format=ImageFormat.BGR, - frame_id="camera", - ts=4.0, +def test_edgetam_refines_detector_output(monkeypatch: Any, module_factory, mocker) -> None: + module = module_factory(target_frame="camera", segmentation_backend="edgetam") + color = image_from_array( + np.zeros((2, 2, 3), dtype=np.uint8), + encoding="bgr8", + header=Header(frame_id="camera", stamp=time_from_seconds(4.0)), ) raw_detections = ImageDetections2D(color, []) segmented_detections = ImageDetections2D(color, []) @@ -265,7 +271,6 @@ def test_edgetam_refines_detector_output(monkeypatch: Any) -> None: module._segmenter.segment.return_value = segmented_detections module._detector = MagicMock() module._detector.process_image.return_value = raw_detections - module.detections_2d = MagicMock() process_3d = MagicMock() monkeypatch.setattr(module, "_process_3d_detections", process_3d) @@ -276,21 +281,22 @@ def test_edgetam_refines_detector_output(monkeypatch: Any) -> None: module.stop() -def test_request_driven_scan_processes_latest_aligned_frame(monkeypatch: Any) -> None: - module = ObjectSceneRegistrationModule( - target_frame="camera", detector_backend="owlv2", detect_on_request=True - ) +def test_request_driven_scan_processes_latest_aligned_frame( + monkeypatch: Any, module_factory, mocker +) -> None: + module = module_factory(target_frame="camera", detector_backend="owlv2", detect_on_request=True) module._object_db = MagicMock() - color = Image( - data=np.zeros((2, 2, 3), dtype=np.uint8), - format=ImageFormat.BGR, - frame_id="camera", - ts=4.0, + color = image_from_array( + np.zeros((2, 2, 3), dtype=np.uint8), + encoding="bgr8", + header=Header(frame_id="camera", stamp=time_from_seconds(4.0)), ) depth = _image(4.0) module._on_aligned_frames((color, depth)) output = MagicMock() - result = MagicMock(spec=Detection3DArray) + result = Detection3DArray( + header=Header(frame_id="camera", stamp=time_from_seconds(4)), detections=[] + ) def process_images(got_color: Image, got_depth: Image) -> list[MagicMock]: assert (got_color, got_depth) == (color, depth) @@ -307,8 +313,8 @@ def process_images(got_color: Image, got_depth: Image) -> list[MagicMock]: module.stop() -def test_request_driven_detect_triggers_scan(monkeypatch: Any) -> None: - module = ObjectSceneRegistrationModule(detector_backend="owlv2", detect_on_request=True) +def test_request_driven_detect_triggers_scan(monkeypatch: Any, module_factory, mocker) -> None: + module = module_factory(detector_backend="owlv2", detect_on_request=True) module._detector = MagicMock() current_object = MagicMock() current_object.agent_encode.return_value = {"name": "mug", "object_id": "current"} @@ -322,18 +328,14 @@ def test_request_driven_detect_triggers_scan(monkeypatch: Any) -> None: module.stop() -def test_scan_output_includes_pending_objects(monkeypatch: Any) -> None: - module = ObjectSceneRegistrationModule(target_frame="camera") - module._camera_info = MagicMock() +def test_scan_output_includes_pending_objects(monkeypatch: Any, module_factory, mocker) -> None: + module = module_factory(target_frame="camera") module._object_db = MagicMock() pending = MagicMock() permanent = MagicMock() module._object_db.get_all_objects.return_value = [pending, permanent] module._object_db.get_objects.return_value = [permanent] module._object_db.add_objects.return_value = [pending] - module.detections_3d = MagicMock() - module.objects = MagicMock() - module.pointcloud = MagicMock() converted_objects = [pending] monkeypatch.setattr( "dimos.perception.experimental.object_scene_registration.Object.from_2d_to_list", @@ -350,7 +352,7 @@ def test_scan_output_includes_pending_objects(monkeypatch: Any) -> None: observed = ObjectSceneRegistrationModule._process_3d_detections( module, - MagicMock(spec=ImageDetections2D), + ImageDetections2D(_image(1), []), _image(4.0), _image(4.0), None, @@ -362,7 +364,7 @@ def test_scan_output_includes_pending_objects(monkeypatch: Any) -> None: module._object_db.add_objects.reset_mock() observed = ObjectSceneRegistrationModule._process_3d_detections( module, - MagicMock(spec=ImageDetections2D), + ImageDetections2D(_image(1), []), _image(5.0), _image(5.0), None, @@ -373,16 +375,13 @@ def test_scan_output_includes_pending_objects(monkeypatch: Any) -> None: module.stop() -def test_concurrent_request_scans_do_not_overlap(monkeypatch: Any) -> None: - module = ObjectSceneRegistrationModule( - target_frame="camera", detector_backend="owlv2", detect_on_request=True - ) +def test_concurrent_request_scans_do_not_overlap(monkeypatch: Any, module_factory, mocker) -> None: + module = module_factory(target_frame="camera", detector_backend="owlv2", detect_on_request=True) module._object_db = MagicMock() - color = Image( - data=np.zeros((2, 2, 3), dtype=np.uint8), - format=ImageFormat.BGR, - frame_id="camera", - ts=4.0, + color = image_from_array( + np.zeros((2, 2, 3), dtype=np.uint8), + encoding="bgr8", + header=Header(frame_id="camera", stamp=time_from_seconds(4.0)), ) module._on_aligned_frames((color, _image(4.0))) first_started = Event() @@ -407,12 +406,15 @@ def process_images(_: Image, __: Image) -> list[MagicMock]: first = Thread(target=lambda: module.scan_scene(text=["cup"])) second = Thread(target=lambda: module.scan_scene(text=["mug"])) first.start() - assert first_started.wait(timeout=1.0) - second.start() - assert not second_started.wait(timeout=0.05) - release_first.set() - first.join(timeout=1.0) - second.join(timeout=1.0) + try: + assert first_started.wait(timeout=1.0) + second.start() + assert not second_started.wait(timeout=0.05) + finally: + release_first.set() + first.join(timeout=1.0) + if second.ident is not None: + second.join(timeout=1.0) assert not first.is_alive() assert not second.is_alive() @@ -420,104 +422,53 @@ def process_images(_: Image, __: Image) -> list[MagicMock]: module.stop() -def test_object_db_uses_wall_clock_for_pending_ttl(monkeypatch: Any) -> None: - object_db = ObjectDB(pending_ttl_s=5.0) - detected = MagicMock() - detected.object_id = "stable-id" - detected.track_id = -1 - detected.ts = 4.0 # Hardware timestamps are relative to camera boot. - detected.last_seen_ts = None - detected.center = None - detected.detections_count = 1 - now = [1000.0] +def test_object_db_uses_wall_clock_for_pending_ttl(monkeypatch): + database = ObjectDB(pending_ttl_s=5) + detected = _object("stable-id", 0, 4) + now = [1000] monkeypatch.setattr("dimos.perception.experimental.objectDB.time.time", lambda: now[0]) - - object_db.add_objects([detected]) - assert detected.last_seen_ts == 1000.0 - now[0] = 1004.0 - object_db.add_objects([]) - assert object_db.find_by_object_id("stable-id") is detected - now[0] = 1006.0 - object_db.add_objects([]) - assert object_db.find_by_object_id("stable-id") is None - - -def test_object_db_counts_each_source_frame_once(monkeypatch: Any) -> None: - object_db = ObjectDB(min_detections_for_permanent=10) - now = [1000.0] + database.add_objects([detected]) + assert detected.last_seen_ts == 1000 + now[0] = 1004 + database.add_objects([]) + assert database.find_by_object_id("stable-id") is detected + now[0] = 1006 + database.add_objects([]) + assert database.find_by_object_id("stable-id") is None + + +def test_object_db_counts_each_source_frame_once(monkeypatch): + database = ObjectDB(min_detections_for_permanent=10) + now = [1000] monkeypatch.setattr("dimos.perception.experimental.objectDB.time.time", lambda: now[0]) - - first = MagicMock( - object_id="first-id", - track_id=-1, - ts=4.0, - last_seen_ts=None, - detections_count=1, - ) - first.center = MagicMock() - duplicate = MagicMock(object_id="duplicate-id", track_id=-1, ts=4.0) - duplicate.center = MagicMock() - duplicate.center.distance.return_value = 0.0 - - observed = object_db.add_objects([first, duplicate]) - - assert observed == [first] - first.update_object.assert_not_called() - assert first.last_seen_ts == 1000.0 - - newer = MagicMock(object_id="newer-id", track_id=-1, ts=5.0) - newer.center = MagicMock() - newer.center.distance.return_value = 0.0 - first.update_object.side_effect = lambda _: setattr(first, "detections_count", 2) - now[0] = 1001.0 - - assert object_db.add_objects([newer]) == [first] - first.update_object.assert_called_once_with(newer) - assert first.last_seen_ts == 1001.0 - - -def test_object_db_promotes_a_first_sighting_when_threshold_is_one() -> None: - object_db = ObjectDB(min_detections_for_permanent=1) - detected = MagicMock( - object_id="first-id", - track_id=-1, - ts=1.0, - last_seen_ts=None, - detections_count=1, - name="cup", - ) - detected.center = None - - object_db.add_objects([detected]) - - assert object_db.get_objects() == [detected] - assert object_db.get_stats() == { - "pending_count": 0, - "permanent_count": 1, - "total_count": 1, - } - - -def test_object_db_keeps_first_sighting_pending_above_threshold() -> None: - object_db = ObjectDB(min_detections_for_permanent=2) - detected = MagicMock( - object_id="first-id", - track_id=7, - ts=1.0, - last_seen_ts=None, - detections_count=1, - name="cup", - ) - detected.center = None - newer = MagicMock(object_id="newer-id", track_id=7, ts=2.0) - newer.center = None - detected.update_object.side_effect = lambda _: setattr(detected, "detections_count", 2) - - object_db.add_objects([detected]) - - assert object_db.get_objects() == [] - assert object_db.find_by_object_id("first-id") is detected - - object_db.add_objects([newer]) - - assert object_db.get_objects() == [detected] + first = _object("first-id", 0, 4) + duplicate = _object("duplicate-id", 0, 4) + assert database.add_objects([first, duplicate]) == [first] + assert first.detections_count == 1 + assert first.last_seen_ts == 1000 + newer = _object("newer-id", 0, 5) + now[0] = 1001 + assert database.add_objects([newer]) == [first] + assert first.detections_count == 2 + assert first.last_seen_ts == 1001 + + +def test_object_db_promotes_a_first_sighting_when_threshold_is_one(): + database = ObjectDB(min_detections_for_permanent=1) + detected = _object("first-id", 0, 1) + database.add_objects([detected]) + assert database.get_objects() == [detected] + assert database.get_stats() == {"pending_count": 0, "permanent_count": 1, "total_count": 1} + + +def test_object_db_keeps_first_sighting_pending_above_threshold(): + database = ObjectDB(min_detections_for_permanent=2) + detected = _object("first-id", 0, 1) + detected.track_id = 7 + newer = _object("newer-id", 0, 2) + newer.track_id = 7 + database.add_objects([detected]) + assert database.get_objects() == [] + assert database.find_by_object_id("first-id") is detected + database.add_objects([newer]) + assert database.get_objects() == [detected] diff --git a/dimos/perception/experimental/test_spatial_memory.py b/dimos/perception/experimental/test_spatial_memory.py index 211644847cf..53cc7f266ac 100644 --- a/dimos/perception/experimental/test_spatial_memory.py +++ b/dimos/perception/experimental/test_spatial_memory.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from contextlib import ExitStack import os import shutil import tempfile @@ -22,9 +23,15 @@ from reactivex import of, operators as ops from reactivex.scheduler import ThreadPoolScheduler -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.core.cdr_transport import CdrTransport +from dimos.messages.builtin_interfaces.msg import Time +from dimos.messages.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos.messages.std_msgs.msg import Header +from dimos.messages.tf2_msgs.msg import TFMessage +from dimos.msgs.image import image_from_array +from dimos.perception.experimental import spatial_perception from dimos.perception.experimental.spatial_perception import SpatialMemory +from dimos.perception.experimental.visual_memory import VisualMemory from dimos.stream.video_provider import VideoProvider @@ -109,7 +116,7 @@ def process_frame(frame): nonlocal frame_counter # Generate a unique position for this frame to ensure minimum distance threshold is met - pos = Pose(frame_counter * 0.5, frame_counter * 0.5, 0) + pos = Vector3(x=frame_counter * 0.5, y=frame_counter * 0.5, z=0.0) transform = {"position": pos, "timestamp": time.time()} frame_counter += 1 @@ -221,11 +228,63 @@ def flaky_get_embedding(image): spatial_memory.process_stream( of( - {"frame": frame, "position": Vector3(1000, 0, 0), "rotation": Vector3()}, - {"frame": frame, "position": Vector3(1001, 0, 0), "rotation": Vector3()}, + { + "frame": frame, + "position": Vector3(x=1000.0, y=0.0, z=0.0), + "rotation": Vector3(x=0.0, y=0.0, z=0.0), + }, + { + "frame": frame, + "position": Vector3(x=1001.0, y=0.0, z=0.0), + "rotation": Vector3(x=0.0, y=0.0, z=0.0), + }, ) ).subscribe(on_next=results.append, on_error=errors.append) assert errors == [] assert [r["position"] for r in results] == [(1001, 0, 0)] assert spatial_memory.stored_frame_count == stored_before + 1 + + +def test_generated_image_and_tf_reach_spatial_backend_with_correct_colors_and_pose( + tmp_path, mocker +): + embedding = mocker.Mock() + embedding.get_embedding.return_value = np.ones(512, dtype=np.float32) / np.sqrt(512) + database = mocker.Mock() + mocker.patch.object(spatial_perception, "ImageEmbeddingProvider", return_value=embedding) + mocker.patch.object(spatial_perception, "SpatialVectorDB", return_value=database) + with ExitStack() as cleanup: + wire = CdrTransport("spatial_tf", TFMessage, backend="lcm", url="memq://") + cleanup.callback(wire.stop) + memory = SpatialMemory( + chroma_client=mocker.Mock(), + db_path=None, + visual_memory_path=None, + visual_memory=VisualMemory(output_dir=str(tmp_path)), + min_time_threshold=0, + ) + cleanup.callback(memory.stop) + memory.tf.transport = wire + memory.tfbuffer.receive_transform( + TransformStamped( + header=Header(frame_id="world", stamp=Time(sec=1, nanosec=0)), + child_frame_id="base_link", + transform=Transform( + translation=Vector3(x=1, y=2, z=3), rotation=Quaternion(x=0, y=0, z=0, w=1) + ), + ) + ) + pixels = np.zeros((4, 4, 3), dtype=np.uint8) + pixels[:, :, 0] = 255 + image = image_from_array( + pixels, encoding="rgb8", header=Header(frame_id="camera", stamp=Time(sec=0, nanosec=1)) + ) + memory._receive_frame(image) + memory._process_frame() + database.add_image_vector.assert_called_once() + stored = database.add_image_vector.call_args.kwargs + np.testing.assert_array_equal(stored["image"][0, 0], [0, 0, 255]) + assert [stored["metadata"][key] for key in ["pos_x", "pos_y", "pos_z"]] == [1, 2, 3] + assert [stored["metadata"][key] for key in ["rot_x", "rot_y", "rot_z"]] == [0, 0, 0] + assert memory.stored_frame_count == 1 diff --git a/docs/capabilities/perception/index.md b/docs/capabilities/perception/index.md index 7fa9fef98cb..adcd7e3d468 100644 --- a/docs/capabilities/perception/index.md +++ b/docs/capabilities/perception/index.md @@ -158,6 +158,21 @@ converts that value to a generated Detection3DArray. Live module/blueprint execution requires the remaining default transport cutover. Historical recordings still require explicit conversion; this slice does not add a legacy decoder. +## Experimental object memory + +The experimental object, scene-registration and tracker modules now use generated +image, calibration, detection and TF values on their typed ports. EdgeTAM and +YOLO-E segmentation retain their existing model APIs and use the same image +helpers. OpenCV tracking keeps its existing geometry and processing-time policy. + +Object geometry keeps the color source stamp on its pose/detection and the depth +source stamp on its projected cloud. Scene scans preserve their capture Header; +spatial deduplication compares source frame and integer nanoseconds rather than +rounded application seconds. Wall-clock pending-object TTL remains separate, and +source timestamp zero is not replaced by wall time. Until default encoding +cutover, these generated ports require the explicit CDR transports described +above. Offline boundary tests do not constitute full model or robot acceptance. + ## Experimental WorldBelief The experimental xArm6 WorldBelief stack records RGB-D observations and processes