diff --git a/.github/workflows/message-codegen.yml b/.github/workflows/message-codegen.yml index 7c1dd07e4b9..6d9ac01c383 100644 --- a/.github/workflows/message-codegen.yml +++ b/.github/workflows/message-codegen.yml @@ -132,7 +132,7 @@ jobs: if: runner.os == 'Linux' run: | sudo apt-get update - sudo apt-get install -y liblcm-dev + sudo apt-get install -y liblcm-dev libportaudio2 mkdir -p build/sdk-deps curl --fail --location --retry 3 https://github.com/eclipse-zenoh/zenoh-c/releases/download/1.10.1/zenoh-c-1.10.1-x86_64-unknown-linux-gnu-standalone.zip -o build/sdk-deps/zenoh-c.zip echo '9ee0f2d732b0f3042a7e1cd3076042a2bc3ac0415587c40bc3ed7b8b62fbde11 build/sdk-deps/zenoh-c.zip' | sha256sum --check @@ -153,6 +153,13 @@ jobs: .venv/bin/python -m pytest --noconftest --import-mode=importlib -o addopts='' -o 'markers=native_e2e: compiled native consumer acceptance' -q dimos/message_codegen/test_project_native.py -k actual_coordinator .native-review/bin/python -I -c 'from dimos.memory import module; from dimos.memory.store import mcap; import pytest; import sys; assert all("site-packages" in str(value.__file__) for value in (module,mcap)); raise SystemExit(pytest.main(sys.argv[1:]))' --noconftest --import-mode=importlib -o addopts='' -q "$PWD/dimos/memory/test_module.py" "$PWD/dimos/memory/store/test_cdr_recording.py" "$PWD/dimos/memory/store/test_mcap.py" cargo test --manifest-path native/rust/dimos-module/Cargo.toml --locked --features cdr --lib --jobs 2 + - name: Verify Go2 CDR boundaries with stubbed drivers and installed code + if: runner.os == 'Linux' + run: | + # Import-only SDK subset, pinned to uv.lock. Audio/hardware APIs are never called. + uv pip install --python .native-review/bin/python 'aiortc==1.14.0' + uv pip install --python .native-review/bin/python --no-deps 'unitree-webrtc-connect==2.1.2' 'wasmtime==41.0.0' 'pycryptodome==3.23.0' 'curl-cffi==0.15.0' 'sounddevice==0.5.5' + LCM_DEFAULT_URL=memq:// .native-review/bin/python -I -c 'from dimos.robot.unitree.go2 import connection; from dimos.robot.unitree.go2.blueprints.basic import unitree_go2_basic; import pytest; import sys; assert all("site-packages" in str(value.__file__) for value in (connection,unitree_go2_basic)); raise SystemExit(pytest.main(sys.argv[1:]))' --noconftest --import-mode=importlib -o addopts='' -q "$PWD/dimos/robot/unitree/go2/test_connection.py" "$PWD/dimos/robot/unitree/test_connection.py" "$PWD/dimos/robot/unitree/go2/blueprints/basic/test_unitree_go2_basic.py" - name: Verify generated browser schemas, relay encoding and web consumers if: runner.os == 'Linux' run: | diff --git a/dimos/models/embedding/clip.py b/dimos/models/embedding/clip.py index ad5e3e0cc6c..f429dc62a93 100644 --- a/dimos/models/embedding/clip.py +++ b/dimos/models/embedding/clip.py @@ -17,12 +17,12 @@ from functools import cached_property from typing import overload -from dimos.messages.sensor_msgs.msg import Image from PIL import Image as PILImage import torch import torch.nn.functional as functional from transformers import CLIPModel as HFCLIPModel, CLIPProcessor +from dimos.messages.sensor_msgs.msg import Image from dimos.models.base import HuggingFaceModel from dimos.models.embedding.base import Embedding, EmbeddingModel, HuggingFaceEmbeddingModelConfig from dimos.msgs.image import image_to_rgb diff --git a/dimos/models/embedding/dino.py b/dimos/models/embedding/dino.py index 5b0d24fd137..022a442ebf7 100644 --- a/dimos/models/embedding/dino.py +++ b/dimos/models/embedding/dino.py @@ -17,12 +17,12 @@ from functools import cached_property from typing import Any, overload -from dimos.messages.sensor_msgs.msg import Image from PIL import Image as PILImage import torch import torch.nn.functional as functional from transformers import AutoImageProcessor, AutoModel +from dimos.messages.sensor_msgs.msg import Image from dimos.models.base import HuggingFaceModel from dimos.models.embedding.base import Embedding, EmbeddingModel, HuggingFaceEmbeddingModelConfig from dimos.msgs.image import image_to_rgb diff --git a/dimos/models/embedding/siglip.py b/dimos/models/embedding/siglip.py index bb0ed762943..8cabb1afdf4 100644 --- a/dimos/models/embedding/siglip.py +++ b/dimos/models/embedding/siglip.py @@ -17,12 +17,12 @@ from functools import cached_property from typing import overload -from dimos.messages.sensor_msgs.msg import Image from PIL import Image as PILImage import torch import torch.nn.functional as functional from transformers import SiglipModel as HFSiglipModel, SiglipProcessor +from dimos.messages.sensor_msgs.msg import Image from dimos.models.base import HuggingFaceModel from dimos.models.embedding.base import Embedding, EmbeddingModel, HuggingFaceEmbeddingModelConfig from dimos.msgs.image import image_to_rgb diff --git a/dimos/models/embedding/siglip2.py b/dimos/models/embedding/siglip2.py index f8f1fe8861b..c85932f9176 100644 --- a/dimos/models/embedding/siglip2.py +++ b/dimos/models/embedding/siglip2.py @@ -17,12 +17,12 @@ from functools import cached_property from typing import overload -from dimos.messages.sensor_msgs.msg import Image from PIL import Image as PILImage import torch import torch.nn.functional as functional from transformers import SiglipModel as HFSiglipModel, SiglipProcessor +from dimos.messages.sensor_msgs.msg import Image from dimos.models.base import HuggingFaceModel from dimos.models.embedding.base import Embedding, EmbeddingModel, HuggingFaceEmbeddingModelConfig from dimos.msgs.image import image_to_rgb diff --git a/dimos/models/embedding/test_generated_images.py b/dimos/models/embedding/test_generated_images.py index e3c1245ab07..a51f9778ba4 100644 --- a/dimos/models/embedding/test_generated_images.py +++ b/dimos/models/embedding/test_generated_images.py @@ -29,20 +29,25 @@ from dimos.msgs.image import image_from_array, image_view -@pytest.mark.parametrize("model_type", [CLIPModel,DINOModel,SigLIPModel,SigLIP2Model]) -def test_generated_embeddings_keep_zero_source_time_and_convert_rgb(mocker,model_type): - processor,network=mocker.Mock(),mocker.Mock() - features=torch.tensor([[3.0,4.0]]) - processor.return_value=BatchFeature(data={"pixel_values":torch.zeros(1,3,4,8)}) - network.get_image_features.return_value=features - network.return_value=BaseModelOutputWithPooling(last_hidden_state=features[:,None,:],pooler_output=features) - mocker.patch.object(model_type,"_processor",new_callable=PropertyMock,return_value=processor) - mocker.patch.object(model_type,"_model",new_callable=PropertyMock,return_value=network) - image=image_from_array(np.full((4,8,3),[0,0,255],dtype=np.uint8),encoding="bgr8", - header=Header(frame_id="camera",stamp=Time(sec=0,nanosec=0))) - model=model_type(device="cpu") - result=model.embed(image) +@pytest.mark.parametrize("model_type", [CLIPModel, DINOModel, SigLIPModel, SigLIP2Model]) +def test_generated_embeddings_keep_zero_source_time_and_convert_rgb(mocker, model_type): + processor, network = mocker.Mock(), mocker.Mock() + features = torch.tensor([[3.0, 4.0]]) + processor.return_value = BatchFeature(data={"pixel_values": torch.zeros(1, 3, 4, 8)}) + network.get_image_features.return_value = features + network.return_value = BaseModelOutputWithPooling( + last_hidden_state=features[:, None, :], pooler_output=features + ) + mocker.patch.object(model_type, "_processor", new_callable=PropertyMock, return_value=processor) + mocker.patch.object(model_type, "_model", new_callable=PropertyMock, return_value=network) + image = image_from_array( + np.full((4, 8, 3), [0, 0, 255], dtype=np.uint8), + encoding="bgr8", + header=Header(frame_id="camera", stamp=Time(sec=0, nanosec=0)), + ) + model = model_type(device="cpu") + result = model.embed(image) assert result.timestamp == 0.0 - torch.testing.assert_close(result.vector,torch.tensor([0.6,0.8])) - assert processor.call_args.kwargs["images"][0].getpixel((0,0)) == (255,0,0) - np.testing.assert_array_equal(image_view(image)[0,0],[0,0,255]) + torch.testing.assert_close(result.vector, torch.tensor([0.6, 0.8])) + assert processor.call_args.kwargs["images"][0].getpixel((0, 0)) == (255, 0, 0) + np.testing.assert_array_equal(image_view(image)[0, 0], [0, 0, 255]) diff --git a/dimos/models/embedding/treid.py b/dimos/models/embedding/treid.py index ce36b5d004b..3ce7331eec3 100644 --- a/dimos/models/embedding/treid.py +++ b/dimos/models/embedding/treid.py @@ -19,11 +19,11 @@ from functools import cached_property -from dimos.messages.sensor_msgs.msg import Image import torch import torch.nn.functional as functional from torchreid import utils as torchreid_utils +from dimos.messages.sensor_msgs.msg import Image from dimos.models.base import LocalModel from dimos.models.embedding.base import Embedding, EmbeddingModel, EmbeddingModelConfig from dimos.msgs.image import image_to_bgr diff --git a/dimos/models/segmentation/edge_tam.py b/dimos/models/segmentation/edge_tam.py index e240a6529d1..b65977283c0 100644 --- a/dimos/models/segmentation/edge_tam.py +++ b/dimos/models/segmentation/edge_tam.py @@ -20,7 +20,6 @@ import tempfile from typing import TYPE_CHECKING, Any, Protocol, TypedDict -from dimos.messages.sensor_msgs.msg import Image from hydra.utils import instantiate import numpy as np from numpy.typing import NDArray @@ -28,6 +27,7 @@ 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.image import image_to_bgr, image_to_jpeg from dimos.perception.detection.detectors.base import Detector diff --git a/dimos/models/segmentation/test_yoloe.py b/dimos/models/segmentation/test_yoloe.py index 345fc79ac1c..5d2efaffa50 100644 --- a/dimos/models/segmentation/test_yoloe.py +++ b/dimos/models/segmentation/test_yoloe.py @@ -16,11 +16,11 @@ from typing import Any from unittest.mock import MagicMock -from dimos.messages.sensor_msgs.msg import Image -from dimos.messages.std_msgs.msg import Header 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.image import image_from_array, image_to_bgr from dimos.msgs.time import time_from_seconds, to_seconds diff --git a/dimos/models/segmentation/yoloe.py b/dimos/models/segmentation/yoloe.py index d19ecc15701..fd075227d27 100644 --- a/dimos/models/segmentation/yoloe.py +++ b/dimos/models/segmentation/yoloe.py @@ -19,11 +19,11 @@ from dataclasses import replace from typing import Protocol -from dimos.messages.sensor_msgs.msg import Image import numpy as np from numpy.typing import NDArray from ultralytics.models.yolo.yoloe import YOLOEVPSegPredictor +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 diff --git a/dimos/robot/unitree/go2/blueprints/basic/test_unitree_go2_basic.py b/dimos/robot/unitree/go2/blueprints/basic/test_unitree_go2_basic.py new file mode 100644 index 00000000000..7ef3cd126a4 --- /dev/null +++ b/dimos/robot/unitree/go2/blueprints/basic/test_unitree_go2_basic.py @@ -0,0 +1,61 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Go2 blueprint rendering consumes the same generated values as live ports.""" + +import numpy as np + +from dimos.messages.builtin_interfaces.msg import Time +from dimos.messages.dimos_msgs.msg import LineSegment3D, LineSegments3D +from dimos.messages.geometry_msgs.msg import Point, Pose, PoseArray, PoseStamped, Quaternion +from dimos.messages.std_msgs.msg import Header +from dimos.msgs.occupancy import occupancy_from_array +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.robot.unitree.go2.blueprints.basic import unitree_go2_basic as view +from dimos.robot.unitree.go2.connection import GO2Connection + + +def test_camera_and_pose_overrides_accept_generated_port_values(): + camera = view._convert_camera_info(GO2Connection.camera_info_static) + assert camera[0][0] == "/world/color_image" + assert camera[0][1].resolution.as_arrow_array().to_pylist() == [[1280.0, 720.0]] + pose = Pose(Point(1.0, 2.0, 3.0), Quaternion(0.0, 0.0, 0.0, 1.0)) + entries = dict(view._plot_odom(PoseStamped(Header(Time(0, 0), "world"), pose))) + assert "world/odom/message" in entries + assert "world/odom" in entries + assert entries["plots/odom/x"].scalars.as_arrow_array().to_pylist() == [1.0] + assert entries["plots/odom/y"].scalars.as_arrow_array().to_pylist() == [2.0] + + +def test_map_and_pgo_overrides_preserve_rendered_geometry(): + header = Header(Time(0, 0), "world") + cloud = pointcloud_from_xyz(np.array([[1.0, 2.0, 0.2], [9.0, 8.0, -0.1]]), header=header) + points = view._convert_global_map(cloud).positions.as_arrow_array().to_pylist() + np.testing.assert_allclose(points, [[1.0, 2.0, 0.2]]) + pose = Pose(Point(1.0, 2.0, 3.0), Quaternion(0.0, 0.0, 0.0, 1.0)) + keyframes = dict(view._convert_pgo_keyframes(PoseArray(header, [pose]))) + assert keyframes["world/pgo/keyframes"].positions.as_arrow_array().to_pylist() == [ + [1.0, 2.0, 3.0] + ] + edges = LineSegments3D(header, [LineSegment3D(Point(1.0, 2.0, 3.0), Point(4.0, 5.0, 6.0), 0.9)]) + loops = dict(view._convert_pgo_loops(edges)) + assert loops["world/pgo/loops"].strips.as_arrow_array().to_pylist() == [ + [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]] + ] + grid = occupancy_from_array(np.array([[0, 100]], dtype=np.int8), header=header, resolution=0.5) + mesh = view._convert_navigation_costmap(grid) + np.testing.assert_allclose( + mesh.vertex_positions.as_arrow_array().to_pylist(), + [[0.0, 0.0, 0.015], [1.0, 0.0, 0.015], [1.0, 0.5, 0.015], [0.0, 0.5, 0.015]], + ) diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py index 2a0cbe8830c..8cc3d0638f5 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py @@ -23,7 +23,11 @@ from dimos.core.coordination.blueprints import autoconnect from dimos.core.global_config import global_config +from dimos.msgs.geometry_msgs.Pose import Pose as NumericalPose +from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid as NumericalGrid +from dimos.msgs.occupancy import occupancy_view from dimos.robot.unitree.go2.connection import GO2Connection +from dimos.visualization.rerun.cdr import camera_pinhole, cloud_archetype, generated_archetypes from dimos.visualization.vis_module import vis_module if TYPE_CHECKING: @@ -32,30 +36,32 @@ from rerun._baseclasses import Archetype from rerun.blueprint import Blueprint - from dimos.msgs.geometry_msgs.PoseArray import PoseArray - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped - from dimos.msgs.geometry_msgs.Twist import Twist - from dimos.msgs.nav_msgs.LineSegments3D import LineSegments3D - from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid - from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 - from dimos.msgs.std_msgs.Float32 import Float32 + from dimos.messages.dimos_msgs.msg import LineSegments3D + from dimos.messages.geometry_msgs.msg import PoseArray, PoseStamped, Twist + from dimos.messages.nav_msgs.msg import OccupancyGrid + from dimos.messages.sensor_msgs.msg import CameraInfo, PointCloud2 + from dimos.messages.std_msgs.msg import Float32 from dimos.visualization.rerun.bridge import RerunData, RerunMulti def _convert_camera_info(camera_info: CameraInfo) -> RerunData: - return camera_info.to_rerun( - image_topic="/world/color_image", - optical_frame="camera_optical", - ) + return [("/world/color_image", camera_pinhole(camera_info))] def _convert_global_map(grid: PointCloud2) -> Archetype: - return grid.to_rerun(bottom_cutoff=0) + return cloud_archetype(grid, bottom_cutoff=0) def _convert_navigation_costmap(grid: OccupancyGrid) -> Archetype: - return grid.to_rerun( + # Reuse the existing numerical texture renderer behind an explicit value adapter. + p, q = grid.info.origin.position, grid.info.origin.orientation + numerical = NumericalGrid( + grid=occupancy_view(grid), + resolution=grid.info.resolution, + origin=NumericalPose(p.x, p.y, p.z, q.x, q.y, q.z, q.w), + frame_id=grid.header.frame_id, + ) + return numerical.to_rerun( colormap="Accent", z_offset=0.015, opacity=0.2, @@ -66,7 +72,10 @@ def _convert_navigation_costmap(grid: OccupancyGrid) -> Archetype: def _convert_pgo_keyframes(keyframes: PoseArray) -> RerunMulti: import rerun as rr - positions = keyframes.positions() + positions = np.asarray( + [[pose.position.x, pose.position.y, pose.position.z] for pose in keyframes.poses], + dtype=np.float32, + ).reshape(-1, 3) return [ ("world/pgo/keyframes", rr.Points3D(positions, colors=[[255, 0, 0]], radii=[0.025])), ("world/pgo/path", rr.LineStrips3D([positions], colors=[[255, 255, 255]], radii=[0.01])), @@ -76,7 +85,13 @@ def _convert_pgo_keyframes(keyframes: PoseArray) -> RerunMulti: def _convert_pgo_loops(edges: LineSegments3D) -> RerunMulti: import rerun as rr - segments = edges.segments.astype(np.float32) + segments = np.asarray( + [ + [[edge.start.x, edge.start.y, edge.start.z], [edge.end.x, edge.end.y, edge.end.z]] + for edge in edges.segments + ], + dtype=np.float32, + ).reshape(-1, 2, 3) strips = rr.LineStrips3D(segments, colors=[[231, 76, 60]], radii=[0.008]) return [("world/pgo/loops", strips)] @@ -90,10 +105,12 @@ def _plot_telemetry(name: str, msg: Float32) -> RerunMulti: def _plot_odom(odom: PoseStamped) -> RerunMulti: import rerun as rr + entries = generated_archetypes(odom, "world/odom") + assert entries is not None return [ - ("world/odom", odom.to_rerun()), - ("plots/odom/x", rr.Scalars(odom.x)), - ("plots/odom/y", rr.Scalars(odom.y)), + *entries, + ("plots/odom/x", rr.Scalars(odom.pose.position.x)), + ("plots/odom/y", rr.Scalars(odom.pose.position.y)), ] diff --git a/dimos/robot/unitree/go2/connection.py b/dimos/robot/unitree/go2/connection.py index 6d95f99e75c..a63e87a046c 100644 --- a/dimos/robot/unitree/go2/connection.py +++ b/dimos/robot/unitree/go2/connection.py @@ -20,6 +20,7 @@ import time from typing import Any, Protocol +import numpy as np from pydantic import Field from reactivex import defer, empty, operators as ops from reactivex.abc import SchedulerBase @@ -35,15 +36,35 @@ from dimos.core.stream import In, Out from dimos.memory.replay import Replay, ReplayStream, resolve_db_path from dimos.memory.store.sqlite import SqliteStore +from dimos.messages.geometry_msgs.msg import ( + Point as CdrPoint, + Pose as CdrPose, + PoseStamped as CdrPoseStamped, + Quaternion as CdrQuaternion, + Transform as CdrTransform, + TransformStamped as CdrTransformStamped, + Twist as CdrTwist, + Vector3 as CdrVector3, +) +from dimos.messages.sensor_msgs.msg import ( + CameraInfo as CdrCameraInfo, + Image as CdrImage, + PointCloud2 as CdrCloud, + PointField, + RegionOfInterest, +) +from dimos.messages.std_msgs.msg import Header +from dimos.messages.tf2_msgs.msg import TFMessage 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.Twist import Twist from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.image import image_from_array from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo from dimos.msgs.sensor_msgs.Image import Image from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.time import time_from_seconds from dimos.robot.unitree.connection import UnitreeWebRTCConnection from dimos.robot.unitree.type.lowstate import LowStateMsg from dimos.spec.perception import Camera, Pointcloud @@ -108,11 +129,69 @@ def publish_request(self, topic: str, data: dict) -> dict: ... # type: ignore[t ) +def _camera_info_message(info: CameraInfo) -> CdrCameraInfo: + return CdrCameraInfo( + header=Header(time_from_seconds(info.ts), info.frame_id), + height=info.height, + width=info.width, + distortion_model=info.distortion_model, + d=np.asarray(info.D, dtype=np.float64), + k=np.asarray(info.K, dtype=np.float64), + r=np.asarray(info.R, dtype=np.float64), + p=np.asarray(info.P, dtype=np.float64), + binning_x=info.binning_x, + binning_y=info.binning_y, + roi=RegionOfInterest(0, 0, 0, 0, False), + ) + + def _camera_info_static() -> CameraInfo: with resources.as_file(_FRONT_CAMERA_720_YAML) as yaml_path: return CameraInfo.from_yaml(str(yaml_path)) +def _cloud_message(cloud: PointCloud2) -> CdrCloud: + """Adapt the driver's numerical cloud, preserving every supported point attribute.""" + points, colors = cloud.as_numpy() + arrays = {name: points[:, index] for index, name in enumerate(("x", "y", "z"))} + types = {name: (" str: """Apply a TF namespace prefix (ModuleConfig.frame_id_prefix) to a frame name.""" if not prefix or not name: @@ -343,16 +422,16 @@ class GO2Connection(Module, Camera, Pointcloud): dedicated_worker = True config: ConnectionConfig - cmd_vel: In[Twist] - pointcloud: Out[PointCloud2] - odom: Out[PoseStamped] - lidar: Out[PointCloud2] - color_image: Out[Image] - camera_info: Out[CameraInfo] + cmd_vel: In[CdrTwist] + pointcloud: Out[CdrCloud] + odom: Out[CdrPoseStamped] + lidar: Out[CdrCloud] + color_image: Out[CdrImage] + camera_info: Out[CdrCameraInfo] tf: Out[TFMessage] connection: Go2ConnectionProtocol - camera_info_static: CameraInfo = _camera_info_static() + camera_info_static: CdrCameraInfo = _camera_info_message(_camera_info_static()) _camera_info_thread: Thread | None = None _latest_lowstate: LowStateMsg | None = None @@ -378,14 +457,36 @@ def __init__(self, **kwargs: Any) -> None: ) if hasattr(self.connection, "camera_info_static"): - self.camera_info_static = self.connection.camera_info_static + self.camera_info_static = _camera_info_message(self.connection.camera_info_static) - if self.config.frame_id_prefix and self.camera_info_static.frame_id: + if self.config.frame_id_prefix and self.camera_info_static.header.frame_id: # Copy so the class-level default is not mutated. - self.camera_info_static = copy.copy(self.camera_info_static) - self.camera_info_static.frame_id = _prefixed( - self.config.frame_id_prefix, self.camera_info_static.frame_id + self.camera_info_static = copy.deepcopy(self.camera_info_static) + self.camera_info_static.header.frame_id = _prefixed( + self.config.frame_id_prefix, self.camera_info_static.header.frame_id + ) + + def _publish_image(self, image: Image) -> None: + encodings = { + "RGB": "rgb8", + "BGR": "bgr8", + "RGBA": "rgba8", + "BGRA": "bgra8", + "GRAY": "mono8", + "GRAY16": "mono16", + "DEPTH": "32FC1", + "DEPTH16": "16UC1", + } + self.color_image.publish( + image_from_array( + image.data, + encoding=encodings[image.format.value], + header=Header( + time_from_seconds(image.ts), + _prefixed(self.config.frame_id_prefix, image.frame_id), + ), ) + ) @rpc def start(self) -> None: @@ -394,18 +495,18 @@ def start(self) -> None: return self.connection.start() - def onimage(image: Image) -> None: - image.frame_id = _prefixed(self.config.frame_id_prefix, image.frame_id) - self.color_image.publish(image) - if self.config.lidar: - self.register_disposable(self.connection.lidar_stream().subscribe(self.lidar.publish)) + self.register_disposable( + self.connection.lidar_stream() + .pipe(ops.map(_cloud_message)) + .subscribe(self.lidar.publish) + ) self.register_disposable(self.connection.odom_stream().subscribe(self._publish_tf)) self.register_disposable(self.connection.lowstate_stream().subscribe(self._on_lowstate)) self.register_disposable(Disposable(self.cmd_vel.subscribe(self.move))) if self.config.camera: - self.register_disposable(self.connection.video_stream().subscribe(onimage)) + self.register_disposable(self.connection.video_stream().subscribe(self._publish_image)) self._camera_info_thread = Thread( target=self.publish_camera_info, daemon=True, @@ -480,18 +581,44 @@ def _publish_tf(self, msg: PoseStamped) -> None: msg.frame_id = self.config.odom_frame_id if self.config.publish_tf: transforms = self._odom_to_tf(msg, prefix=self.config.frame_id_prefix or "") - self.tf.publish(TFMessage(*transforms)) + self.tf.publish( + TFMessage( + [ + CdrTransformStamped( + Header(time_from_seconds(transform.ts), transform.frame_id), + transform.child_frame_id, + CdrTransform( + CdrVector3(*transform.translation), + CdrQuaternion(*transform.rotation), + ), + ) + for transform in transforms + ] + ) + ) if self.odom.transport: - self.odom.publish(msg) + self.odom.publish( + CdrPoseStamped( + Header(time_from_seconds(msg.ts), msg.frame_id), + CdrPose(CdrPoint(*msg.position), CdrQuaternion(*msg.orientation)), + ) + ) def publish_camera_info(self) -> None: while True: - self.camera_info.publish(self.camera_info_static.with_ts(time.time())) + info = copy.deepcopy(self.camera_info_static) + info.header.stamp = time_from_seconds(time.time()) + self.camera_info.publish(info) time.sleep(1.0) @rpc - def move(self, twist: Twist, duration: float = 0.0) -> bool: + def move(self, twist: CdrTwist | Twist, duration: float = 0.0) -> bool: """Send movement command to robot.""" + if isinstance(twist, CdrTwist): + twist = Twist( + linear=Vector3(twist.linear.x, twist.linear.y, twist.linear.z), + angular=Vector3(twist.angular.x, twist.angular.y, twist.angular.z), + ) return self.connection.move(twist, duration) @rpc diff --git a/dimos/robot/unitree/go2/test_connection.py b/dimos/robot/unitree/go2/test_connection.py index 0698e12f94d..e23d849f252 100644 --- a/dimos/robot/unitree/go2/test_connection.py +++ b/dimos/robot/unitree/go2/test_connection.py @@ -19,18 +19,41 @@ """ from collections.abc import Callable, Iterator +from contextlib import ExitStack +from functools import partial import threading import time from types import SimpleNamespace from unittest.mock import MagicMock +from dimos_cdrgen.registry import decode, encode +import numpy as np +import open3d as o3d import pytest from pytest_mock import MockerFixture from reactivex import Subject +import dimos.core.cdr_transport as cdr_module from dimos.core.global_config import GlobalConfig +from dimos.core.transport_factory import make_transport +from dimos.messages.builtin_interfaces.msg import Time +from dimos.messages.geometry_msgs.msg import ( + PoseStamped as CdrPoseStamped, + Twist as CdrTwist, + Vector3, +) +from dimos.messages.sensor_msgs.msg import PointCloud2 as CdrCloud +from dimos.messages.std_msgs.msg import Header +from dimos.messages.tf2_msgs.msg import TFMessage from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo +from dimos.msgs.image import image_view +from dimos.msgs.pointcloud import pointcloud_view +from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo as DriverCameraInfo +from dimos.msgs.sensor_msgs.Image import Image as DriverImage, ImageFormat +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.time import to_seconds +from dimos.protocol.rpc.pubsubrpc import LCMRPC +from dimos.protocol.service.lcmservice import LCMConfig from dimos.robot.unitree.go2 import connection as go2_conn from dimos.robot.unitree.go2.connection import ConnectionConfig, GO2Connection @@ -40,6 +63,7 @@ def stub_webrtc(monkeypatch: pytest.MonkeyPatch) -> MagicMock: """Replace UnitreeWebRTCConnection in go2.connection so the webrtc branch runs without dialing out.""" stub = MagicMock(name="UnitreeWebRTCConnection") + stub.return_value.camera_info_static = DriverCameraInfo(frame_id="camera_optical", ts=0.0) monkeypatch.setattr(go2_conn, "UnitreeWebRTCConnection", stub) return stub @@ -148,7 +172,9 @@ def test_odom_to_tf_unprefixed_by_default() -> None: @pytest.fixture -def connection(stub_webrtc: MagicMock) -> Iterator[Callable[[bool], GO2Connection]]: +def connection( + stub_webrtc: MagicMock, mocker: MockerFixture +) -> Iterator[Callable[[bool], GO2Connection]]: """Build GO2Connections with the tf and odom ports stubbed, and stop them after.""" built: list[GO2Connection] = [] @@ -158,8 +184,8 @@ def build(publish_tf: bool) -> GO2Connection: publish_tf=publish_tf, odom_frame_id="go2_odom", ) - conn.tf = MagicMock() - conn.odom = MagicMock() + mocker.patch.object(conn, "tf") + mocker.patch.object(conn, "odom") built.append(conn) return conn @@ -194,8 +220,10 @@ def test_camera_info_is_restamped_on_each_publish( ) -> None: """A frozen stamp collapses a recording's camera_info onto one instant before the run.""" conn = connection(publish_tf=True) - conn.camera_info = MagicMock() - conn.camera_info_static = CameraInfo(frame_id="camera_optical", ts=1.0) + mocker.patch.object(conn, "camera_info") + conn.camera_info_static = go2_conn._camera_info_message( + DriverCameraInfo(ts=1.0, frame_id="camera_optical") + ) def sleep(_seconds: float) -> None: if conn.camera_info.publish.call_count >= 2: @@ -207,9 +235,11 @@ def sleep(_seconds: float) -> None: conn.publish_camera_info() published = [call.args[0] for call in conn.camera_info.publish.call_args_list] - assert [info.ts >= before for info in published] == [True, True] - assert [info.frame_id for info in published] == [conn.camera_info_static.frame_id] * 2 - assert conn.camera_info_static.ts == 1.0 + assert [to_seconds(info.header.stamp) >= before for info in published] == [True, True] + assert [info.header.frame_id for info in published] == [ + conn.camera_info_static.header.frame_id + ] * 2 + assert conn.camera_info_static.header.stamp == Time(1, 0) def test_odom_to_tf_prefixed() -> None: @@ -226,3 +256,136 @@ def test_odom_to_tf_prefixed() -> None: "robot0/camera_link", "robot0/camera_optical", ) + + +def test_driver_cloud_boundary_preserves_all_supported_attributes(): + source = PointCloud2.from_numpy( + points=np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), + intensities=np.array([0.25, 0.75]), + offset_times=np.array([3, 9], dtype=np.uint32), + tags=np.array([1, 2], dtype=np.uint8), + lines=np.array([2, 3], dtype=np.uint8), + stamps=np.array([10.0, 10.1]), + timestamp=0.0, + frame_id="lidar", + ) + source.pointcloud_tensor.point["colors"] = o3d.core.Tensor( + np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=np.float32) + ) + generated = go2_conn._cloud_message(source) + restored = decode(encode(generated), CdrCloud) + records = pointcloud_view(restored) + assert restored.header == Header(Time(0, 0), "lidar") + assert records["intensity"].ravel().tolist() == [0.25, 0.75] + assert records["offset_time"].ravel().tolist() == [3, 9] + assert records["tag"].ravel().tolist() == [1, 2] + assert records["line"].ravel().tolist() == [2, 3] + assert records["stamps"].ravel().tolist() == [10.0, 10.1] + assert records["rgb"].ravel().tolist() == [0xFF0000, 0x00FF00] + np.testing.assert_array_equal(records["x"].ravel(), [1.0, 4.0]) + np.testing.assert_array_equal(source.as_numpy()[0], [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + + +def test_driver_odometry_boundary_publishes_generated_tf_and_pose(connection): + conn = connection(publish_tf=True) + conn._publish_tf(PoseStamped(ts=0.0, frame_id="source")) + pose = conn.odom.publish.call_args.args[0] + transforms = conn.tf.publish.call_args.args[0] + assert type(pose) is CdrPoseStamped + assert type(transforms) is TFMessage + assert pose.header == Header(Time(0, 0), "go2_odom") + assert [(t.header.frame_id, t.child_frame_id) for t in transforms.transforms] == [ + ("go2_odom", "base_link"), + ("base_link", "camera_link"), + ("camera_link", "camera_optical"), + ] + assert all(t.header.stamp == Time(0, 0) for t in transforms.transforms) + + +def test_generated_command_is_adapted_only_at_driver_boundary(connection): + conn = connection(publish_tf=False) + conn.move(CdrTwist(Vector3(1.0, 2.0, 3.0), Vector3(4.0, 5.0, 6.0)), 0.5) + command, duration = conn.connection.move.call_args.args + assert command.linear.x == 1.0 + assert command.linear.y == 2.0 + assert command.linear.z == 3.0 + assert command.angular.x == 4.0 + assert command.angular.y == 5.0 + assert command.angular.z == 6.0 + assert duration == 0.5 + + +@pytest.mark.parametrize( + "format,encoding,dtype", + [ + (ImageFormat.RGB, "rgb8", np.uint8), + (ImageFormat.DEPTH16, "16UC1", np.uint16), + ], +) +def test_driver_image_boundary_copies_pixels_and_prefixes_frame( + connection, mocker, format, encoding, dtype +): + conn = connection(publish_tf=False) + conn.config.frame_id_prefix = "robot0" + mocker.patch.object(conn, "color_image") + shape = (2, 3, 3) if format is ImageFormat.RGB else (2, 3) + pixels = np.arange(np.prod(shape), dtype=dtype).reshape(shape) + source = DriverImage(data=pixels, format=format, ts=0.0, frame_id="camera_optical") + conn._publish_image(source) + message = conn.color_image.publish.call_args.args[0] + assert message.header == Header(Time(0, 0), "robot0/camera_optical") + assert message.encoding == encoding + np.testing.assert_array_equal(image_view(message), pixels) + pixels[:] = 0 + assert image_view(message).max() > 0 + assert source.frame_id == "camera_optical" + + +class MemqRPC(LCMRPC): + def __init__(self, **kwargs): + super().__init__(url="memq://", **kwargs) + + +def test_go2_default_ports_deliver_cdr_without_robot(stub_webrtc, mocker, monkeypatch): + mocker.patch("dimos.core.module.get_loop", return_value=(None, None)) + mocker.patch.object(LCMRPC, "start") + mocker.patch.object(LCMRPC, "serve_module_rpc") + monkeypatch.setattr(cdr_module, "LCMConfig", partial(LCMConfig, url="memq://")) + config = GlobalConfig(transport="lcm", robot_ip="127.0.0.1") + received = {} + condition = threading.Condition() + + def on_message(name, value): + with condition: + received[name] = value + condition.notify_all() + + with ExitStack() as cleanup: + conn = GO2Connection(g=config, rpc_transport=MemqRPC) + cleanup.callback(conn.stop) + for name in ("tf", "odom", "color_image", "lidar"): + port = conn.outputs[name] + transport = make_transport(name, port.type, g=config) + cleanup.callback(transport.stop) + conn.set_transport(name, transport) + cleanup.callback(transport.subscribe(partial(on_message, name))) + source = DriverImage( + data=np.arange(18, dtype=np.uint8).reshape(2, 3, 3), + format=ImageFormat.RGB, + ts=0.0, + frame_id="camera_optical", + ) + conn._publish_image(source) + conn._publish_tf(PoseStamped(ts=0.0, frame_id="world")) + cloud = PointCloud2.from_numpy(np.array([[1.0, 2.0, 3.0]]), timestamp=0.0) + conn.lidar.publish(go2_conn._cloud_message(cloud)) + with condition: + assert condition.wait_for(lambda: len(received) == 4, timeout=5) + assert type(received["odom"]) is CdrPoseStamped + assert type(received["tf"]) is TFMessage + assert type(received["lidar"]) is CdrCloud + assert received["color_image"].encoding == "rgb8" + assert received["odom"].header.stamp == Time(0, 0) + assert len(received["tf"].transforms) == 3 + np.testing.assert_array_equal(image_view(received["color_image"]), source.data) + np.testing.assert_array_equal(pointcloud_view(received["lidar"])["z"], [[3.0]]) diff --git a/dimos/spec/perception.py b/dimos/spec/perception.py index 29c715710be..7ef34e6ebc1 100644 --- a/dimos/spec/perception.py +++ b/dimos/spec/perception.py @@ -15,10 +15,9 @@ from typing import Protocol from dimos.core.stream import Out -from dimos.messages.sensor_msgs.msg import CameraInfo, Image as ImageMsg +from dimos.messages.sensor_msgs.msg import CameraInfo, Image as ImageMsg, PointCloud2 from dimos.msgs.nav_msgs.Odometry import Odometry as OdometryMsg from dimos.msgs.sensor_msgs.Imu import Imu -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 class Image(Protocol):