Mini-Nav/mini-nav/compressors/pipeline.py

"""SAM + DINO + Hash compression pipeline."""

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

import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image


def create_pipeline_from_config(config) -> "HashPipeline":
    """Create HashPipeline from a config object.

    Args:
        config: Configuration object with model settings

    Returns:
        Initialized 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,
    )


class HashPipeline(nn.Module):
    """Pipeline: SAM segmentation + DINO features + Hash compression.

    Pipeline flow:
        PIL Image -> SAM (largest object mask) -> DINO (features) -> Hash (binary codes)

    Usage:
        # Initialize with config
        pipeline = HashPipeline(
            dino_model="facebook/dinov2-large",
            hash_bits=512,
        )

        # Process image
        image = Image.open("path/to/image.jpg")
        hash_bits = pipeline(image)  # [1, 512] binary bits
    """

    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,
    ):
        """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
        """
        super().__init__()

        # Auto detect 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

        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 = get_dino_dim(dino_model)

        # Initialize HashCompressor
        self.hash_compressor = load_hash_compressor(
            input_dim=self.dino_dim,
            hash_bits=hash_bits,
            compressor_path=compressor_path,
        )

    @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.

        Args:
            image: Input PIL Image

        Returns:
            Binary hash codes [1, hash_bits] as int32
        """
        return self._encode_image(image, apply_sam=True)

    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.

        Alias for forward().

        Args:
            image: Input PIL Image

        Returns:
            Binary hash codes [1, hash_bits] as int32
        """
        return self.forward(image)

    def extract_features(self, image: Image.Image) -> torch.Tensor:
        """Extract DINO features from an image.

        Args:
            image: Input PIL Image

        Returns:
            DINO features [1, dino_dim], normalized
        """
        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)
            features = outputs.last_hidden_state.mean(dim=1)  # [1, dim]
            features = F.normalize(features, dim=-1)

        return features