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

mini-nav/compressors/pipeline.py

@@ -1,8 +1,13 @@
-"""Hash compression pipeline with DINO feature extraction.
+"""SAM + DINO + Hash compression pipeline."""
 
-This pipeline extracts features using DINOv2 and compresses them
-to binary hash codes using HashCompressor.
-"""
+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
 
@@ -10,7 +15,6 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from PIL import Image
-from transformers import AutoImageProcessor, AutoModel
 
 
 def create_pipeline_from_config(config) -> "HashPipeline":
@@ -24,17 +28,20 @@ def create_pipeline_from_config(config) -> "HashPipeline":
     """
     return HashPipeline(
         dino_model=config.model.dino_model,
+        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,
-        device=config.model.device if config.model.device != "auto" else None,
     )
 
 
 class HashPipeline(nn.Module):
-    """Pipeline: DINO features + Hash compression.
+    """Pipeline: SAM segmentation + DINO features + Hash compression.
 
     Pipeline flow:
-        PIL Image -> DINO (features) -> Hash (binary codes)
+        PIL Image -> SAM (largest object mask) -> DINO (features) -> Hash (binary codes)
 
     Usage:
         # Initialize with config
@@ -51,14 +58,22 @@ class HashPipeline(nn.Module):
     def __init__(
         self,
         dino_model: str = "facebook/dinov2-large",
+        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,
-        device: 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
@@ -66,53 +81,66 @@ class HashPipeline(nn.Module):
         super().__init__()
 
         # Auto detect device
-        if device is None:
-            device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.device = torch.device(device)
+        self.device = get_device()
 
         self.dino_model = dino_model
+        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
 
-        # Initialize DINO processor and model
-        self.processor = AutoImageProcessor.from_pretrained(dino_model)
-        self.dino = AutoModel.from_pretrained(dino_model).to(self.device)
-        self.dino.eval()
+        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
-        self.dino_dim = 1024 if "large" in dino_model else 768
+        self.dino_dim = get_dino_dim(dino_model)
 
         # Initialize HashCompressor
-        self.hash_compressor = nn.Module()  # Placeholder, will be replaced
-        self._init_hash_compressor(hash_bits, compressor_path)
-
-    def _init_hash_compressor(
-        self, hash_bits: int, compressor_path: Optional[str] = None
-    ):
-        """Initialize the hash compressor module.
-
-        This is called during __init__ but we need to replace it properly.
-        """
-        # Import here to avoid circular imports
-        from .hash_compressor import HashCompressor
-
-        compressor = HashCompressor(input_dim=self.dino_dim, hash_bits=hash_bits).to(
-            self.device
+        self.hash_compressor = load_hash_compressor(
+            input_dim=self.dino_dim,
+            hash_bits=hash_bits,
+            compressor_path=compressor_path,
         )
 
-        # Load pretrained compressor if provided
-        if compressor_path is not None:
-            compressor.load_state_dict(
-                torch.load(compressor_path, map_location=self.device)
-            )
-            print(f"[OK] Loaded HashCompressor from {compressor_path}")
-
-        # Replace the placeholder
-        self.hash_compressor = compressor
-
     @property
     def hash_bits(self):
         """Return the number of hash bits."""
         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
+
+        masks = segment_image(
+            self.mask_generator,
+            image,
+            min_area=self.sam_min_mask_area,
+            max_masks=self.sam_max_masks,
+        )
+
+        if not masks:
+            return image
+
+        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)
+
+        with torch.no_grad():
+            outputs = self.dino(**inputs)
+            tokens = outputs.last_hidden_state
+
+        _, _, bits = self.hash_compressor(tokens)
+        return bits
+
     def forward(self, image: Image.Image) -> torch.Tensor:
         """Process a single image through the pipeline.
 
@@ -122,17 +150,11 @@ class HashPipeline(nn.Module):
         Returns:
             Binary hash codes [1, hash_bits] as int32
         """
-        # Extract DINO features
-        inputs = self.processor(image, return_tensors="pt").to(self.device)
+        return self._encode_image(image, apply_sam=True)
 
-        with torch.no_grad():
-            outputs = self.dino(**inputs)
-            tokens = outputs.last_hidden_state  # [1, N, dim]
-
-        # Compress to hash codes
-        _, _, bits = self.hash_compressor(tokens)
-
-        return bits
+    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.
@@ -156,7 +178,8 @@ class HashPipeline(nn.Module):
         Returns:
             DINO features [1, dino_dim], normalized
         """
-        inputs = self.processor(image, return_tensors="pt").to(self.device)
+        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)
@@ -164,7 +187,3 @@ class HashPipeline(nn.Module):
             features = F.normalize(features, dim=-1)
 
         return features
-
-
-# Backward compatibility alias
-SAMHashPipeline = HashPipeline