refactor(compressors): consolidate pipeline and improve mask handling

mini-nav/benchmarks/tasks/retrieval.py

@@ -60,7 +60,7 @@ def _establish_eval_database(
                 {
                     "id": global_idx + j,
                     "label": labels_list[j],
-                    "vector": all_features[global_idx + j].numpy(),
+                    "vector": all_features[global_idx + j].detach().cpu().numpy(),
                 }
                 for j in range(batch_size)
             ]

mini-nav/compressors/__init__.py

@@ -1,6 +1,12 @@
-from .common import BinarySign, bits_to_hash, hamming_distance, hamming_similarity, hash_to_bits
+from .common import (
+    BinarySign,
+    bits_to_hash,
+    hamming_distance,
+    hamming_similarity,
+    hash_to_bits,
+)
 from .hash_compressor import HashCompressor, HashLoss, VideoPositiveMask
-from .pipeline import HashPipeline, SAMHashPipeline, create_pipeline_from_config
+from .pipeline import HashPipeline, create_pipeline_from_config
 from .train import train
 
 __all__ = [
@@ -9,7 +15,6 @@ __all__ = [
     "HashLoss",
     "VideoPositiveMask",
     "HashPipeline",
-    "SAMHashPipeline",  # Backward compatibility alias
     "create_pipeline_from_config",
     "BinarySign",
     "hamming_distance",

mini-nav/compressors/pipeline.py

@@ -1,14 +1,5 @@
 """SAM + DINO + Hash compression pipeline."""
 
-from utils import get_device
-from utils.model import (
-    get_dino_dim,
-    load_dino_model,
-    load_hash_compressor,
-    load_sam_model,
-)
-from utils.image import extract_masked_region, segment_image
-
 from typing import Optional
 
 import torch
@@ -16,15 +7,24 @@ import torch.nn as nn
 import torch.nn.functional as F
 from PIL import Image
 
+from utils import get_device
+from utils.image import extract_masked_region, segment_image
+from utils.model import (
+    get_dino_dim,
+    load_dino_model,
+    load_hash_compressor,
+    load_sam_model,
+)
+
 
 def create_pipeline_from_config(config) -> "HashPipeline":
     """Create HashPipeline from a config object.
 
     Args:
-        config: Configuration object with model settings
+        config: Configuration object with model settings.
 
     Returns:
-        Initialized HashPipeline
+        Initialized HashPipeline.
     """
     return HashPipeline(
         dino_model=config.model.dino_model,
@@ -38,21 +38,15 @@ def create_pipeline_from_config(config) -> "HashPipeline":
 
 
 class HashPipeline(nn.Module):
-    """Pipeline: SAM segmentation + DINO features + Hash compression.
+    """Pipeline for SAM segmentation + DINO features + Hash compression.
 
     Pipeline flow:
         PIL Image -> SAM (largest object mask) -> DINO (features) -> Hash (binary codes)
 
-    Usage:
-        # Initialize with config
-        pipeline = HashPipeline(
-            dino_model="facebook/dinov2-large",
-            hash_bits=512,
-        )
-
-        # Process image
+    Example:
+        pipeline = HashPipeline(dino_model="facebook/dinov2-large", hash_bits=512)
         image = Image.open("path/to/image.jpg")
-        hash_bits = pipeline(image)  # [1, 512] binary bits
+        hash_bits = pipeline(image)  # Returns [1, 512] binary bits
     """
 
     def __init__(
@@ -65,38 +59,25 @@ class HashPipeline(nn.Module):
         hash_bits: int = 512,
         compressor_path: Optional[str] = None,
     ):
-        """Initialize the pipeline.
-
-        Args:
-            dino_model: DINOv2 model name from HuggingFace
-            sam_model: SAM2.1 model name from HuggingFace
-            sam_min_mask_area: Minimum area threshold for valid SAM masks
-            sam_max_masks: Maximum number of SAM masks to keep
-            sam_points_per_batch: Prompt points batch size for SAM2 mask generation
-            sam_checkpoint_dir: Optional local cache directory for SAM2 weights
-            hash_bits: Number of bits in hash code
-            compressor_path: Optional path to trained HashCompressor weights
-            device: Device to run models on
-        """
         super().__init__()
 
-        # Auto detect device
+        # Device for model placement.
         self.device = get_device()
 
-        self.dino_model = dino_model
+        # SAM2 settings.
         self.sam_model_name = sam_model
         self.sam_min_mask_area = sam_min_mask_area
         self.sam_max_masks = sam_max_masks
         self.sam_points_per_batch = sam_points_per_batch
 
+        # Load models.
         self.mask_generator = load_sam_model(model_name=sam_model)
-
         self.processor, self.dino = load_dino_model(model_name=dino_model)
 
-        # Determine DINO feature dimension
+        # DINO feature dimension based on model size.
         self.dino_dim = get_dino_dim(dino_model)
 
-        # Initialize HashCompressor
+        # Hash compressor for binarizing DINO features.
         self.hash_compressor = load_hash_compressor(
             input_dim=self.dino_dim,
             hash_bits=hash_bits,
@@ -104,18 +85,21 @@ class HashPipeline(nn.Module):
         )
 
     @property
-    def hash_bits(self):
-        """Return the number of hash bits."""
+    def hash_bits(self) -> int:
+        """Number of bits in the hash code."""
         return self.hash_compressor.hash_bits
 
-    def _prepare_image_for_encoding(
-        self,
-        image: Image.Image,
-        apply_sam: bool,
-    ) -> Image.Image:
-        if not apply_sam:
-            return image
+    def _segment_with_sam(self, image: Image.Image) -> Image.Image:
+        """Segment image with SAM and extract the largest object mask.
 
+        If no valid masks are found, returns the original image.
+
+        Args:
+            image: Input PIL Image.
+
+        Returns:
+            Masked image containing only the largest object, or original if no masks.
+        """
         masks = segment_image(
             self.mask_generator,
             image,
@@ -128,62 +112,45 @@ class HashPipeline(nn.Module):
 
         return extract_masked_region(image, masks[0]["segment"])
 
-    def _encode_image(self, image: Image.Image, apply_sam: bool) -> torch.Tensor:
-        image_for_encoding = self._prepare_image_for_encoding(
-            image, apply_sam=apply_sam
-        )
-        inputs = self.processor(image_for_encoding, return_tensors="pt").to(self.device)
+    def _dino_forward(self, image: Image.Image) -> torch.Tensor:
+        """Extract DINO tokens from an image.
+
+        Args:
+            image: Input PIL Image.
+
+        Returns:
+            Last hidden state tokens of shape [1, N, dim].
+        """
+        inputs = self.processor(image, return_tensors="pt").to(self.device)
 
         with torch.no_grad():
             outputs = self.dino(**inputs)
-            tokens = outputs.last_hidden_state
+            return outputs.last_hidden_state
 
+    def forward(self, image: Image.Image) -> torch.Tensor:
+        """Process a single image through the full pipeline.
+
+        Args:
+            image: Input PIL Image.
+
+        Returns:
+            Binary hash codes of shape [1, hash_bits] as int32.
+        """
+        image = self._segment_with_sam(image)
+        tokens = self._dino_forward(image)
         _, _, bits = self.hash_compressor(tokens)
         return bits
 
-    def forward(self, image: Image.Image) -> torch.Tensor:
-        """Process a single image through the pipeline.
-
-        Args:
-            image: Input PIL Image
-
-        Returns:
-            Binary hash codes [1, hash_bits] as int32
-        """
-        return self._encode_image(image, apply_sam=True)
-
-    def encode_masked_region(self, image: Image.Image) -> torch.Tensor:
-        """Encode a pre-masked region using DINO+Hash without SAM stage."""
-        return self._encode_image(image, apply_sam=False)
-
-    def encode(self, image: Image.Image) -> torch.Tensor:
-        """Encode an image to binary hash bits.
-
-        Alias for forward().
-
-        Args:
-            image: Input PIL Image
-
-        Returns:
-            Binary hash codes [1, hash_bits] as int32
-        """
-        return self.forward(image)
-
     def extract_features(self, image: Image.Image) -> torch.Tensor:
-        """Extract DINO features from an image.
+        """Extract normalized DINO features from an image.
 
         Args:
-            image: Input PIL Image
+            image: Input PIL Image.
 
         Returns:
-            DINO features [1, dino_dim], normalized
+            Normalized DINO features of shape [1, dino_dim].
         """
-        image_for_encoding = self._prepare_image_for_encoding(image, apply_sam=True)
-        inputs = self.processor(image_for_encoding, return_tensors="pt").to(self.device)
-
-        with torch.no_grad():
-            outputs = self.dino(**inputs)
-            features = outputs.last_hidden_state.mean(dim=1)  # [1, dim]
-            features = F.normalize(features, dim=-1)
-
-        return features
+        image = self._segment_with_sam(image)
+        tokens = self._dino_forward(image)
+        features = tokens.mean(dim=1)
+        return F.normalize(features, dim=-1)

mini-nav/configs/config.yaml

@@ -1,7 +1,7 @@
 model:
   dino_model: "facebook/dinov2-large"
   compression_dim: 512
-  device: "auto" # auto-detect GPU
+  device: "cuda:3" # auto-detect GPU
   sam_model: "facebook/sam2.1-hiera-large" # SAM model name
   sam_min_mask_area: 100 # Minimum mask area threshold
   sam_max_masks: 10 # Maximum number of masks to keep

mini-nav/configs/models.py

@@ -18,7 +18,12 @@ class ModelConfig(BaseModel):
     compression_dim: int = Field(
         default=512, gt=0, description="Output feature dimension"
     )
-    device: str = "auto"
+    device: str = Field(
+        default="auto",
+        description=(
+            "Device to use for model inference (e.g., 'cuda:1,3', 'auto', 'cpu')"
+        ),
+    )
     sam_model: str = Field(
         default="facebook/sam2.1-hiera-large",
         description="SAM model name from HuggingFace",

mini-nav/feature_retrieval.py

@@ -102,7 +102,7 @@ class FeatureRetrieval:
                     {
                         "id": i,
                         "label": batch_label,
-                        "vector": cls_tokens[i].numpy(),
+                        "vector": cls_tokens[i].detach().cpu().numpy(),
                         "binary": pil_image_to_bytes(images[i]),
                     }
                 ]

mini-nav/tests/test_compressors.py

@@ -7,7 +7,6 @@ from compressors import (
     BinarySign,
     HashCompressor,
     HashPipeline,
-    SAMHashPipeline,
     VideoPositiveMask,
     bits_to_hash,
     create_pipeline_from_config,
@@ -257,10 +256,6 @@ class TestHashPipeline:
         pipeline = HashPipeline(hash_bits=256)
         assert pipeline.hash_bits == 256
 
-    def test_pipeline_alias(self):
-        """Verify SAMHashPipeline is alias for HashPipeline."""
-        assert SAMHashPipeline is HashPipeline
-
 
 class TestConfigIntegration:
     """Test suite for config integration with pipeline."""

mini-nav/tests/test_image_utils.py

@@ -0,0 +1,82 @@
+from unittest.mock import Mock
+
+import torch
+from PIL import Image
+
+from utils.image import segment_image
+
+
+def test_segment_image_passes_pil_image_to_mask_generator() -> None:
+    mock_generator = Mock(return_value={"masks": []})
+
+    segment_image(
+        mock_generator,
+        Image.new("RGBA", (16, 16), color=(255, 0, 0, 255)),
+        points_per_batch=32,
+    )
+
+    image_arg = mock_generator.call_args.args[0]
+    assert isinstance(image_arg, Image.Image)
+    assert image_arg.mode == "RGB"
+    assert mock_generator.call_args.kwargs["points_per_batch"] == 32
+
+
+def test_segment_image_supports_tensor_masks_output() -> None:
+    masks_tensor = torch.tensor(
+        [
+            [
+                [1, 1, 0],
+                [1, 1, 0],
+                [0, 0, 0],
+            ],
+            [
+                [1, 1, 1],
+                [1, 1, 1],
+                [1, 1, 1],
+            ],
+        ],
+        dtype=torch.float32,
+    )
+    mock_generator = Mock(return_value={"masks": masks_tensor})
+
+    result = segment_image(
+        mock_generator,
+        Image.new("RGB", (3, 3), color=(0, 0, 0)),
+        min_area=3,
+        max_masks=5,
+    )
+
+    assert len(result) == 2
+    assert result[0]["area"] == 9
+    assert result[0]["bbox"] == [0, 0, 3, 3]
+    assert result[1]["area"] == 4
+    assert result[1]["bbox"] == [0, 0, 2, 2]
+
+
+def test_segment_image_filters_tensor_masks_by_min_area() -> None:
+    masks_tensor = torch.tensor(
+        [
+            [
+                [1, 0, 0],
+                [0, 0, 0],
+                [0, 0, 0],
+            ],
+            [
+                [1, 1, 0],
+                [1, 1, 0],
+                [0, 0, 0],
+            ],
+        ],
+        dtype=torch.float32,
+    )
+    mock_generator = Mock(return_value={"masks": masks_tensor})
+
+    result = segment_image(
+        mock_generator,
+        Image.new("RGB", (3, 3), color=(0, 0, 0)),
+        min_area=2,
+        max_masks=5,
+    )
+
+    assert len(result) == 1
+    assert result[0]["area"] == 4

mini-nav/utils/image.py

@@ -9,7 +9,7 @@ def segment_image(
     image: Image.Image,
     min_area: int = 32 * 32,
     max_masks: int = 5,
-    points_per_batch=64,
+    points_per_batch: int = 64,
 ) -> list[dict[str, Any]]:
     """Segment image using SAM to extract object masks.
 
@@ -27,26 +27,104 @@ def segment_image(
             - predicted_iou: Model's confidence in the mask
             - stability_score: Stability score for the mask
     """
-    # Convert PIL Image to numpy array
-    image_np = np.array(image.convert("RGB"))
+    image_rgb = image.convert("RGB")
+    raw_output = mask_generator(image_rgb, points_per_batch=points_per_batch)
+    raw_masks = raw_output.get("masks", raw_output)
 
-    # Generate masks
-    masks = mask_generator(image_np, points_per_batch=points_per_batch)["masks"]
+    normalized_masks: list[dict[str, Any]] = []
 
-    if not masks:
+    if isinstance(raw_masks, list):
+        if raw_masks and isinstance(raw_masks[0], dict):
+            normalized_masks = raw_masks
+        else:
+            for mask_like in raw_masks:
+                mask_dict = _to_mask_dict(mask_like)
+                if mask_dict is not None:
+                    normalized_masks.append(mask_dict)
+    else:
+        mask_array = _to_numpy_mask_array(raw_masks)
+        if mask_array is not None:
+            if mask_array.ndim == 2:
+                mask_array = np.expand_dims(mask_array, axis=0)
+            if mask_array.ndim == 3:
+                for single_mask in mask_array:
+                    mask_dict = _to_mask_dict(single_mask)
+                    if mask_dict is not None:
+                        normalized_masks.append(mask_dict)
+
+    if not normalized_masks:
         return []
 
-    # Filter by minimum area
-    filtered_masks = [m for m in masks if m["area"] >= min_area]
+    filtered_masks = [m for m in normalized_masks if int(m["area"]) >= min_area]
 
     if not filtered_masks:
         return []
 
-    # Sort by area (largest first) and limit to max_masks
     sorted_masks = sorted(filtered_masks, key=lambda x: x["area"], reverse=True)
     return sorted_masks[:max_masks]
 
 
+def _to_numpy_mask_array(mask_like: Any) -> np.ndarray | None:
+    if mask_like is None:
+        return None
+    if isinstance(mask_like, np.ndarray):
+        return mask_like
+
+    try:
+        import torch
+
+        if isinstance(mask_like, torch.Tensor):
+            return mask_like.detach().cpu().numpy()
+    except ImportError:
+        pass
+
+    return None
+
+
+def _to_mask_dict(mask_like: Any) -> dict[str, Any] | None:
+    if isinstance(mask_like, dict):
+        if "area" in mask_like and "bbox" in mask_like and "segment" in mask_like:
+            return mask_like
+
+        segment = mask_like.get("segment")
+        if segment is None and "mask" in mask_like:
+            segment = mask_like["mask"]
+        if segment is None:
+            return None
+
+        mask_array = _to_numpy_mask_array(segment)
+        if mask_array is None:
+            return None
+        return _build_mask_dict(mask_array)
+
+    mask_array = _to_numpy_mask_array(mask_like)
+    if mask_array is None:
+        return None
+    return _build_mask_dict(mask_array)
+
+
+def _build_mask_dict(mask_array: np.ndarray) -> dict[str, Any] | None:
+    if mask_array.ndim != 2:
+        return None
+    segment = mask_array.astype(bool)
+    area = int(segment.sum())
+    if area <= 0:
+        return None
+
+    ys, xs = np.where(segment)
+    min_y, max_y = int(ys.min()), int(ys.max())
+    min_x, max_x = int(xs.min()), int(xs.max())
+    bbox = [min_x, min_y, max_x - min_x + 1, max_y - min_y + 1]
+
+    return {
+        "segment": segment,
+        "area": area,
+        "bbox": bbox,
+        "predicted_iou": None,
+        "stability_score": None,
+    }
+
+
 def extract_masked_region(
     image: Image.Image,
     mask: np.ndarray,

mini-nav/utils/model.py

@@ -1,6 +1,6 @@
 """Model loading utilities for DINO, SAM2 and HashCompressor."""
 
-from compressors import HashCompressor
+from typing import TYPE_CHECKING
 
 
 import torch
@@ -8,6 +8,9 @@ from transformers import AutoImageProcessor, AutoModel, pipeline, MaskGeneration
 
 from .common import get_device
 
+if TYPE_CHECKING:
+    from compressors.hash_compressor import HashCompressor
+
 
 def load_sam_model(
     model_name: str = "facebook/sam2.1-hiera-large",
@@ -44,7 +47,7 @@ def load_hash_compressor(
     input_dim: int = 1024,
     hash_bits: int = 512,
     compressor_path: str | None = None,
-) -> HashCompressor:
+) -> "HashCompressor":
     from compressors.hash_compressor import HashCompressor
 
     device = get_device()

notebooks/verification.py

@@ -9,11 +9,11 @@
 import marimo
 
 __generated_with = "0.21.1"
-app = marimo.App()
+app = marimo.App(app_title="Pipeline Verification")
 
 
 @app.cell
-def _():
+def import_packages():
     import habitat_sim
     import numpy as np
     import polars as pl
@@ -28,13 +28,16 @@ def _():
     return (
         HashPipeline,
         Image,
+        ObjectNode,
         RoomNode,
         SimpleSceneGraph,
+        extract_masked_region,
         habitat_sim,
         maps,
         np,
         pl,
         plt,
+        segment_image,
     )
 
 
@@ -172,11 +175,21 @@ def _(HashPipeline, hash_bits, sam_max_masks, sam_min_area):
         sam_max_masks=sam_max_masks,
         hash_bits=hash_bits,
     )
-    return
+    return (hash_pipeline,)
 
 
 @app.cell
-def _(Image, SimpleSceneGraph, all_room_views, np, room_nodes):
+def _(
+    Image,
+    ObjectNode,
+    SimpleSceneGraph,
+    all_room_views,
+    extract_masked_region,
+    hash_pipeline,
+    np,
+    room_nodes,
+    segment_image,
+):
     scene_graph = SimpleSceneGraph(
         rooms={_room.room_id: _room for _room in room_nodes}, objects={}
     )
@@ -188,6 +201,39 @@ def _(Image, SimpleSceneGraph, all_room_views, np, room_nodes):
             _rgb3 = _rgb[..., :3] if _rgb.shape[-1] > 3 else _rgb
             _image = Image.fromarray(_rgb3.astype(np.uint8))
 
+            _masks = segment_image(
+                hash_pipeline.mask_generator,
+                _image,
+                min_area=hash_pipeline.sam_min_mask_area,
+                max_masks=hash_pipeline.sam_max_masks,
+                points_per_batch=hash_pipeline.sam_points_per_batch,
+            )
+            total_masks += len(_masks)
+
+            for _mask in _masks:
+                _masked_image = extract_masked_region(_image, _mask["segment"])
+                _bits = hash_pipeline(_masked_image)
+
+                _bbox = _mask["bbox"]
+                _obj_center = np.array(
+                    [_bbox[0] + _bbox[2] / 2, _bbox[1] + _bbox[3] / 2, 0.0],
+                    dtype=np.float32,
+                )
+
+                _obj_id = f"obj_{_obj_index:04d}"
+                _obj_index += 1
+                _bits_np = _bits.squeeze().detach().cpu().numpy()
+
+                _obj_node = ObjectNode(
+                    obj_id=_obj_id,
+                    room_id=_room_id,
+                    position=_obj_center,
+                    visual_hash=_bits_np,
+                    semantic_hash=_bits_np,
+                    hit_count=1,
+                    last_seen_frame=0,
+                )
+                scene_graph.objects[_obj_id] = _obj_node
 
     print(f"Total objects created: {len(scene_graph.objects)}")
     print(f"Total processed masks: {total_masks}")