From 26ce4cbf164af48e2793f9141deda0d8af5d273a Mon Sep 17 00:00:00 2001 From: cc Date: Tue, 20 Oct 2026 01:04:42 -0700 Subject: [PATCH 1/2] refactor(memory): accept generated poses in lossless pose-fill --- .github/workflows/message-codegen.yml | 7 ++ dimos/mapping/cli/pose_fill.py | 108 ++++++++++++++------------ dimos/mapping/cli/test_pose_fill.py | 97 +++++++++++++++++++++-- dimos/memory/type/observation.py | 14 +++- dimos/memory/type/test_observation.py | 92 ++++++++++++++++++++++ 5 files changed, 258 insertions(+), 60 deletions(-) create mode 100644 dimos/memory/type/test_observation.py diff --git a/.github/workflows/message-codegen.yml b/.github/workflows/message-codegen.yml index f1b4b144ead..cf06937b10e 100644 --- a/.github/workflows/message-codegen.yml +++ b/.github/workflows/message-codegen.yml @@ -64,6 +64,8 @@ on: - native/rust/dimos-module/** - native/rust/dimos-module-macros/** - native/message-source.nix + - dimos/mapping/cli/pose_fill.py + - dimos/mapping/cli/test_pose_fill.py - native/cpp/include/dimos/native/cdr_codec.hpp - docs/development/message-source-packages.md - docs/development/messages-external-project.md @@ -366,3 +368,8 @@ jobs: export DIMOS_NATIVE_MLS_BINARY="$PWD/dimos/navigation/global_planner/mls_planner/rust/target/debug/mls_planner" .native-review/bin/python -I -c 'from dimos.navigation.global_planner.mls_planner import mls_planner_native,transformer,viz; import dimos_mls_planner; import pytest; import sys; assert all("site-packages" in str(value.__file__) for value in (mls_planner_native,transformer,viz,dimos_mls_planner)); raise SystemExit(pytest.main(sys.argv[1:]))' --noconftest --import-mode=importlib -o addopts='' -q "$PWD/dimos/navigation/global_planner/mls_planner/test_generated_native.py" "$PWD/dimos/navigation/global_planner/mls_planner/test_transformer.py" "$PWD/dimos/navigation/global_planner/mls_planner/test_viz.py" MYPYPATH="$PWD/build/message-source/python" .native-review/bin/python -m mypy --follow-imports=silent dimos/navigation/global_planner/mls_planner/{mls_planner_native,transformer,viz}.py dimos/navigation/spec.py + - name: Verify canonical observation pose payloads and lossless pose-fill + if: runner.os == 'Linux' + run: | + .native-review/bin/python -I -c 'from dimos.memory.type import observation; from dimos.mapping.cli import pose_fill; import pytest; import sys; assert all("site-packages" in str(value.__file__) for value in (observation,pose_fill)); raise SystemExit(pytest.main(sys.argv[1:]))' --noconftest --import-mode=importlib -o addopts='' -q "$PWD/dimos/memory/type/test_observation.py" "$PWD/dimos/mapping/cli/test_pose_fill.py" "$PWD/dimos/memory/test_stream.py::TestWithPose" "$PWD/dimos/memory/store/test_cdr_recording.py" + MYPYPATH="$PWD/build/message-source/python" .native-review/bin/python -m mypy --follow-imports=silent dimos/memory/type/observation.py dimos/mapping/cli/pose_fill.py diff --git a/dimos/mapping/cli/pose_fill.py b/dimos/mapping/cli/pose_fill.py index 38465ff59fb..4d1220081eb 100644 --- a/dimos/mapping/cli/pose_fill.py +++ b/dimos/mapping/cli/pose_fill.py @@ -40,7 +40,7 @@ from dimos.memory.backend import Backend from dimos.memory.stream import Stream from dimos.memory.type.observation import Observation - from dimos.msgs.geometry_msgs.Transform import Transform + from dimos.messages.geometry_msgs.msg import Transform def pose_fill( @@ -54,10 +54,10 @@ def pose_fill( Pairs are formed by :meth:`Stream.align` (nearest ``|Δts| <= tolerance``); primaries with no match in tolerance are dropped. The pose is read from the - matched pose observation's *payload* (``PoseStamped`` / ``Odometry`` / any - object exposing ``.position`` + ``.orientation``) because pose-message - streams carry the pose in the value, not the indexed pose columns. The - target stream's payload stays lazy. + matched pose observation's payload: generated ``Pose``, ``Transform`` or + their standard stamped/covariance/odometry forms. Pose-message streams + carry the pose in the value, not the indexed pose columns. The target + stream's payload stays lazy. *mount* composes a static child transform onto each source pose (``world_base ∘ mount``) before attaching — e.g. ``base_link → @@ -65,9 +65,8 @@ def pose_fill( image pose. """ from dimos.memory.type.observation import _to_tuple - 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.messages.geometry_msgs.msg import Quaternion, Transform, Vector3 + from dimos.msgs.geometry import compose_transform_values def _fill(pair_obs: Observation[Any]) -> Observation[Any]: primary, secondary = cast( @@ -77,7 +76,7 @@ def _fill(pair_obs: Observation[Any]) -> Observation[Any]: return primary.with_pose(secondary.data) x, y, z, qx, qy, qz, qw = cast("tuple[float, ...]", _to_tuple(secondary.data)) world_base = Transform(translation=Vector3(x, y, z), rotation=Quaternion(qx, qy, qz, qw)) - return primary.with_pose(world_base + mount) + return primary.with_pose(compose_transform_values(world_base, mount)) return stream.align(pose_stream, tolerance=tolerance).map(_fill) @@ -112,48 +111,57 @@ def pose_fill_db( "collide on ts). Delete it (and any -wal/-shm sidecars) first." ) - src = SqliteStore(path=str(src_path), must_exist=True) - dest = SqliteStore(path=str(dest_path)) - names = streams if streams is not None else src.list_streams() - if target not in names: - raise ValueError(f"target stream {target!r} not in {names}") - if pose_source not in src.list_streams(): - raise ValueError(f"pose_source stream {pose_source!r} not found in source dataset") - - written: dict[str, int] = {} - for name in names: - src_b = cast("Backend[Any]", src.stream(name)._source) - dest_b = cast("Backend[Any]", dest.stream(name, src_b.data_type, codec=src_b.codec)._source) - - # Re-posed observations for the target (ids preserved from source, so - # the source blob still resolves); plain metadata for everything else. - if name == target: - rows: Iterable[Observation[Any]] = pose_fill( - src.stream(name).from_time(seek or None).to_time(duration).order_by("ts"), - src.stream(pose_source).from_time(seek or None).to_time(duration).order_by("ts"), - tolerance=tolerance, + with ( + SqliteStore(path=str(src_path), must_exist=True) as src, + SqliteStore(path=str(dest_path)) as dest, + ): + names = streams if streams is not None else src.list_streams() + if target not in names: + raise ValueError(f"target stream {target!r} not in {names}") + if pose_source not in src.list_streams(): + raise ValueError(f"pose_source stream {pose_source!r} not found in source dataset") + + written: dict[str, int] = {} + for name in names: + src_b = cast("Backend[Any]", src.stream(name)._source) + dest_b = cast( + "Backend[Any]", dest.stream(name, src_b.data_type, codec=src_b.codec)._source ) - else: - rows = src.stream(name).from_time(seek or None).to_time(duration).order_by("ts") - - scalar = src_b.data_type in (int, float) - n = 0 - for obs in rows: - blob = ( - None if scalar or src_b.blob_store is None else src_b.blob_store.get(name, obs.id) - ) - row_id = dest_b.metadata_store.insert(obs) - if blob is not None: - assert dest_b.blob_store is not None - dest_b.blob_store.put(name, row_id, blob) - n += 1 - commit = getattr(dest_b.metadata_store, "commit", None) - if commit is not None: - commit() - written[name] = n - - src.stop() - dest.stop() + + # Re-posed observations for the target (ids preserved from source, so + # the source blob still resolves); plain metadata for everything else. + if name == target: + rows: Iterable[Observation[Any]] = pose_fill( + src.stream(name).from_time(seek or None).to_time(duration).order_by("ts"), + src.stream(pose_source) + .from_time(seek or None) + .to_time(duration) + .order_by("ts"), + tolerance=tolerance, + ) + else: + rows = src.stream(name).from_time(seek or None).to_time(duration).order_by("ts") + + scalar = src_b.data_type in (int, float) + n = 0 + for obs in rows: + blob = ( + None + if scalar or src_b.blob_store is None + else src_b.blob_store.get(name, obs.id) + ) + row_id = dest_b.metadata_store.insert(obs) + if blob is not None: + assert dest_b.blob_store is not None + dest_b.blob_store.put(name, row_id, blob) + n += 1 + commit = getattr(dest_b.metadata_store, "commit", None) + if commit is not None: + commit() + written[name] = n + + src.stop() + dest.stop() return written diff --git a/dimos/mapping/cli/test_pose_fill.py b/dimos/mapping/cli/test_pose_fill.py index 2953d9e13e6..a2adc2a4982 100644 --- a/dimos/mapping/cli/test_pose_fill.py +++ b/dimos/mapping/cli/test_pose_fill.py @@ -20,11 +20,28 @@ from __future__ import annotations -from dimos.mapping.cli.pose_fill import pose_fill +from copy import deepcopy +import math + +from dimos_cdrgen.registry import encode +import numpy as np + +from dimos.mapping.cli.pose_fill import pose_fill, pose_fill_db from dimos.memory.store.memory import MemoryStore -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.memory.store.sqlite import SqliteStore +from dimos.messages.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Transform, + Vector3, +) +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.messages.std_msgs.msg import Header +from dimos.msgs.geometry import quaternion_from_euler +from dimos.msgs.pointcloud import pointcloud_from_xyz +from dimos.msgs.time import time_from_nanoseconds def test_pose_fill_attaches_nearest_pose() -> None: @@ -33,7 +50,12 @@ def test_pose_fill_attaches_nearest_pose() -> None: target = store.stream("image", str) poses = store.stream("odom", Transform) target.append("img0", ts=0.0) - poses.append(Transform(translation=Vector3(1.0, 0.0, 0.0)), ts=0.001) + poses.append( + Transform( + translation=Vector3(x=1.0, y=0.0, z=0.0), rotation=Quaternion(x=0, y=0, z=0, w=1) + ), + ts=0.001, + ) out = pose_fill(target, poses, tolerance=0.05).to_list() @@ -50,8 +72,15 @@ def test_pose_fill_mount_composes_static_child_transform() -> None: poses = store.stream("odom", Transform) target.append("img0", ts=0.0) # Base pose 1m forward in x; mount offsets 1m in y (identity rotations). - poses.append(Transform(translation=Vector3(1.0, 0.0, 0.0)), ts=0.001) - mount = Transform(translation=Vector3(0.0, 1.0, 0.0), rotation=Quaternion()) + poses.append( + Transform( + translation=Vector3(x=1.0, y=0.0, z=0.0), rotation=Quaternion(x=0, y=0, z=0, w=1) + ), + ts=0.001, + ) + mount = Transform( + translation=Vector3(x=0, y=1, z=0), rotation=Quaternion(x=0, y=0, z=0, w=1) + ) out = pose_fill(target, poses, tolerance=0.05, mount=mount).to_list() @@ -59,3 +88,57 @@ def test_pose_fill_mount_composes_static_child_transform() -> None: assert out[0].pose_tuple is not None # Identity rotations → composed translation is the component-wise sum. assert out[0].pose_tuple[:3] == (1.0, 1.0, 0.0) + + +def test_stamped_pose_mount_rotates_the_offset_and_keeps_payload(): + header = Header(frame_id="world", stamp=time_from_nanoseconds(1700000000123456789)) + pose = PoseStamped( + header=header, + pose=Pose( + position=Point(x=1, y=2, z=3), orientation=quaternion_from_euler(0, 0, math.pi / 2) + ), + ) + mount = Transform(translation=Vector3(x=1, y=0, z=0), rotation=Quaternion(x=0, y=0, z=0, w=1)) + with MemoryStore() as store: + target = store.stream("target", str) + source = store.stream("pose", PoseStamped) + target.append("frame", ts=10.0) + source.append(pose, ts=10.01) + before = encode(pose) + [result] = pose_fill(target, source, mount=mount).to_list() + assert result.data == "frame" + np.testing.assert_allclose(result.pose_tuple[:3], [1, 3, 3], atol=1e-12) + assert encode(pose) == before + + +def test_sqlite_pose_fill_copies_original_cdr_bytes_and_exact_header(tmp_path): + source_path = tmp_path / "source.db" + destination = tmp_path / "posed.db" + header = Header(frame_id="lidar", stamp=time_from_nanoseconds(1700000000123456789)) + cloud = pointcloud_from_xyz(np.array([[1, 2, 3]], dtype=np.float32), header=header) + pose_header = deepcopy(header) + pose_header.frame_id = "odom" + pose = PoseStamped( + header=pose_header, + pose=Pose(position=Point(x=4, y=5, z=6), orientation=Quaternion(x=0, y=0, z=0, w=1)), + ) + with SqliteStore(path=str(source_path)) as store: + store.stream("lidar", PointCloud2).append(cloud, ts=10.0) + store.stream("odom", PoseStamped).append(pose, ts=10.001) + written = pose_fill_db(source_path, destination, target="lidar", pose_source="odom") + assert written == {"lidar": 1, "odom": 1} + with ( + SqliteStore(path=str(source_path), must_exist=True) as original, + SqliteStore(path=str(destination), must_exist=True) as store, + ): + result = store.stream("lidar", PointCloud2).first() + assert result.pose_tuple == (4.0, 5.0, 6.0, 0.0, 0.0, 0.0, 1.0) + assert result.ts == 10.0 + assert encode(result.data) == encode(cloud) + assert result.data.header.stamp.nanosec == 123456789 + assert encode(store.stream("odom", PoseStamped).first().data) == encode(pose) + before = original.stream("lidar", PointCloud2) + after = store.stream("lidar", PointCloud2) + assert before._source.blob_store.get( + "lidar", before.first().id + ) == after._source.blob_store.get("lidar", result.id) diff --git a/dimos/memory/type/observation.py b/dimos/memory/type/observation.py index 3a42b4c80d0..543ed3ac17b 100644 --- a/dimos/memory/type/observation.py +++ b/dimos/memory/type/observation.py @@ -65,9 +65,8 @@ def _to_tuple(p: Any) -> PoseTuple | None: Accepts: `None`, a 3- or 7-tuple, or any object with `translation`+`rotation` (Transform) or `position`+`orientation` - (Pose / PoseStamped) attributes. 3-tuples are padded with the identity - quaternion. Duck-typed to avoid importing `dimos.msgs.geometry_msgs` - at module load. + (Pose) attributes, including the standard generated stamped/covariance/ + odometry nesting. 3-tuples are padded with the identity quaternion. """ if p is None: return None @@ -78,6 +77,15 @@ def _to_tuple(p: Any) -> PoseTuple | None: x, y, z = p return (float(x), float(y), float(z), 0.0, 0.0, 0.0, 1.0) raise TypeError(f"Pose tuple must have length 3 or 7, got {len(p)}") + # Generated ROS messages retain their declared nesting; only unwrap the + # pose-bearing shapes, leaving the indexed storage tuple contract unchanged. + name = getattr(type(p), "__msgtype__", None) + if name in ("nav_msgs/msg/Odometry", "geometry_msgs/msg/PoseWithCovarianceStamped"): + p = p.pose.pose + elif name in ("geometry_msgs/msg/PoseStamped", "geometry_msgs/msg/PoseWithCovariance"): + p = p.pose + elif name == "geometry_msgs/msg/TransformStamped": + p = p.transform # Use ``is not None`` rather than ``or`` — a Vector3 at the origin is # falsy but valid, and falling through to ``position`` on a Transform # would crash. diff --git a/dimos/memory/type/test_observation.py b/dimos/memory/type/test_observation.py new file mode 100644 index 00000000000..25e3df391ad --- /dev/null +++ b/dimos/memory/type/test_observation.py @@ -0,0 +1,92 @@ +# 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 copy import deepcopy + +from dimos_cdrgen.registry import encode +import numpy as np +import pytest + +from dimos.memory.type.observation import Observation +from dimos.messages.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + PoseWithCovariance, + Quaternion, + Transform, + TransformStamped, + Twist, + TwistWithCovariance, + Vector3, +) +from dimos.messages.nav_msgs.msg import Odometry +from dimos.messages.std_msgs.msg import Header +from dimos.msgs.time import time_from_nanoseconds + + +def canonical_pose_values(): + header = Header(frame_id="odom", stamp=time_from_nanoseconds(1700000000123456789)) + pose = Pose( + position=Point(x=1.5, y=-2, z=0.7), orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.927) + ) + transform = Transform( + translation=Vector3(x=1.5, y=-2, z=0.7), rotation=deepcopy(pose.orientation) + ) + covariance = PoseWithCovariance(pose=deepcopy(pose), covariance=np.zeros(36, dtype=np.float64)) + return [ + pose, + PoseStamped(header=deepcopy(header), pose=deepcopy(pose)), + transform, + TransformStamped( + header=deepcopy(header), child_frame_id="base", transform=deepcopy(transform) + ), + covariance, + Odometry( + header=deepcopy(header), + child_frame_id="base", + pose=deepcopy(covariance), + twist=TwistWithCovariance( + twist=Twist(linear=Vector3(x=0, y=0, z=0), angular=Vector3(x=0, y=0, z=0)), + covariance=np.zeros(36, dtype=np.float64), + ), + ), + ] + + +@pytest.mark.parametrize("index", range(6)) +def test_canonical_pose_payloads_normalize_metadata_without_loading_or_mutation(index): + value = canonical_pose_values()[index] + before = encode(value) + loads = [] + + def load(): + loads.append("loaded") + return "payload" + + observation = Observation(id=7, ts=4.5, _loader=load) + posed = observation.with_pose(value) + assert posed.pose_tuple == (1.5, -2.0, 0.7, 0.1, 0.2, 0.3, 0.927) + assert observation.pose_tuple is None and posed.id == 7 and posed.ts == 4.5 + assert loads == [] + assert posed.data == observation.data == "payload" + assert loads == ["loaded"] + assert encode(value) == before + assert observation.with_pose(None).pose_tuple is None + + +def test_unsupported_message_is_not_silently_treated_as_a_pose(): + header = Header(frame_id="odom", stamp=time_from_nanoseconds(0)) + with pytest.raises(TypeError, match="Cannot coerce"): + Observation(pose=header, _data="unchanged") From b6700a02db57422bc96f77d2d9a4fd84854fcb5c Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Tue, 20 Oct 2026 08:05:48 +0000 Subject: [PATCH 2/2] [autofix.ci] apply automated fixes --- dimos/models/embedding/clip.py | 2 +- dimos/models/embedding/dino.py | 2 +- dimos/models/embedding/siglip.py | 2 +- dimos/models/embedding/siglip2.py | 2 +- .../models/embedding/test_generated_images.py | 37 +++++++++++-------- dimos/models/embedding/treid.py | 2 +- dimos/models/segmentation/edge_tam.py | 2 +- dimos/models/segmentation/test_yoloe.py | 4 +- dimos/models/segmentation/yoloe.py | 2 +- 9 files changed, 30 insertions(+), 25 deletions(-) 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