refactor(compressors): switch SAM from automatic mask generation to bbox-prompted segmentation

- Replace SAM2AutomaticMaskGenerator pipeline with Sam2Processor+Sam2Model
- Freeze SAM model parameters at load time, removing torch.no_grad() at call sites
- Rewrite proposal/core.py to use bbox prompts instead of automatic point sampling
- Add bboxes parameter to all HashPipeline public methods (forward, forward_dataset, extract_features, extract_features_dataset)
- Extract mask filtering logic (_filter_masks) from proposal into pipeline
- Rename object_score/ to filter/
- Add load_owlv2_model to model_loader
- Rename notebooks/test.py to habitat_sim_setup.py

mini-nav/compressors/model_loader.py

@@ -1,11 +1,17 @@
-"""Model loading utilities for DINO, SAM2 and HashCompressor."""
+"""Model loading utilities for DINO, SAM2, OWLv2 and HashCompressor."""
 
 from typing import TYPE_CHECKING, Any
 
-
 import torch
-from transformers import AutoImageProcessor, AutoModel, pipeline, MaskGenerationPipeline
-
+from transformers import (
+    AutoImageProcessor,
+    AutoModel,
+    Owlv2ForObjectDetection,
+    Owlv2Model,
+    Owlv2Processor,
+    Sam2Model,
+    Sam2Processor,
+)
 from utils import get_device
 
 if TYPE_CHECKING:
@@ -14,14 +20,17 @@ if TYPE_CHECKING:
 
 def load_sam_model(
     model_name: str = "facebook/sam2.1-hiera-large",
-) -> MaskGenerationPipeline:
+) -> tuple[Sam2Processor, Sam2Model]:
+    """Load SAM2 processor and model with frozen parameters."""
     device = get_device()
 
-    return pipeline(
-        task="mask-generation",
-        model=model_name,
-        device=device,
-    )
+    processor = Sam2Processor.from_pretrained(model_name)
+    model = Sam2Model.from_pretrained(model_name).to(device)
+    model.eval()
+    for param in model.parameters():
+        param.requires_grad = False
+
+    return processor, model
 
 
 def load_dino_model(
@@ -36,6 +45,18 @@ def load_dino_model(
     return processor, dino
 
 
+def load_owlv2_model(
+    model_name: str = "google/owlv2-base-patch16-ensemble",
+) -> tuple[Owlv2Processor, Owlv2Model]:
+    device = get_device()
+
+    processor = Owlv2Processor.from_pretrained(model_name)
+    model = Owlv2ForObjectDetection.from_pretrained(model_name).to(device)
+    model.eval()
+
+    return processor, model
+
+
 def get_dino_dim(model_name: str) -> int:
     if "large" in model_name.lower():
         return 1024

mini-nav/compressors/pipeline.py

@@ -6,20 +6,20 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from PIL import Image
-
-from .object_score import select_best_mask
-from .proposal import (
-    extract_masked_region,
-    generate_proposals,
-    generate_proposals_batch,
-)
 from utils import get_device
+
+from .filter import select_best_mask
 from .model_loader import (
     get_dino_dim,
     load_dino_model,
     load_hash_compressor,
     load_sam_model,
 )
+from .proposal import (
+    extract_masked_region,
+    generate_proposals,
+    generate_proposals_batch,
+)
 
 
 def create_pipeline_from_config(config) -> "HashPipeline":
@@ -36,7 +36,6 @@ def create_pipeline_from_config(config) -> "HashPipeline":
         sam_model=config.model.sam_model,
         sam_min_mask_area=config.model.sam_min_mask_area,
         sam_max_masks=config.model.sam_max_masks,
-        sam_points_per_batch=config.model.sam_points_per_batch,
         hash_bits=config.model.compression_dim,
         compressor_path=config.model.compressor_path,
     )
@@ -60,7 +59,6 @@ class HashPipeline(nn.Module):
         sam_model: str = "facebook/sam2.1-hiera-large",
         sam_min_mask_area: int = 100,
         sam_max_masks: int = 10,
-        sam_points_per_batch: int = 64,
         hash_bits: int = 512,
         compressor_path: Optional[str] = None,
     ):
@@ -69,15 +67,13 @@ class HashPipeline(nn.Module):
         # Device for model placement.
         self.device = get_device()
 
-        # SAM2 settings.
-        self.sam_model_name = sam_model
+        # SAM2 filter settings.
         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)
+        self.sam_processor, self.sam_model = load_sam_model(model_name=sam_model)
+        self.dino_processor, self.dino = load_dino_model(model_name=dino_model)
 
         # DINO feature dimension based on model size.
         self.dino_dim = get_dino_dim(dino_model)
