Add offline object inventory and localize perception pipeline DIM1343

dimos/models/base.py

@@ -24,12 +24,21 @@
 from dimos.core.resource import Resource
 from dimos.protocol.service.spec import BaseConfig, Configurable
 
-# Device string type - 'cuda', 'cpu', 'cuda:0', 'cuda:1', etc.
-DeviceType = Annotated[str, "Device identifier (e.g., 'cuda', 'cpu', 'cuda:0')"]
+# Device string type - 'cuda', 'cpu', 'cuda:0', 'cuda:1', 'mps', etc.
+DeviceType = Annotated[str, "Device identifier (e.g., 'cuda', 'cpu', 'cuda:0', 'mps')"]
+
+
+def default_torch_device() -> str:
+    """Best available torch device: CUDA, then Apple Metal (MPS), then CPU."""
+    if torch.cuda.is_available():
+        return "cuda"
+    if torch.backends.mps.is_available():
+        return "mps"
+    return "cpu"
 
 
 class LocalModelConfig(BaseConfig):
-    device: DeviceType = "cuda" if torch.cuda.is_available() else "cpu"
+    device: DeviceType = default_torch_device()
     dtype: torch.dtype = torch.float32
     warmup: bool = False
     autostart: bool = False
@@ -54,8 +63,8 @@ def __init__(self, **kwargs: object) -> None:
         """Initialize local model with device and dtype configuration.
 
         Args:
-            device: Device to run on ('cuda', 'cpu', 'cuda:0', etc.).
-                    Auto-detects CUDA availability if None.
+            device: Device to run on ('cuda', 'cpu', 'cuda:0', 'mps', etc.).
+                    Defaults to the best available device (CUDA, then MPS, then CPU).
             dtype: Model dtype (torch.float16, torch.bfloat16, etc.).
                    Uses class _default_dtype if None.
             autostart: If True, immediately load the model.
@@ -109,6 +118,8 @@ def stop(self) -> None:
         gc.collect()
         if self.config.device.startswith("cuda") and torch.cuda.is_available():
             torch.cuda.empty_cache()
+        elif self.config.device.startswith("mps") and torch.backends.mps.is_available():
+            torch.mps.empty_cache()
 
     def _ensure_cuda_initialized(self) -> None:
         """Initialize CUDA context to prevent cuBLAS allocation failures.

dimos/models/embedding/siglip.py

@@ -0,0 +1,101 @@
+# 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 __future__ import annotations
+
+from functools import cached_property
+from typing import overload
+
+from PIL import Image as PILImage
+import torch
+import torch.nn.functional as functional
+from transformers import SiglipModel as HFSiglipModel, SiglipProcessor
+
+from dimos.models.base import HuggingFaceModel
+from dimos.models.embedding.base import Embedding, EmbeddingModel, HuggingFaceEmbeddingModelConfig
+from dimos.msgs.sensor_msgs.Image import Image
+
+
+class SigLIPModelConfig(HuggingFaceEmbeddingModelConfig):
+    model_name: str = "google/siglip-base-patch16-224"
+    dtype: torch.dtype = torch.float32
+
+
+class SigLIPModel(EmbeddingModel, HuggingFaceModel):
+    """SigLIP image/text embedding model for frame-level semantic retrieval.
+
+    Weights come from the Hugging Face hub cache: a miss downloads once, a
+    hit reuses the cache. Text inputs use ``padding="max_length"`` - SigLIP
+    was trained with max-length padded text, and unpadded prompts measurably
+    degrade the text-image alignment.
+    """
+
+    config: SigLIPModelConfig
+    _model_class = HFSiglipModel
+
+    @cached_property
+    def _model(self) -> HFSiglipModel:
+        self._ensure_cuda_initialized()
+        return HFSiglipModel.from_pretrained(self.config.model_name).eval().to(self.config.device)
+
+    @cached_property
+    def _processor(self) -> SiglipProcessor:
+        return SiglipProcessor.from_pretrained(self.config.model_name, use_fast=True)
+
+    @property
+    def logit_scale(self) -> float:
+        """Trained sigmoid temperature: exp of the model's raw logit scale."""
+        return float(self._model.logit_scale.exp())
+
+    @overload
+    def embed(self, image: Image, /) -> Embedding: ...
+    @overload
+    def embed(self, *images: Image) -> list[Embedding]: ...
+    def embed(self, *images: Image) -> Embedding | list[Embedding]:
+        """Embed one or more images into the shared image-text space."""
+        pil_images = [PILImage.fromarray(img.to_rgb().data) for img in images]
+
+        with torch.inference_mode():
+            inputs = self._processor(images=pil_images, return_tensors="pt").to(self.config.device)
+            image_features = self._model.get_image_features(**inputs)
+            if self.config.normalize:
+                image_features = functional.normalize(image_features, dim=-1)
+
+        embeddings = [
+            Embedding(vector=feat, timestamp=images[i].ts) for i, feat in enumerate(image_features)
+        ]
+        return embeddings[0] if len(images) == 1 else embeddings
+
+    @overload
+    def embed_text(self, text: str, /) -> Embedding: ...
+    @overload
+    def embed_text(self, *texts: str) -> list[Embedding]: ...
+    def embed_text(self, *texts: str) -> Embedding | list[Embedding]:
+        """Embed one or more text strings into the shared image-text space."""
+        with torch.inference_mode():
+            inputs = self._processor(
+                text=list(texts), return_tensors="pt", padding="max_length", truncation=True
+            ).to(self.config.device)
+            text_features = self._model.get_text_features(**inputs)
+            if self.config.normalize:
+                text_features = functional.normalize(text_features, dim=-1)
+
+        embeddings = [Embedding(vector=feat) for feat in text_features]
+        return embeddings[0] if len(texts) == 1 else embeddings
+
+    def stop(self) -> None:
+        """Release model and free GPU memory."""
+        if "_processor" in self.__dict__:
+            del self.__dict__["_processor"]
+        super().stop()

