refactor(compressors): reorganize SAM utilities into proposal module

mini-nav/compressors/model_loader.py

@@ -0,0 +1,59 @@
+"""Model loading utilities for DINO, SAM2 and HashCompressor."""
+
+from typing import TYPE_CHECKING, Any
+
+
+import torch
+from transformers import AutoImageProcessor, AutoModel, pipeline, MaskGenerationPipeline
+
+from utils import get_device
+
+if TYPE_CHECKING:
+    from compressors.hash_compressor import HashCompressor
+
+
+def load_sam_model(
+    model_name: str = "facebook/sam2.1-hiera-large",
+) -> MaskGenerationPipeline:
+    device = get_device()
+
+    return pipeline(
+        task="mask-generation",
+        model=model_name,
+        device=device,
+    )
+
+
+def load_dino_model(
+    model_name: str = "facebook/dinov2-large",
+) -> tuple[Any, Any]:
+    device = get_device()
+
+    processor = AutoImageProcessor.from_pretrained(model_name)
+    dino = AutoModel.from_pretrained(model_name).to(device)
+    dino.eval()
+
+    return processor, dino
+
+
+def get_dino_dim(model_name: str) -> int:
+    if "large" in model_name.lower():
+        return 1024
+    return 768
+
+
+def load_hash_compressor(
+    input_dim: int = 1024,
+    hash_bits: int = 512,
+    compressor_path: str | None = None,
+) -> "HashCompressor":
+    from compressors.hash_compressor import HashCompressor
+
+    device = get_device()
+    compressor = HashCompressor(input_dim=input_dim, hash_bits=hash_bits).to(device)
+
+    if compressor_path is not None:
+        compressor.load_state_dict(torch.load(compressor_path, map_location=device))
+        print(f"[OK] Loaded HashCompressor from {compressor_path}")
+
+    return compressor

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 utils.model import (
+from .model_loader 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

@@ -0,0 +1,206 @@
+"""SAM mask proposal generation."""
+
+from typing import Any, Sequence
+import torch
+
+import numpy as np
+from PIL import Image
+
+
+def generate_proposals(
+    mask_generator: Any,
+    image: Image.Image,
+    min_area: int = 32 * 32,
+    max_masks: int = 5,
+    points_per_batch: int = 64,
+) -> list[dict[str, Any]]:
+    """Segment image using SAM to extract object masks.
+
+    Args:
+        mask_generator: SAM2 mask generator.
+        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.
+
+    Returns:
+        List of mask dictionaries with keys:
+            - segment: Binary mask (numpy array)
+            - area: Mask area in pixels
+            - bbox: Bounding box [x, y, width, height]
+            - predicted_iou: Model's confidence in the mask
+            - stability_score: Stability score for the mask
+    """
+    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
+    )
+
+
+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]
+    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
+        ]
+
+    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
+    ]
+
+
+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
+    )
+
+    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:
+            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 []
+
+    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]
+
+
+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
+
+    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:
+        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,
+    }

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))