@@ -94,25 +90,27 @@ class HashPipeline(nn.Module):
         """Number of bits in the hash code."""
         return self.hash_compressor.hash_bits
 
-    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.
+    def _segment_with_sam(
+        self, image: Image.Image, bboxes: list[list[float]]
+    ) -> Image.Image:
+        """Segment image with SAM and extract the best object mask.
 
         Args:
             image: Input PIL Image.
+            bboxes: Bounding boxes from object detector as [[x1,y1,x2,y2], ...].
 
         Returns:
-            Masked image containing only the largest object, or original if no masks.
+            Masked image containing only the best object, or original if no masks.
         """
         masks = generate_proposals(
-            self.mask_generator,
+            self.sam_model,
+            self.sam_processor,
             image,
-            min_area=self.sam_min_mask_area,
-            max_masks=self.sam_max_masks,
-            points_per_batch=self.sam_points_per_batch,
+            bboxes,
         )
 
+        masks = _filter_masks(masks, self.sam_min_mask_area, self.sam_max_masks)
+
         if not masks:
             return image
 
@@ -124,15 +122,19 @@ class HashPipeline(nn.Module):
     def _segment_with_sam_dataset(
         self,
         images: Sequence[Image.Image],
+        bboxes_per_image: list[list[list[float]]],
     ) -> list[Image.Image]:
         image_list = list(images)
         masks_dataset = generate_proposals_batch(
-            self.mask_generator,
+            self.sam_model,
+            self.sam_processor,
             image_list,
-            min_area=self.sam_min_mask_area,
-            max_masks=self.sam_max_masks,
-            points_per_batch=self.sam_points_per_batch,
+            bboxes_per_image,
         )
+        masks_dataset = [
+            _filter_masks(masks, self.sam_min_mask_area, self.sam_max_masks)
+            for masks in masks_dataset
+        ]
         selected_images: list[Image.Image] = []
         for image, masks in zip(image_list, masks_dataset):
             if not masks:
@@ -156,14 +158,14 @@ class HashPipeline(nn.Module):
         Returns:
             Last hidden state tokens of shape [1, N, dim].
         """
-        inputs = self.processor(image, return_tensors="pt").to(self.device)
+        inputs = self.dino_processor(image, return_tensors="pt").to(self.device)
 
         with torch.no_grad():
             outputs = self.dino(**inputs)
             return outputs.last_hidden_state
 
     def _dino_forward_batch(self, images: Sequence[Image.Image]) -> torch.Tensor:
-        inputs = self.processor(images=list(images), return_tensors="pt").to(
+        inputs = self.dino_processor(images=list(images), return_tensors="pt").to(
             self.device
         )
 
@@ -171,16 +173,17 @@ class HashPipeline(nn.Module):
             outputs = self.dino(**inputs)
             return outputs.last_hidden_state
 
-    def forward(self, image: Image.Image) -> torch.Tensor:
+    def forward(self, image: Image.Image, bboxes: list[list[float]]) -> torch.Tensor:
         """Process a single image through the full pipeline.
 
         Args:
             image: Input PIL Image.
+            bboxes: Bounding boxes from object detector as [[x1,y1,x2,y2], ...].
 
         Returns:
             Binary hash codes of shape [1, hash_bits] as int32.
         """
-        image = self._segment_with_sam(image)
+        image = self._segment_with_sam(image, bboxes)
         tokens = self._dino_forward(image)
         _, _, bits = self.hash_compressor(tokens)
         return bits
@@ -188,6 +191,7 @@ class HashPipeline(nn.Module):
     def forward_dataset(
         self,
         images: Sequence[Image.Image],
+        bboxes_per_image: list[list[list[float]]],
         batch_size: int = 32,
         apply_sam: bool = True,
     ) -> torch.Tensor:
@@ -202,7 +206,9 @@ class HashPipeline(nn.Module):
             )
 
         if apply_sam:
-            processed_images = self._segment_with_sam_dataset(image_list)
+            processed_images = self._segment_with_sam_dataset(
+                image_list, bboxes_per_image
+            )
         else:
             processed_images = image_list
 
@@ -215,16 +221,19 @@ class HashPipeline(nn.Module):
 
         return torch.cat(batch_bits, dim=0)
 
-    def extract_features(self, image: Image.Image) -> torch.Tensor:
+    def extract_features(
+        self, image: Image.Image, bboxes: list[list[float]]
+    ) -> torch.Tensor:
         """Extract normalized DINO features from an image.
 
         Args:
             image: Input PIL Image.