dimos/models/segmentation/edge_tam.py

@@ -13,7 +13,8 @@
 # limitations under the License.
 
 from collections.abc import Generator
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
+from functools import lru_cache
 from pathlib import Path
 import shutil
 import tempfile
@@ -26,6 +27,7 @@
 from PIL import Image as PILImage
 import torch
 
+from dimos.models.base import default_torch_device
 from dimos.msgs.sensor_msgs.Image import Image
 from dimos.perception.detection.detectors.base import Detector
 from dimos.perception.detection.type.detection2d.bbox import Detection2DBBox
@@ -40,6 +42,20 @@
 
 logger = setup_logger()
 
+# Hole/sprinkle cleanup in SAM2 needs the optional CUDA extension sam2._C.
+# Requesting it without _C triggers a per-call UserWarning and then skips.
+_FILL_HOLE_AREA = 8
+_MIN_MASK_REGION_AREA = 150
+
+
+@lru_cache(maxsize=1)
+def _sam2_cuda_postprocess_available() -> bool:
+    try:
+        from sam2 import _C  # noqa: F401
+    except ImportError:
+        return False
+    return True
+
 
 class SAM2InferenceState(TypedDict):
     images: list[torch.Tensor | None]
@@ -54,17 +70,19 @@ def _build_model() -> "SAM2VideoPredictor":
     if not local_config_path.exists():
         raise FileNotFoundError(f"EdgeTAM config not found at {local_config_path}")
 
-    if not torch.cuda.is_available():
-        raise RuntimeError("EdgeTAM requires a CUDA-capable GPU")
+    device = default_torch_device()
+    if device == "cpu":
+        raise RuntimeError("EdgeTAM requires a CUDA or MPS device")
 
     cfg = OmegaConf.load(local_config_path)
 
+    fill_hole_area = _FILL_HOLE_AREA if _sam2_cuda_postprocess_available() else 0
     overrides = {
         "model.sam_mask_decoder_extra_args.dynamic_multimask_via_stability": True,
         "model.sam_mask_decoder_extra_args.dynamic_multimask_stability_delta": 0.05,
         "model.sam_mask_decoder_extra_args.dynamic_multimask_stability_thresh": 0.98,
         "model.binarize_mask_from_pts_for_mem_enc": True,
-        "model.fill_hole_area": 8,
+        "model.fill_hole_area": fill_hole_area,
     }
 
     for key, value in overrides.items():
@@ -90,7 +108,7 @@ def _build_model() -> "SAM2VideoPredictor":
     if unexpected_keys:
         raise RuntimeError("Unexpected keys in checkpoint")
 
-    predictor = predictor.to("cuda")
+    predictor = predictor.to(device)
     predictor.eval()
     return predictor
 
@@ -102,6 +120,59 @@ def __init__(self) -> None:
         from sam2.sam2_image_predictor import SAM2ImagePredictor
 
         self._predictor = SAM2ImagePredictor(_build_model())
