From e86cbe81130bc8e60284b68d14f5cd4426b046d4 Mon Sep 17 00:00:00 2001 From: cc Date: Mon, 19 Oct 2026 22:56:20 -0700 Subject: [PATCH 1/2] refactor: retain generated CDR values through navigation search --- .github/workflows/message-codegen.yml | 7 + dimos/mapping/occupancy/path_map.py | 7 +- .../go2/replanning_a_star/min_cost_astar.py | 165 +++++++++++------- .../replanning_a_star/test_generated_astar.py | 129 ++++++++++++++ 4 files changed, 240 insertions(+), 68 deletions(-) create mode 100644 dimos/navigation/go2/replanning_a_star/test_generated_astar.py diff --git a/.github/workflows/message-codegen.yml b/.github/workflows/message-codegen.yml index b3a41a22629..567c681a633 100644 --- a/.github/workflows/message-codegen.yml +++ b/.github/workflows/message-codegen.yml @@ -21,6 +21,8 @@ on: - dimos/memory/** - dimos/mapping/occupancy/** - dimos/mapping/voxels/** + - dimos/navigation/go2/replanning_a_star/min_cost_astar.py + - dimos/navigation/go2/replanning_a_star/test_generated_astar.py - dimos/mapping/pointclouds/occupancy.py - dimos/mapping/odometry_hist.py - dimos/mapping/test_generated_odometry_hist.py @@ -255,3 +257,8 @@ jobs: run: | .native-review/bin/python -I -c 'from dimos.mapping.voxels import grid, module; from dimos.mapping.voxels.impl import packed, o3d; import pytest; import sys; assert all("site-packages" in str(value.__file__) for value in (grid, module, packed, o3d)); raise SystemExit(pytest.main(sys.argv[1:]))' --noconftest --import-mode=importlib -o addopts='' -q "$PWD/dimos/mapping/voxels/test_generated_stream.py" "$PWD/dimos/mapping/voxels/test_stamped.py" MYPYPATH="$PWD/build/message-source/python" .native-review/bin/python -m mypy --follow-imports=silent dimos/mapping/voxels/grid.py dimos/mapping/voxels/module.py dimos/mapping/voxels/impl/packed.py dimos/mapping/voxels/impl/o3d.py + - name: Verify generated navigation costmaps and Python/C++ A* paths + if: runner.os == 'Linux' + run: | + .native-review/bin/python -I -c 'from dimos.mapping.occupancy import path_map; from importlib import import_module; astar = import_module("dimos.navigation.go2.replanning_a_star.min_cost_astar"); import pytest; import sys; assert all("site-packages" in str(value.__file__) for value in (path_map, astar)); raise SystemExit(pytest.main(sys.argv[1:]))' --noconftest --import-mode=importlib -o addopts='' -q "$PWD/dimos/navigation/go2/replanning_a_star/test_generated_astar.py" + MYPYPATH="$PWD/build/message-source/python" .native-review/bin/python -m mypy --follow-imports=silent dimos/mapping/occupancy/path_map.py dimos/navigation/go2/replanning_a_star/min_cost_astar.py diff --git a/dimos/mapping/occupancy/path_map.py b/dimos/mapping/occupancy/path_map.py index 69a1f937385..a1f5e29ddba 100644 --- a/dimos/mapping/occupancy/path_map.py +++ b/dimos/mapping/occupancy/path_map.py @@ -14,17 +14,16 @@ from dimos.mapping.occupancy.gradient import GradientStrategy, gradient, voronoi_gradient from dimos.mapping.occupancy.inflation import simple_inflate -from dimos.mapping.occupancy.operations import overlay_occupied, smooth_occupied +from dimos.mapping.occupancy.operations import _GridT, overlay_occupied, smooth_occupied from dimos.mapping.occupancy.types import NavigationStrategy -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid def make_navigation_map( - occupancy_grid: OccupancyGrid, + occupancy_grid: _GridT, robot_width: float, strategy: NavigationStrategy, gradient_strategy: GradientStrategy, -) -> OccupancyGrid: +) -> _GridT: half_width = robot_width / 2 gradient_distance = 1.5 diff --git a/dimos/navigation/go2/replanning_a_star/min_cost_astar.py b/dimos/navigation/go2/replanning_a_star/min_cost_astar.py index 03311db6897..bd75418d4e0 100644 --- a/dimos/navigation/go2/replanning_a_star/min_cost_astar.py +++ b/dimos/navigation/go2/replanning_a_star/min_cost_astar.py @@ -12,13 +12,26 @@ # See the License for the specific language governing permissions and # limitations under the License. +from copy import deepcopy import heapq - +from typing import overload + +import numpy as np +from numpy.typing import NDArray + +from dimos.messages.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped as CdrPoseStamped, + Quaternion as CdrQuaternion, +) +from dimos.messages.nav_msgs.msg import OccupancyGrid as CdrOccupancyGrid, Path as CdrPath from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import VectorLike +from dimos.msgs.geometry_msgs.Vector3 import Vector3, VectorLike from dimos.msgs.nav_msgs.OccupancyGrid import CostValues, OccupancyGrid from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.occupancy import grid_to_world, occupancy_view, world_to_grid from dimos.utils.logging_config import setup_logger # Try to import C++ extension for faster pathfinding @@ -64,89 +77,99 @@ def _heuristic(x1: int, y1: int, x2: int, y2: int) -> float: def _reconstruct_path( parents: dict[tuple[int, int], tuple[int, int]], current: tuple[int, int], - costmap: OccupancyGrid, start_tuple: tuple[int, int], - goal_tuple: tuple[int, int], -) -> Path: - frame_id = costmap.frame_id - waypoints: list[PoseStamped] = [] +) -> list[tuple[int, int]]: + coordinates = [] while current in parents: - world_point = costmap.grid_to_world(current) - pose = PoseStamped( - frame_id=frame_id, - position=[world_point.x, world_point.y, 0.0], - orientation=Quaternion(0, 0, 0, 1), # Identity quaternion - ) - waypoints.append(pose) + coordinates.append(current) current = parents[current] + coordinates.append(start_tuple) + coordinates.reverse() + return coordinates - start_world_point = costmap.grid_to_world(start_tuple) - start_pose = PoseStamped( - frame_id=frame_id, - position=[start_world_point.x, start_world_point.y, 0.0], - orientation=Quaternion(0, 0, 0, 1), - ) - waypoints.append(start_pose) - - waypoints.reverse() - - # Add the goal position if it's not already included - goal_point = costmap.grid_to_world(goal_tuple) - if ( - not waypoints - or (waypoints[-1].x - goal_point.x) ** 2 + (waypoints[-1].y - goal_point.y) ** 2 > 1e-10 - ): - goal_pose = PoseStamped( - frame_id=frame_id, - position=[goal_point.x, goal_point.y, 0.0], - orientation=Quaternion(0, 0, 0, 1), +def _reconstruct_path_from_coords( + path_coords: list[tuple[int, int]], + costmap: OccupancyGrid | CdrOccupancyGrid, +) -> Path | CdrPath: + if isinstance(costmap, CdrOccupancyGrid): + return CdrPath( + header=deepcopy(costmap.header), + poses=[ + CdrPoseStamped( + header=deepcopy(costmap.header), + pose=Pose( + position=grid_to_world(costmap, point), + orientation=CdrQuaternion(x=0, y=0, z=0, w=1), + ), + ) + for point in path_coords + ], + ) + waypoints = [] + for coordinate in path_coords: + point = costmap.grid_to_world(coordinate) + waypoints.append( + PoseStamped( + frame_id=costmap.frame_id, + position=[point.x, point.y, 0.0], + orientation=Quaternion(0, 0, 0, 1), + ) ) - waypoints.append(goal_pose) + return Path(frame_id=costmap.frame_id, poses=waypoints) - return Path(frame_id=frame_id, poses=waypoints) +def _cell(costmap: OccupancyGrid | CdrOccupancyGrid, value: Point | VectorLike) -> tuple[int, int]: + vector = Vector3((value.x, value.y, value.z)) if isinstance(value, Point) else Vector3(value) + if isinstance(costmap, CdrOccupancyGrid): + x, y = world_to_grid(costmap, Point(x=vector.x, y=vector.y, z=vector.z)) + else: + coordinates = costmap.world_to_grid(vector) + x, y = coordinates.x, coordinates.y + return int(x), int(y) -def _reconstruct_path_from_coords( - path_coords: list[tuple[int, int]], - costmap: OccupancyGrid, -) -> Path: - frame_id = costmap.frame_id - waypoints: list[PoseStamped] = [] - - for gx, gy in path_coords: - world_point = costmap.grid_to_world((gx, gy)) - pose = PoseStamped( - frame_id=frame_id, - position=[world_point.x, world_point.y, 0.0], - orientation=Quaternion(0, 0, 0, 1), - ) - waypoints.append(pose) - return Path(frame_id=frame_id, poses=waypoints) +@overload +def min_cost_astar( + costmap: CdrOccupancyGrid, + goal: Point | VectorLike, + start: Point | VectorLike = (0.0, 0.0), + cost_threshold: int = 100, + unknown_penalty: float = 0.8, + use_cpp: bool = True, +) -> CdrPath | None: ... +@overload def min_cost_astar( costmap: OccupancyGrid, - goal: VectorLike, - start: VectorLike = (0.0, 0.0), + goal: Point | VectorLike, + start: Point | VectorLike = (0.0, 0.0), cost_threshold: int = 100, unknown_penalty: float = 0.8, use_cpp: bool = True, -) -> Path | None: - start_vector = costmap.world_to_grid(start) - goal_vector = costmap.world_to_grid(goal) +) -> Path | None: ... - start_tuple = (int(start_vector.x), int(start_vector.y)) - goal_tuple = (int(goal_vector.x), int(goal_vector.y)) - if not (0 <= goal_tuple[0] < costmap.width and 0 <= goal_tuple[1] < costmap.height): +def min_cost_astar( + costmap: OccupancyGrid | CdrOccupancyGrid, + goal: Point | VectorLike, + start: Point | VectorLike = (0.0, 0.0), + cost_threshold: int = 100, + unknown_penalty: float = 0.8, + use_cpp: bool = True, +) -> Path | CdrPath | None: + start_tuple = _cell(costmap, start) + goal_tuple = _cell(costmap, goal) + cells = occupancy_view(costmap) if isinstance(costmap, CdrOccupancyGrid) else costmap.grid + height, width = cells.shape + if not (0 <= goal_tuple[0] < width and 0 <= goal_tuple[1] < height): return None if use_cpp: if _USE_CPP: path_coords = _astar_cpp( - costmap.grid, + cells, start_tuple[0], start_tuple[1], goal_tuple[0], @@ -163,6 +186,20 @@ def min_cost_astar( _CPP_IMPORT_ERROR, ) + coordinates = _find_path(cells, start_tuple, goal_tuple, cost_threshold, unknown_penalty) + if coordinates is None: + return None + return _reconstruct_path_from_coords(coordinates, costmap) + + +def _find_path( + cells: NDArray[np.int8], + start_tuple: tuple[int, int], + goal_tuple: tuple[int, int], + cost_threshold: int, + unknown_penalty: float, +) -> list[tuple[int, int]] | None: + height, width = cells.shape open_set: list[tuple[float, float, tuple[int, int]]] = [] # Priority queue for nodes to explore closed_set: set[tuple[int, int]] = set() # Set of explored nodes @@ -185,7 +222,7 @@ def min_cost_astar( continue if current == goal_tuple: - return _reconstruct_path(parents, current, costmap, start_tuple, goal_tuple) + return _reconstruct_path(parents, current, start_tuple) closed_set.add(current) @@ -193,13 +230,13 @@ def min_cost_astar( neighbor_x, neighbor_y = current_x + dx, current_y + dy neighbor = (neighbor_x, neighbor_y) - if not (0 <= neighbor_x < costmap.width and 0 <= neighbor_y < costmap.height): + if not (0 <= neighbor_x < width and 0 <= neighbor_y < height): continue if neighbor in closed_set: continue - neighbor_val = costmap.grid[neighbor_y, neighbor_x] + neighbor_val = cells[neighbor_y, neighbor_x] if neighbor_val >= cost_threshold: continue diff --git a/dimos/navigation/go2/replanning_a_star/test_generated_astar.py b/dimos/navigation/go2/replanning_a_star/test_generated_astar.py new file mode 100644 index 00000000000..21386373155 --- /dev/null +++ b/dimos/navigation/go2/replanning_a_star/test_generated_astar.py @@ -0,0 +1,129 @@ +# 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. + +from copy import deepcopy +from itertools import pairwise +import math + +from dimos_cdrgen.registry import decode, encode +import numpy as np +import pytest + +from dimos.mapping.occupancy.path_map import make_navigation_map +from dimos.messages.geometry_msgs.msg import Point, Pose, Quaternion +from dimos.messages.nav_msgs.msg import Path +from dimos.messages.std_msgs.msg import Header +from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid as LegacyGrid +from dimos.msgs.occupancy import grid_to_world, occupancy_from_array, occupancy_view, world_to_grid +from dimos.msgs.time import time_from_nanoseconds +from dimos.navigation.go2.replanning_a_star.min_cost_astar import _USE_CPP, min_cost_astar + + +@pytest.mark.parametrize("use_cpp", [False, True]) +@pytest.mark.parametrize("stamp", [0, -1, 1700000000123456789]) +@pytest.mark.parametrize("rotated", [False, True]) +def test_generated_astar_routes_around_obstacles_and_retains_exact_metadata( + use_cpp, stamp, rotated +): + if use_cpp: + assert _USE_CPP, "Acceptance requires the actual compiled C++ extension" + cells = np.zeros((6, 7), dtype=np.int8) + cells[:5, 3] = 100 + origin = Pose( + position=Point(x=10, y=-20, z=0), + orientation=Quaternion( + x=0, + y=0, + z=math.sin(math.pi / 4) if rotated else 0, + w=math.cos(math.pi / 4) if rotated else 1, + ), + ) + grid = occupancy_from_array( + cells, + header=Header(frame_id="map", stamp=time_from_nanoseconds(stamp)), + resolution=0.5, + origin=origin, + ) + grid = decode(encode(grid), type(grid)) + before = deepcopy(grid) + start = grid_to_world(grid, (1.25, 1.25)) + goal = grid_to_world(grid, (5.25, 1.25)) + + result = min_cost_astar(grid, goal, start, use_cpp=use_cpp) + + assert type(result) is Path + assert result.header == grid.header + assert all(p.header == grid.header and p.pose.orientation.w == 1 for p in result.poses) + coordinates = [ + tuple(round(v) for v in world_to_grid(grid, p.pose.position)) for p in result.poses + ] + assert coordinates[0] == (1, 1) + assert coordinates[-1] == (5, 1) + assert (3, 5) in coordinates + assert all(cells[y, x] != 100 for x, y in coordinates) + assert all(max(abs(a[0] - b[0]), abs(a[1] - b[1])) == 1 for a, b in pairwise(coordinates)) + assert decode(encode(result), Path) == result + result.header.frame_id = "changed" + result.poses[0].header.frame_id = "changed" + assert result.poses[-1].header == before.header + assert grid.header == before.header and grid.info == before.info + np.testing.assert_array_equal(grid.data, before.data) + + +@pytest.mark.parametrize("use_cpp", [False, True]) +def test_generated_astar_unknown_penalty_and_unreachable_outside_goal(use_cpp): + cells = np.array([[0, -1, 0]], dtype=np.int8) + grid = occupancy_from_array( + cells, header=Header(frame_id="map", stamp=time_from_nanoseconds(0)), resolution=1 + ) + assert min_cost_astar(grid, (2, 0), unknown_penalty=1, use_cpp=use_cpp) is None + path = min_cost_astar(grid, (2, 0), unknown_penalty=0.1, use_cpp=use_cpp) + assert path is not None + assert [p.pose.position.x for p in path.poses] == [0, 1, 2] + assert min_cost_astar(grid, (3, 0), use_cpp=use_cpp) is None + + +@pytest.mark.parametrize("use_cpp", [False, True]) +def test_generated_and_existing_astar_share_identical_cell_route(use_cpp): + cells = np.zeros((7, 7), dtype=np.int8) + cells[:5, 3] = 100 + generated = occupancy_from_array( + cells, header=Header(frame_id="map", stamp=time_from_nanoseconds(-1)), resolution=0.5 + ) + legacy = LegacyGrid(cells, resolution=0.5, frame_id="map", ts=0) + old = min_cost_astar(legacy, (2.5, 0.5), (0.5, 0.5), use_cpp=use_cpp) + new = min_cost_astar(generated, (2.5, 0.5), (0.5, 0.5), use_cpp=use_cpp) + assert old is not None and new is not None + assert [(p.x, p.y, p.z) for p in old.poses] == [ + (p.pose.position.x, p.pose.position.y, p.pose.position.z) for p in new.poses + ] + + +@pytest.mark.parametrize("strategy", ["simple", "mixed"]) +@pytest.mark.parametrize("gradient_strategy", ["gradient", "voronoi"]) +def test_navigation_costmap_generation_matches_existing_math_and_preserves_metadata( + strategy, gradient_strategy +): + cells = np.zeros((10, 10), dtype=np.int8) + cells[4:7, 4:7] = 100 + cells[0, :] = -1 + generated = occupancy_from_array( + cells, header=Header(frame_id="map", stamp=time_from_nanoseconds(-1)), resolution=0.5 + ) + generated.info.map_load_time = time_from_nanoseconds(17) + old = make_navigation_map(LegacyGrid(cells, resolution=0.5), 0.5, strategy, gradient_strategy) + new = make_navigation_map(generated, 0.5, strategy, gradient_strategy) + assert new.header == generated.header and new.info == generated.info + np.testing.assert_array_equal(occupancy_view(new), old.grid) + np.testing.assert_array_equal(occupancy_view(generated), cells) From 1c74b424121aa357172fbaa7f2d87ef2e1892765 Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Tue, 20 Oct 2026 05:57:04 +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 +- 6 files changed, 26 insertions(+), 21 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