From 23f44b4eefaf7f7fb649242a7539ac0d80633393 Mon Sep 17 00:00:00 2001 From: cc Date: Mon, 19 Oct 2026 22:14:31 -0700 Subject: [PATCH 1/2] refactor: support generated occupancy values in existing grid algorithms --- .github/workflows/message-codegen.yml | 6 + dimos/mapping/occupancy/extrude_occupancy.py | 17 +-- dimos/mapping/occupancy/gradient.py | 49 +++---- dimos/mapping/occupancy/inflation.py | 17 +-- dimos/mapping/occupancy/operations.py | 64 +++++---- .../occupancy/test_generated_operations.py | 124 ++++++++++++++++++ 6 files changed, 201 insertions(+), 76 deletions(-) create mode 100644 dimos/mapping/occupancy/test_generated_operations.py diff --git a/.github/workflows/message-codegen.yml b/.github/workflows/message-codegen.yml index 00ab0db96b8..89e37a9519d 100644 --- a/.github/workflows/message-codegen.yml +++ b/.github/workflows/message-codegen.yml @@ -19,6 +19,7 @@ on: - dimos/msgs/** - dimos/protocol/tf/** - dimos/memory/** + - dimos/mapping/occupancy/** - dimos/protocol/cdr_mcap.py - docs/usage/transforms.md - dimos/hardware/sensors/camera/** @@ -221,3 +222,8 @@ jobs: .native-review/bin/python -I -c 'import dimos.msgs.occupancy as values; import pytest; import sys; assert "site-packages" in str(values.__file__); raise SystemExit(pytest.main(sys.argv[1:]))' --noconftest --import-mode=importlib -o addopts='' -q "$PWD/dimos/msgs/test_occupancy.py" .native-review/bin/python -I -c 'import dimos.msgs.nav_msgs.OccupancyGrid; import pytest; import sys; raise SystemExit(pytest.main(sys.argv[1:]))' --noconftest --import-mode=importlib -o addopts='' -q "$PWD/dimos/msgs/nav_msgs/test_OccupancyGrid.py" -k block_max_reduce MYPYPATH="$PWD/build/message-source/python" .native-review/bin/python -m mypy --follow-imports=silent dimos/msgs/occupancy.py + - name: Verify occupancy algorithms preserve generated and existing value contracts + if: runner.os == 'Linux' + run: | + .native-review/bin/python -I -c 'from dimos.mapping.occupancy import operations, gradient, inflation, extrude_occupancy; import pytest; import sys; assert all("site-packages" in str(value.__file__) for value in (operations, gradient, inflation, extrude_occupancy)); raise SystemExit(pytest.main(sys.argv[1:]))' --noconftest --import-mode=importlib -o addopts='' -q "$PWD/dimos/mapping/occupancy/test_generated_operations.py" + MYPYPATH="$PWD/build/message-source/python" .native-review/bin/python -m mypy --follow-imports=silent dimos/mapping/occupancy/operations.py dimos/mapping/occupancy/gradient.py dimos/mapping/occupancy/inflation.py dimos/mapping/occupancy/extrude_occupancy.py diff --git a/dimos/mapping/occupancy/extrude_occupancy.py b/dimos/mapping/occupancy/extrude_occupancy.py index 799319cbf6d..c0f05b9a838 100644 --- a/dimos/mapping/occupancy/extrude_occupancy.py +++ b/dimos/mapping/occupancy/extrude_occupancy.py @@ -16,13 +16,14 @@ import numpy as np from numpy.typing import NDArray -from dimos.msgs.nav_msgs.OccupancyGrid import CostValues, OccupancyGrid +from dimos.mapping.occupancy.operations import _grid_values, _GridT +from dimos.msgs.occupancy import CostValues # Rectangle type: (x, y, width, height) Rect = tuple[int, int, int, int] -def identify_convex_shapes(occupancy_grid: OccupancyGrid) -> list[Rect]: +def identify_convex_shapes(occupancy_grid: _GridT) -> list[Rect]: """Identify occupied zones and decompose them into convex rectangles. This function finds all occupied cells in the occupancy grid and @@ -36,7 +37,7 @@ def identify_convex_shapes(occupancy_grid: OccupancyGrid) -> list[Rect]: Returns: List of rectangles as (x, y, width, height) tuples in grid coords. """ - grid = occupancy_grid.grid + grid = _grid_values(occupancy_grid) # Create binary mask of occupied cells (treat UNKNOWN as OCCUPIED) occupied_mask = ((grid == CostValues.OCCUPIED) | (grid == CostValues.UNKNOWN)).astype( @@ -166,7 +167,7 @@ def _largest_rect_in_histogram( def generate_mujoco_scene( - occupancy_grid: OccupancyGrid, + occupancy_grid: _GridT, ) -> str: """Generate a MuJoCo scene XML from an occupancy grid. @@ -184,9 +185,9 @@ def generate_mujoco_scene( # Get rectangles from the occupancy grid rectangles = identify_convex_shapes(occupancy_grid) - resolution = occupancy_grid.resolution - origin_x = occupancy_grid.origin.position.x - origin_y = occupancy_grid.origin.position.y + resolution = occupancy_grid.info.resolution + origin_x = occupancy_grid.info.origin.position.x + origin_y = occupancy_grid.info.origin.position.y # Build XML xml_lines = [ @@ -208,7 +209,7 @@ def generate_mujoco_scene( # Add each rectangle as a box geom for i, (gx, gy, gw, gh) in enumerate(rectangles): # Convert grid coordinates to world coordinates - # Grid origin is top-left, world origin is at occupancy_grid.origin + # Grid origin is top-left, world origin is at occupancy_grid.info.origin # gx, gy are in grid cells, need to convert to meters world_x = origin_x + (gx + gw / 2) * resolution world_y = origin_y + (gy + gh / 2) * resolution diff --git a/dimos/mapping/occupancy/gradient.py b/dimos/mapping/occupancy/gradient.py index 802b52a9faf..d198af9bcbf 100644 --- a/dimos/mapping/occupancy/gradient.py +++ b/dimos/mapping/occupancy/gradient.py @@ -17,7 +17,8 @@ import numpy as np from scipy import ndimage -from dimos.msgs.nav_msgs.OccupancyGrid import CostValues, OccupancyGrid +from dimos.mapping.occupancy.operations import _grid_values, _GridT, _with_grid +from dimos.msgs.occupancy import CostValues if TYPE_CHECKING: from numpy.typing import NDArray @@ -27,8 +28,8 @@ def gradient( - occupancy_grid: OccupancyGrid, obstacle_threshold: int = 50, max_distance: float = 2.0 -) -> OccupancyGrid: + occupancy_grid: _GridT, obstacle_threshold: int = 50, max_distance: float = 2.0 +) -> _GridT: """Create a gradient OccupancyGrid for path planning. Creates a gradient where free space has value 0 and values increase near obstacles. @@ -49,19 +50,19 @@ def gradient( """ # Remember which cells are unknown - unknown_mask = occupancy_grid.grid == CostValues.UNKNOWN + unknown_mask = _grid_values(occupancy_grid) == CostValues.UNKNOWN # Create binary obstacle map # Consider cells >= threshold as obstacles (1), everything else as free (0) # Unknown cells are not considered obstacles for distance calculation - obstacle_map = (occupancy_grid.grid >= obstacle_threshold).astype(np.float32) + obstacle_map = (_grid_values(occupancy_grid) >= obstacle_threshold).astype(np.float32) # Compute distance transform (distance to nearest obstacle in cells) # Unknown cells are treated as if they don't exist for distance calculation distance_cells = cast("NDArray[np.float64]", ndimage.distance_transform_edt(1 - obstacle_map)) # Convert to meters and clip to max distance - distance_meters = np.clip(distance_cells * occupancy_grid.resolution, 0, max_distance) + distance_meters = np.clip(distance_cells * occupancy_grid.info.resolution, 0, max_distance) # Invert and scale to 0-100 range # Far from obstacles (max_distance) -> 0 @@ -78,20 +79,14 @@ def gradient( gradient_data[unknown_mask] = CostValues.UNKNOWN # Create new OccupancyGrid with gradient - gradient_grid = OccupancyGrid( - grid=gradient_data, - resolution=occupancy_grid.resolution, - origin=occupancy_grid.origin, - frame_id=occupancy_grid.frame_id, - ts=occupancy_grid.ts, - ) + gradient_grid = _with_grid(occupancy_grid, gradient_data) return gradient_grid def voronoi_gradient( - occupancy_grid: OccupancyGrid, obstacle_threshold: int = 50, max_distance: float = 2.0 -) -> OccupancyGrid: + occupancy_grid: _GridT, obstacle_threshold: int = 50, max_distance: float = 2.0 +) -> _GridT: """Create a Voronoi-based gradient OccupancyGrid for path planning. Unlike the regular gradient which can result in suboptimal paths in narrow @@ -122,23 +117,17 @@ def voronoi_gradient( - 100: At obstacles """ # Remember which cells are unknown - unknown_mask = occupancy_grid.grid == CostValues.UNKNOWN + unknown_mask = _grid_values(occupancy_grid) == CostValues.UNKNOWN # Create binary obstacle map - obstacle_map = (occupancy_grid.grid >= obstacle_threshold).astype(np.float32) + obstacle_map = (_grid_values(occupancy_grid) >= obstacle_threshold).astype(np.float32) # Check if there are any obstacles if not np.any(obstacle_map): # No obstacles - everything is free - gradient_data = np.zeros_like(occupancy_grid.grid, dtype=np.int8) + gradient_data = np.zeros_like(_grid_values(occupancy_grid), dtype=np.int8) gradient_data[unknown_mask] = CostValues.UNKNOWN - return OccupancyGrid( - grid=gradient_data, - resolution=occupancy_grid.resolution, - origin=occupancy_grid.origin, - frame_id=occupancy_grid.frame_id, - ts=occupancy_grid.ts, - ) + return _with_grid(occupancy_grid, gradient_data) # Label connected obstacle regions (clusters) # This groups all cells of the same wall/obstacle together @@ -192,7 +181,7 @@ def voronoi_gradient( gradient_values[obstacle_map > 0] = CostValues.OCCUPIED # Apply max_distance clipping - cells beyond max_distance from obstacles get cost 0 - max_distance_cells = max_distance / occupancy_grid.resolution + max_distance_cells = max_distance / occupancy_grid.info.resolution gradient_values[distance_cells > max_distance_cells] = 0 # Convert to int8 @@ -201,10 +190,4 @@ def voronoi_gradient( # Preserve unknown cells gradient_data[unknown_mask] = CostValues.UNKNOWN - return OccupancyGrid( - grid=gradient_data, - resolution=occupancy_grid.resolution, - origin=occupancy_grid.origin, - frame_id=occupancy_grid.frame_id, - ts=occupancy_grid.ts, - ) + return _with_grid(occupancy_grid, gradient_data) diff --git a/dimos/mapping/occupancy/inflation.py b/dimos/mapping/occupancy/inflation.py index 610157ddf62..90ec3c46b95 100644 --- a/dimos/mapping/occupancy/inflation.py +++ b/dimos/mapping/occupancy/inflation.py @@ -15,10 +15,11 @@ import numpy as np from scipy import ndimage -from dimos.msgs.nav_msgs.OccupancyGrid import CostValues, OccupancyGrid +from dimos.mapping.occupancy.operations import _grid_values, _GridT, _with_grid +from dimos.msgs.occupancy import CostValues -def simple_inflate(occupancy_grid: OccupancyGrid, radius: float) -> OccupancyGrid: +def simple_inflate(occupancy_grid: _GridT, radius: float) -> _GridT: """Inflate obstacles by a given radius (binary inflation). Args: radius: Inflation radius in meters @@ -26,10 +27,10 @@ def simple_inflate(occupancy_grid: OccupancyGrid, radius: float) -> OccupancyGri New OccupancyGrid with inflated obstacles """ # Convert radius to grid cells - cell_radius = int(np.ceil(radius / occupancy_grid.resolution)) + cell_radius = int(np.ceil(radius / occupancy_grid.info.resolution)) # Get grid as numpy array - grid_array = occupancy_grid.grid + grid_array = _grid_values(occupancy_grid) # Create circular kernel for binary inflation y, x = np.ogrid[-cell_radius : cell_radius + 1, -cell_radius : cell_radius + 1] @@ -44,10 +45,4 @@ def simple_inflate(occupancy_grid: OccupancyGrid, radius: float) -> OccupancyGri result_grid[inflated] = CostValues.OCCUPIED # Create new OccupancyGrid with inflated data using numpy constructor - return OccupancyGrid( - grid=result_grid, - resolution=occupancy_grid.resolution, - origin=occupancy_grid.origin, - frame_id=occupancy_grid.frame_id, - ts=occupancy_grid.ts, - ) + return _with_grid(occupancy_grid, result_grid) diff --git a/dimos/mapping/occupancy/operations.py b/dimos/mapping/occupancy/operations.py index 8609ac634b3..6b9107f9f0c 100644 --- a/dimos/mapping/occupancy/operations.py +++ b/dimos/mapping/occupancy/operations.py @@ -12,15 +12,43 @@ # See the License for the specific language governing permissions and # limitations under the License. +from copy import deepcopy +from typing import TypeVar + import numpy as np +from numpy.typing import NDArray from scipy import ndimage -from dimos.msgs.nav_msgs.OccupancyGrid import CostValues, OccupancyGrid +from dimos.messages.nav_msgs.msg import OccupancyGrid as CdrOccupancyGrid +from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid +from dimos.msgs.occupancy import CostValues, occupancy_view + +_GridT = TypeVar("_GridT", OccupancyGrid, CdrOccupancyGrid) + + +def _grid_values(message: _GridT) -> NDArray[np.int8]: + if isinstance(message, CdrOccupancyGrid): + return occupancy_view(message) + return message.grid + + +def _with_grid(message: _GridT, cells: NDArray[np.int8]) -> _GridT: + if isinstance(message, CdrOccupancyGrid): + if cells.shape != (message.info.height, message.info.width): + raise ValueError("Replacement occupancy cells must retain map dimensions") + result = deepcopy(message) + result.data = np.array(cells, dtype=np.int8, copy=True).reshape(-1) + return result + return OccupancyGrid( + grid=cells, + resolution=message.resolution, + origin=message.origin, + frame_id=message.frame_id, + ts=message.ts, + ) -def smooth_occupied( - occupancy_grid: OccupancyGrid, min_neighbor_fraction: float = 0.4 -) -> OccupancyGrid: +def smooth_occupied(occupancy_grid: _GridT, min_neighbor_fraction: float = 0.4) -> _GridT: """Smooth occupied zones by removing unsupported protrusions. Removes occupied cells that don't have sufficient neighboring occupied @@ -33,7 +61,7 @@ def smooth_occupied( Returns: New OccupancyGrid with smoothed occupied zones """ - grid_array = occupancy_grid.grid + grid_array = _grid_values(occupancy_grid) occupied_mask = grid_array >= CostValues.OCCUPIED # Count occupied neighbors for each cell (8-connectivity). @@ -49,16 +77,10 @@ def smooth_occupied( result_grid = grid_array.copy() result_grid[unsupported] = CostValues.FREE - return OccupancyGrid( - grid=result_grid, - resolution=occupancy_grid.resolution, - origin=occupancy_grid.origin, - frame_id=occupancy_grid.frame_id, - ts=occupancy_grid.ts, - ) + return _with_grid(occupancy_grid, result_grid) -def overlay_occupied(base: OccupancyGrid, overlay: OccupancyGrid) -> OccupancyGrid: +def overlay_occupied(base: _GridT, overlay: _GridT) -> _GridT: """Overlay occupied zones from one grid onto another. Marks cells as occupied in the base grid wherever they are occupied @@ -70,19 +92,13 @@ def overlay_occupied(base: OccupancyGrid, overlay: OccupancyGrid) -> OccupancyGr Returns: New OccupancyGrid with combined occupied zones """ - if base.grid.shape != overlay.grid.shape: + if _grid_values(base).shape != _grid_values(overlay).shape: raise ValueError( - f"Grid shapes must match: base {base.grid.shape} vs overlay {overlay.grid.shape}" + f"Grid shapes must match: base {_grid_values(base).shape} vs overlay {_grid_values(overlay).shape}" ) - result_grid = base.grid.copy() - overlay_occupied_mask = overlay.grid >= CostValues.OCCUPIED + result_grid = _grid_values(base).copy() + overlay_occupied_mask = _grid_values(overlay) >= CostValues.OCCUPIED result_grid[overlay_occupied_mask] = CostValues.OCCUPIED - return OccupancyGrid( - grid=result_grid, - resolution=base.resolution, - origin=base.origin, - frame_id=base.frame_id, - ts=base.ts, - ) + return _with_grid(base, result_grid) diff --git a/dimos/mapping/occupancy/test_generated_operations.py b/dimos/mapping/occupancy/test_generated_operations.py new file mode 100644 index 00000000000..764c1fe5c2c --- /dev/null +++ b/dimos/mapping/occupancy/test_generated_operations.py @@ -0,0 +1,124 @@ +# 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 decode, encode +import numpy as np +import pytest + +from dimos.mapping.occupancy.extrude_occupancy import generate_mujoco_scene, identify_convex_shapes +from dimos.mapping.occupancy.gradient import gradient, voronoi_gradient +from dimos.mapping.occupancy.inflation import simple_inflate +from dimos.mapping.occupancy.operations import overlay_occupied, smooth_occupied +from dimos.messages.builtin_interfaces.msg import Time +from dimos.messages.nav_msgs.msg import OccupancyGrid +from dimos.messages.std_msgs.msg import Header +from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid as LegacyGrid +from dimos.msgs.occupancy import occupancy_from_array, occupancy_view + + +@pytest.fixture(params=[(0, 0), (-1, 999999999), (1700000000, 123456789)]) +def grids(request): + cells = np.zeros((7, 7), dtype=np.int8) + cells[0, 0] = -1 + cells[3, 3] = cells[4, 4] = 100 + generated = occupancy_from_array( + cells, + header=Header(frame_id="map", stamp=Time(sec=request.param[0], nanosec=request.param[1])), + resolution=0.25, + ) + generated.info.map_load_time = Time(sec=7, nanosec=9) + return generated, LegacyGrid(grid=cells, resolution=0.25, frame_id="map", ts=0) + + +@pytest.mark.parametrize( + "operation,kwargs", + [ + (gradient, {"max_distance": 0.5}), + (voronoi_gradient, {"max_distance": 0.5}), + (smooth_occupied, {}), + (simple_inflate, {"radius": 0.25}), + ], +) +def test_existing_algorithms_preserve_input_identity_and_exact_generated_metadata( + grids, operation, kwargs +): + generated, legacy = grids + original = deepcopy(generated) + + actual = operation(generated, **kwargs) + old = operation(legacy, **kwargs) + received = decode(encode(actual), OccupancyGrid) + + assert type(actual) is OccupancyGrid + assert type(old) is LegacyGrid + np.testing.assert_array_equal(occupancy_view(received), old.grid) + assert received.header == original.header + assert received.info.map_load_time == Time(sec=7, nanosec=9) + assert received.info.origin == original.info.origin + assert occupancy_view(received)[0, 0] == -1 + np.testing.assert_array_equal(generated.data, original.data) + assert generated.header == original.header + assert generated.info == original.info + actual.header.stamp.nanosec = 0 + actual.info.origin.position.x = 99 + actual.data[0] = 42 + np.testing.assert_array_equal(generated.data, original.data) + assert generated.header == original.header + assert generated.info == original.info + assert legacy.ts == 0 + + +def test_gradient_retains_unknown_obstacles_and_exact_distance_cost(grids): + generated, _ = grids + + actual = gradient(generated, max_distance=0.5) + + assert occupancy_view(actual)[0, 0] == -1 + assert occupancy_view(actual)[3, 3] == 100 + assert occupancy_view(actual)[3, 2] == 50 + assert occupancy_view(actual)[1, 1] == 0 + + +def test_overlay_preserves_base_metadata_and_rejects_different_dimensions(grids): + base, _ = grids + cells = np.zeros((7, 7), dtype=np.int8) + cells[2, 5] = 100 + overlay = occupancy_from_array( + cells, header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)), resolution=0.25 + ) + + actual = overlay_occupied(base, overlay) + + assert occupancy_view(actual)[2, 5] == 100 + assert occupancy_view(actual)[3, 3] == 100 + assert actual.header == base.header + assert actual.info.map_load_time == base.info.map_load_time + other = occupancy_from_array( + np.zeros((2, 2), dtype=np.int8), header=overlay.header, resolution=0.25 + ) + with pytest.raises(ValueError, match="Grid shapes must match"): + overlay_occupied(base, other) + + +def test_generated_and_legacy_extrusion_produce_identical_offline_scene(grids): + generated, legacy = grids + + actual = generate_mujoco_scene(generated) + + assert actual == generate_mujoco_scene(legacy) + assert identify_convex_shapes(generated) == identify_convex_shapes(legacy) + assert '' in actual + assert 'name="wall_0"' in actual From f04b81fd18d86cb51709e4f719fd11117e79b18b 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:15:38 +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