+        self._mask_generators: dict[int, Any] = {}
+
+    def propose_all(
+        self,
+        image: Image,
+        points_per_side: int = 24,
+        pred_iou_thresh: float = 0.55,
+        stability_score_thresh: float = 0.85,
+        min_mask_region_area: int = _MIN_MASK_REGION_AREA,
+    ) -> ImageDetections2D:
+        """Class-agnostic mask proposals over the whole image, no vocabulary.
+
+        Wraps ``SAM2AutomaticMaskGenerator`` on the shared EdgeTAM model.
+        Proposals carry no names ("proposal") and ``confidence`` is the
+        predicted mask IoU.
+        """
+        import cv2
+        from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator
+
+        if not _sam2_cuda_postprocess_available():
+            min_mask_region_area = 0
+
+        generator = self._mask_generators.get(points_per_side)
+        if generator is None:
+            generator = SAM2AutomaticMaskGenerator(
+                model=self._predictor.model,
+                points_per_side=points_per_side,
+                pred_iou_thresh=pred_iou_thresh,
+                stability_score_thresh=stability_score_thresh,
+                min_mask_region_area=min_mask_region_area,
+            )
+            self._mask_generators[points_per_side] = generator
+
+        rgb = cv2.cvtColor(image.to_opencv(), cv2.COLOR_BGR2RGB)
+        amp = (
+            torch.autocast("cuda", dtype=torch.bfloat16)
+            if self._predictor.device.type == "cuda"
+            else nullcontext()
+        )
+        with torch.no_grad(), amp:
+            proposals = generator.generate(rgb)
+
+        detections: list[Detection2DBBox] = [
+            Detection2DSeg.from_sam2_result(
+                proposal["segmentation"],
+                i,
+                image,
+                name="proposal",
+                confidence=float(proposal.get("predicted_iou", 1.0)),
+            )
+            for i, proposal in enumerate(proposals)
+        ]
+        return ImageDetections2D(image, [det for det in detections if det.is_valid()])
 
     def segment(self, detections: ImageDetections2D) -> ImageDetections2D:
         """Refine box detections into mask detections (Detection2DSeg)."""
@@ -113,7 +184,13 @@ def segment(self, detections: ImageDetections2D) -> ImageDetections2D:
         image = detections.image
         rgb = cv2.cvtColor(image.to_opencv(), cv2.COLOR_BGR2RGB)
 
-        with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
+        # bfloat16 autocast is CUDA-only; MPS runs in float32
+        amp = (
+            torch.autocast("cuda", dtype=torch.bfloat16)
+            if self._predictor.device.type == "cuda"
+            else nullcontext()
+        )
+        with torch.no_grad(), amp:
             self._predictor.set_image(rgb)
             boxes = np.array([det.bbox for det in detections], dtype=np.float32)
             masks, _, _ = self._predictor.predict(box=boxes, multimask_output=False)
@@ -168,7 +245,7 @@ def _prepare_frame(self, image: Image) -> torch.Tensor:
         img_np /= img_std
 
         img_tensor = torch.from_numpy(img_np).permute(2, 0, 1).float()
-        img_tensor = img_tensor.cuda()
+        img_tensor = img_tensor.to(self._predictor.device)
 
         return img_tensor
 

dimos/models/test_base.py

@@ -16,9 +16,10 @@
 
 from functools import cached_property
 
+import pytest
 import torch
 
-from dimos.models.base import HuggingFaceModel, LocalModel
+from dimos.models.base import HuggingFaceModel, LocalModel, default_torch_device
 
 
 class ConcreteLocalModel(LocalModel):
@@ -37,11 +38,17 @@ def _model(self) -> str:
         return f"hf_model:{self.model_name}"
 
 
-def test_local_model_device_auto_detection() -> None:
-    """Test that device is auto-detected based on CUDA availability."""
-    model = ConcreteLocalModel()
-    expected = "cuda" if torch.cuda.is_available() else "cpu"
-    assert model.device == expected
+def test_default_torch_device_priority(monkeypatch: pytest.MonkeyPatch) -> None:
+    """CUDA wins over MPS, MPS wins over CPU, CPU is the last resort."""
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.backends.mps, "is_available", lambda: True)
+    assert default_torch_device() == "cuda"
+
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+    assert default_torch_device() == "mps"
+
+    monkeypatch.setattr(torch.backends.mps, "is_available", lambda: False)
+    assert default_torch_device() == "cpu"
 
 
 def test_local_model_explicit_device() -> None: