refactor(compressors): reorganize SAM utilities into proposal module

mini-nav/compressors/model_loader.py

@@ -1,12 +1,12 @@
 """Model loading utilities for DINO, SAM2 and HashCompressor."""
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 
 import torch
 from transformers import AutoImageProcessor, AutoModel, pipeline, MaskGenerationPipeline
 
-from .common import get_device
+from utils import get_device
 
 if TYPE_CHECKING:
     from compressors.hash_compressor import HashCompressor
@@ -26,7 +26,7 @@ def load_sam_model(
 
 def load_dino_model(
     model_name: str = "facebook/dinov2-large",
-) -> tuple[AutoImageProcessor, AutoModel]:
+) -> tuple[Any, Any]:
     device = get_device()
 
     processor = AutoImageProcessor.from_pretrained(model_name)

mini-nav/compressors/pipeline.py

@@ -8,9 +8,13 @@ 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 utils.image import extract_masked_region, segment_image, segment_image_dataset
+from .model_loader import (
-from utils.model import (
     get_dino_dim,
     load_dino_model,
     load_hash_compressor,
@@ -101,7 +105,7 @@ class HashPipeline(nn.Module):
         Returns:
             Masked image containing only the largest object, or original if no masks.
         """
-        masks = segment_image(
+        masks = generate_proposals(
             self.mask_generator,
             image,
             min_area=self.sam_min_mask_area,
@@ -122,7 +126,7 @@ class HashPipeline(nn.Module):
         images: Sequence[Image.Image],
     ) -> list[Image.Image]:
         image_list = list(images)
-        masks_dataset = segment_image_dataset(
+        masks_dataset = generate_proposals_batch(
             self.mask_generator,
             image_list,
             min_area=self.sam_min_mask_area,

mini-nav/compressors/proposal/__init__.py

@@ -0,0 +1,10 @@
+"""Proposal module — SAM mask generation and extraction."""
+
+from .core import generate_proposals, generate_proposals_batch
+from .utils import extract_masked_region
+
+__all__ = [
+    "generate_proposals",
+    "generate_proposals_batch",
+    "extract_masked_region",
+]

mini-nav/compressors/proposal/core.py

@@ -1,10 +1,13 @@
+"""SAM mask proposal generation."""
+
 from typing import Any, Sequence
+import torch
 
 import numpy as np
 from PIL import Image
 
 
-def segment_image(
+def generate_proposals(
     mask_generator: Any,
     image: Image.Image,
     min_area: int = 32 * 32,
@@ -19,6 +22,7 @@ def segment_image(
         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.
+
     Returns:
         List of mask dictionaries with keys:
             - segment: Binary mask (numpy array)
@@ -34,26 +38,35 @@ def segment_image(
     )
 
 
-def segment_image_dataset(
+def generate_proposals_batch(
     mask_generator: Any,
     images: Sequence[Image.Image],
     min_area: int = 32 * 32,
     max_masks: int = 5,
     points_per_batch: int = 64,
 ) -> list[list[dict[str, Any]]]:
+    """Segment a batch of images using SAM.
+
+    Args:
+        mask_generator: SAM2 mask generator.
+        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.
+
+    Returns:
+        List of lists of mask dictionaries, one inner list per image.
+    """
     image_list = list(images)
     if not image_list:
         return []
 
     image_rgb_list = [image.convert("RGB") for image in image_list]
-    try:
     raw_batch_output = mask_generator(
         image_rgb_list,
         points_per_batch=points_per_batch,
     )
-        batch_items = _split_batch_output(
+    batch_items = _split_batch_output(raw_batch_output, expected_size=len(image_list))
-            raw_batch_output, expected_size=len(image_list)
-        )
     if batch_items is not None:
         return [
             _normalize_and_filter_masks(
@@ -63,8 +76,6 @@ def segment_image_dataset(
             )
             for batch_item in batch_items
         ]
-    except TypeError:
-        pass
 
     return [
         _normalize_and_filter_masks(
@@ -77,6 +88,7 @@ def segment_image_dataset(
 
 
 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
@@ -95,6 +107,7 @@ def _normalize_and_filter_masks(
     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)
@@ -135,23 +148,20 @@ def _normalize_and_filter_masks(
 
 
 def _to_numpy_mask_array(mask_like: Any) -> np.ndarray | None:
+    """Convert mask-like object to numpy array."""
     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:
+    """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
@@ -174,6 +184,7 @@ def _to_mask_dict(mask_like: Any) -> dict[str, Any] | None:
 
 
 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:
         return None
     segment = mask_array.astype(bool)
@@ -193,24 +204,3 @@ def _build_mask_dict(mask_array: np.ndarray) -> dict[str, Any] | None:
         "predicted_iou": None,
         "stability_score": None,
     }
-
-
-def extract_masked_region(
-    image: Image.Image,
-    mask: np.ndarray,
-) -> Image.Image:
-    """Extract masked region from image.
-
-    Args:
-        image: Original PIL Image.
-        mask: Binary mask as numpy array (True = keep).
-
-    Returns:
-        PIL Image with only the masked region visible.
-    """
-    image_np = np.array(image.convert("RGB"))
-
-    # Apply mask
-    masked_np = image_np * mask[:, :, np.newaxis]
-
-    return Image.fromarray(masked_np.astype(np.uint8))

mini-nav/compressors/proposal/utils.py

@@ -0,0 +1,25 @@
+"""Mask extraction utilities."""
+
+import numpy as np
+from PIL import Image
+
+
+def extract_masked_region(
+    image: Image.Image,
+    mask: np.ndarray,
+) -> Image.Image:
+    """Extract masked region from image.
+
+    Args:
+        image: Original PIL Image.
+        mask: Binary mask as numpy array (True = keep).
+
+    Returns:
+        PIL Image with only the masked region visible.
+    """
+    image_np = np.array(image.convert("RGB"))
+
+    # Apply mask.
+    masked_np = image_np * mask[:, :, np.newaxis]
+
+    return Image.fromarray(masked_np.astype(np.uint8))

mini-nav/utils/__init__.py

@@ -4,8 +4,6 @@ from .feature_extractor import (
     extract_single_image_feature,
     infer_vector_dim,
 )
-from .image import extract_masked_region, segment_image, segment_image_dataset
-from .model import get_dino_dim, load_dino_model, load_hash_compressor, load_sam_model
 
 __all__ = [
     "get_device",
@@ -13,11 +11,4 @@ __all__ = [
     "infer_vector_dim",
     "extract_single_image_feature",
     "extract_batch_features",
-    "segment_image",
-    "segment_image_dataset",
-    "extract_masked_region",
-    "load_dino_model",
-    "load_sam_model",
-    "get_dino_dim",
-    "load_hash_compressor",
 ]

notebooks/verification.py

@@ -33,7 +33,7 @@ def import_packages():
         create_habitat_simulator,
         render_topdown_scene_map,
     )
-    from utils.image import extract_masked_region, segment_image_dataset
+    from compressors.proposal import extract_masked_region, generate_proposals_batch
 
     return (
         HabitatSimulatorConfig,
@@ -53,7 +53,7 @@ def import_packages():
         pl,
         plt,
         render_topdown_scene_map,
-        segment_image_dataset,
+        generate_proposals_batch,
     )
 
 
@@ -150,7 +150,7 @@ def build_scene_graph_pipeline(
     room_nodes,
     sam_max_masks,
     sam_min_area,
-    segment_image_dataset,
+    generate_proposals_batch,
     sim,
     views_per_room,
 ):
@@ -191,7 +191,7 @@ def build_scene_graph_pipeline(
         rgb3 = rgb[..., :3] if rgb.shape[-1] > 3 else rgb
         room_view_images.append(Image.fromarray(rgb3.astype(np.uint8)))
 
-    masks_dataset = segment_image_dataset(
+    masks_dataset = generate_proposals_batch(
         hash_pipeline.mask_generator,
         room_view_images,
         min_area=hash_pipeline.sam_min_mask_area,