+            bboxes: Bounding boxes from object detector as [[x1,y1,x2,y2], ...].
 
         Returns:
             Normalized DINO features of shape [1, dino_dim].
         """
-        image = self._segment_with_sam(image)
+        image = self._segment_with_sam(image, bboxes)
         tokens = self._dino_forward(image)
         features = tokens.mean(dim=1)
         return F.normalize(features, dim=-1)
@@ -232,6 +241,7 @@ class HashPipeline(nn.Module):
     def extract_features_dataset(
         self,
         images: Sequence[Image.Image],
+        bboxes_per_image: list[list[list[float]]],
         batch_size: int = 32,
         apply_sam: bool = True,
     ) -> torch.Tensor:
@@ -246,7 +256,9 @@ class HashPipeline(nn.Module):
             )
 
         if apply_sam:
-            processed_images = self._segment_with_sam_dataset(image_list)
+            processed_images = self._segment_with_sam_dataset(
+                image_list, bboxes_per_image
+            )
         else:
             processed_images = image_list
 
@@ -258,3 +270,16 @@ class HashPipeline(nn.Module):
             all_features.append(F.normalize(features, dim=-1))
 
         return torch.cat(all_features, dim=0)
+
+
+def _filter_masks(
+    masks: list[dict],
+    min_area: int,
+    max_masks: int,
+) -> list[dict]:
+    """Filter masks by area and keep top-N largest."""
+    filtered = [m for m in masks if int(m["area"]) >= min_area]
+    if not filtered:
+        return []
+    sorted_masks = sorted(filtered, key=lambda m: m["area"], reverse=True)
+    return sorted_masks[:max_masks]

mini-nav/compressors/proposal/core.py

@@ -1,27 +1,26 @@
-"""SAM mask proposal generation."""
+"""SAM mask proposal generation via bounding box prompts."""
 
 from typing import Any, Sequence
-import torch
 
 import numpy as np
 from PIL import Image
+from transformers import Sam2Model, Sam2Processor
+from utils import get_device
 
 
 def generate_proposals(
-    mask_generator: Any,
+    model: Sam2Model,
+    processor: Sam2Processor,
     image: Image.Image,
-    min_area: int = 32 * 32,
-    max_masks: int = 5,
-    points_per_batch: int = 64,
+    bboxes: list[list[float]],
 ) -> list[dict[str, Any]]:
-    """Segment image using SAM to extract object masks.
+    """Segment regions in image using SAM2 with bounding box prompts.
 
     Args:
-        mask_generator: SAM2 mask generator.
+        model: Sam2Model instance.
+        processor: Sam2Processor instance.
         image: PIL Image to segment.
-        min_area: Minimum mask area threshold in pixels.
-        max_masks: Maximum number of masks to return.
-        points_per_batch: Number of prompt points to process in each batch.
+        bboxes: Bounding boxes as [[x1, y1, x2, y2], ...].
 
     Returns:
         List of mask dictionaries with keys:
@@ -31,28 +30,40 @@ def generate_proposals(
             - predicted_iou: Model's confidence in the mask
             - stability_score: Stability score for the mask
     """
+    if not bboxes:
+        return []
+
+    device = get_device()
     image_rgb = image.convert("RGB")
-    raw_output = mask_generator(image_rgb, points_per_batch=points_per_batch)
-    return _normalize_and_filter_masks(
-        raw_output, min_area=min_area, max_masks=max_masks
-    )
+    input_boxes = [bboxes]
+    inputs = processor(
+        images=image_rgb,
+        input_boxes=input_boxes,
+        return_tensors="pt",
+    ).to(device)
+
+    outputs = model(**inputs, multimask_output=False)
+    masks = processor.post_process_masks(
+        outputs.pred_masks.cpu(),
+        inputs["original_sizes"],
+    )[0]
+
+    return _masks_to_proposals(masks)
 
 
 def generate_proposals_batch(
-    mask_generator: Any,
+    model: Sam2Model,
+    processor: Sam2Processor,
     images: Sequence[Image.Image],
-    min_area: int = 32 * 32,
-    max_masks: int = 5,
-    points_per_batch: int = 64,
+    bboxes_per_image: list[list[list[float]]],
 ) -> list[list[dict[str, Any]]]:
-    """Segment a batch of images using SAM.
+    """Segment a batch of images using SAM2 with bounding box prompts.
 
     Args:
-        mask_generator: SAM2 mask generator.
+        model: Sam2Model instance.
+        processor: Sam2Processor instance.
         images: Sequence of PIL Images to segment.
