refactor(pipeline): integrate SAM segmentation and modularize model loading

mini-nav/utils/__init__.py

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

mini-nav/utils/common.py

@@ -3,11 +3,10 @@ from pathlib import Path
 
 import torch
 from configs import cfg_manager
-from torch.types import Device
 
 
 @lru_cache(maxsize=1)
-def get_device() -> Device:
+def get_device() -> torch.device:
     config = cfg_manager.get()
     device = config.model.device
     if device == "auto":

mini-nav/utils/image.py

@@ -0,0 +1,68 @@
+from typing import Any
+
+import numpy as np
+from PIL import Image
+
+
+def segment_image(
+    mask_generator: Any,
+    image: Image.Image,
+    min_area: int = 32 * 32,
+    max_masks: int = 5,
+    points_per_batch=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
+    """
+    # Convert PIL Image to numpy array
+    image_np = np.array(image.convert("RGB"))
+
+    # Generate masks
+    masks = mask_generator(image_np, points_per_batch=points_per_batch)["masks"]
+
+    if not masks:
+        return []
+
+    # Filter by minimum area
+    filtered_masks = [m for m in masks if 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 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/model.py

@@ -0,0 +1,57 @@
+"""Model loading utilities for DINO, SAM2 and HashCompressor."""
+
+from compressors import HashCompressor
+
+
+import torch
+from transformers import AutoImageProcessor, AutoModel, pipeline, MaskGenerationPipeline
+
+from .common import get_device
+
+
+def load_sam_model(
+    model_name: str = "facebook/sam2.1-hiera-large",
+) -> MaskGenerationPipeline:
+    device = str(get_device())
+    device_id = 0 if device.startswith("cuda") else -1
+
+    return pipeline(
+        task="mask-generation",
+        model=model_name,
+        device=device_id,
+    )
+
+
+def load_dino_model(
+    model_name: str = "facebook/dinov2-large",
+) -> tuple[AutoImageProcessor, AutoModel]:
+    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/utils/sam.py

@@ -1,100 +0,0 @@
-"""SAM (Segment Anything Model) utilities for object segmentation."""
-
-from pathlib import Path
-from typing import Any
-
-import numpy as np
-import torch
-from PIL import Image
-from sam2.build_sam import build_sam2
-from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator
-
-
-def load_sam_model(
-    model_name: str = "facebook/sam2.1-hiera-large",
-    device: str = "cuda",
-    checkpoint_dir: Path | None = None,
-) -> tuple[Any, Any]:
-    """Load SAM 2.1 model and mask generator.
-
-    Args:
-        model_name: SAM model name (currently supports facebook/sam2.1-hiera-*).
-        device: Device to load model on (cuda or cpu).
-        checkpoint_dir: Optional directory for model checkpoint cache.
-
-    Returns:
-        Tuple of (sam_model, mask_generator).
-    """
-    if device == "cuda" and not torch.cuda.is_available():
-        device = "cpu"
-
-    # Build SAM2 model
-    sam_model = build_sam2(model_name, device=device)
-
-    # Create automatic mask generator
-    mask_generator = SAM2AutomaticMaskGenerator(sam_model)
-
-    return sam_model, mask_generator
-
-
-def segment_image(
-    mask_generator: Any,
-    image: Image.Image,
-    min_area: int = 32 * 32,
-    max_masks: int = 5,
-) -> list[dict[str, Any]]:
-    """Segment image using SAM to extract object masks.
-
-    Args:
-        mask_generator: SAM2AutomaticMaskGenerator instance.
-        image: PIL Image to segment.
-        min_area: Minimum mask area threshold in pixels.
-        max_masks: Maximum number of masks to return.
-
-    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
-    """
-    # Convert PIL Image to numpy array
-    image_np = np.array(image.convert("RGB"))
-
-    # Generate masks
-    masks = mask_generator.generate(image_np)
-
-    if not masks:
-        return []
-
-    # Filter by minimum area
-    filtered_masks = [m for m in masks if 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 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))