-        min_area: Minimum mask area threshold in pixels.
-        max_masks: Maximum number of masks to return per image.
-        points_per_batch: Number of prompt points to process in each batch.
+        bboxes_per_image: Bounding boxes per image, outer list matches images length.
 
     Returns:
         List of lists of mask dictionaries, one inner list per image.
@@ -61,90 +72,48 @@ def generate_proposals_batch(
     if not image_list:
         return []
 
-    image_rgb_list = [image.convert("RGB") for image in image_list]
-    raw_batch_output = mask_generator(
-        image_rgb_list,
-        points_per_batch=points_per_batch,
-    )
-    batch_items = _split_batch_output(raw_batch_output, expected_size=len(image_list))
-    if batch_items is not None:
-        return [
-            _normalize_and_filter_masks(
-                batch_item,
-                min_area=min_area,
-                max_masks=max_masks,
-            )
-            for batch_item in batch_items
-        ]
+    device = get_device()
+    image_rgb_list = [img.convert("RGB") for img in image_list]
 
-    return [
-        _normalize_and_filter_masks(
-            mask_generator(image_rgb, points_per_batch=points_per_batch),
-            min_area=min_area,
-            max_masks=max_masks,
-        )
-        for image_rgb in image_rgb_list
-    ]
+    inputs = processor(
+        images=image_rgb_list,
+        input_boxes=bboxes_per_image,
+        return_tensors="pt",
+    ).to(device)
 
-
-def _split_batch_output(raw_output: Any, expected_size: int) -> list[Any] | None:
-    """Attempt to split raw batch output into per-image results."""
-    if isinstance(raw_output, list):
-        if len(raw_output) == expected_size:
-            return raw_output
-        return None
-
-    if isinstance(raw_output, dict):
-        raw_masks = raw_output.get("masks", raw_output)
-        if isinstance(raw_masks, list) and len(raw_masks) == expected_size:
-            return raw_masks
-
-    return None
-
-
-def _normalize_and_filter_masks(
-    raw_output: Any,
-    min_area: int,
-    max_masks: int,
-) -> list[dict[str, Any]]:
-    """Normalize raw SAM output into mask dicts and filter by area/count."""
-    raw_masks = (
-        raw_output.get("masks", raw_output)
-        if isinstance(raw_output, dict)
-        else raw_output
+    outputs = model(**inputs, multimask_output=False)
+    all_masks = processor.post_process_masks(
+        outputs.pred_masks.cpu(),
+        inputs["original_sizes"],
     )
 
-    normalized_masks: list[dict[str, Any]] = []
-    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:
+    return [_masks_to_proposals(image_masks) for image_masks in all_masks]
+
+
+def _masks_to_proposals(masks: Any) -> list[dict[str, Any]]:
+    """Convert model output masks to list of mask dicts."""
+    mask_array = _to_numpy_mask_array(masks)
+    if mask_array is None:
+        return []
+
+    # Ensure 3D: [num_masks, H, W]
     if mask_array.ndim == 2:
         mask_array = np.expand_dims(mask_array, axis=0)
-            if mask_array.ndim == 3:
+
+    if mask_array.ndim != 3:
+        return []
+
+    # Remove batch dim if present: [1, num_masks, H, W] → [num_masks, H, W]
+    if mask_array.ndim == 3 and mask_array.shape[0] == 1:
+        mask_array = mask_array[0]
+
+    proposals: list[dict[str, Any]] = []
     for single_mask in mask_array:
-                    mask_dict = _to_mask_dict(single_mask)
+        mask_dict = _build_mask_dict(single_mask)
         if mask_dict is not None:
-                        normalized_masks.append(mask_dict)
+            proposals.append(mask_dict)
 
-    if not normalized_masks:
-        return []
-
-    filtered_masks = [
-        mask for mask in normalized_masks if int(mask["area"]) >= min_area
-    ]
-    if not filtered_masks:
-        return []
-
-    sorted_masks = sorted(filtered_masks, key=lambda mask: mask["area"], reverse=True)
-    return sorted_masks[:max_masks]
+    return proposals
 
 
 def _to_numpy_mask_array(mask_like: Any) -> np.ndarray | None:
@@ -154,35 +123,14 @@ def _to_numpy_mask_array(mask_like: Any) -> np.ndarray | None:
     if isinstance(mask_like, np.ndarray):
         return mask_like
 
+    import torch
+
     if isinstance(mask_like, torch.Tensor):
         return mask_like.detach().cpu().numpy()
 
     return None
 
 
-def _to_mask_dict(mask_like: Any) -> dict[str, Any] | None:
-    """Convert a single mask-like object to a standardized mask dict."""
-    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:
     """Build a mask dictionary from a 2D boolean numpy array."""
     if mask_array.ndim != 2: