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

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

from typing import Optional

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

from utils import get_device
from utils.image import extract_masked_region, segment_image
from utils.model import (
    get_dino_dim,
    load_dino_model,
    load_hash_compressor,
    load_sam_model,
)


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 for SAM segmentation + DINO features + Hash compression.

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

    Example:
        pipeline = HashPipeline(dino_model="facebook/dinov2-large", hash_bits=512)
        image = Image.open("path/to/image.jpg")
        hash_bits = pipeline(image)  # Returns [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,
    ):
        super().__init__()

        # Device for model placement.
        self.device = get_device()

        # SAM2 settings.
        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

        # Load models.
        self.mask_generator = load_sam_model(model_name=sam_model)
        self.processor, self.dino = load_dino_model(model_name=dino_model)

        # DINO feature dimension based on model size.
        self.dino_dim = get_dino_dim(dino_model)

        # Hash compressor for binarizing DINO features.
        self.hash_compressor = load_hash_compressor(
            input_dim=self.dino_dim,
            hash_bits=hash_bits,
            compressor_path=compressor_path,
        )

    @property
    def hash_bits(self) -> int:
        """Number of bits in the hash code."""
        return self.hash_compressor.hash_bits

    def _segment_with_sam(self, image: Image.Image) -> Image.Image:
        """Segment image with SAM and extract the largest object mask.

        If no valid masks are found, returns the original image.

        Args:
            image: Input PIL Image.

        Returns:
            Masked image containing only the largest object, or original if no masks.
        """
        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 _dino_forward(self, image: Image.Image) -> torch.Tensor:
        """Extract DINO tokens from an image.

        Args:
            image: Input PIL Image.

        Returns:
            Last hidden state tokens of shape [1, N, dim].
        """
        inputs = self.processor(image, return_tensors="pt").to(self.device)

        with torch.no_grad():
            outputs = self.dino(**inputs)
            return outputs.last_hidden_state

    def forward(self, image: Image.Image) -> torch.Tensor:
        """Process a single image through the full pipeline.

        Args:
            image: Input PIL Image.

        Returns:
            Binary hash codes of shape [1, hash_bits] as int32.
        """
        image = self._segment_with_sam(image)
        tokens = self._dino_forward(image)
        _, _, bits = self.hash_compressor(tokens)
        return bits

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

        Args:
            image: Input PIL Image.

        Returns:
            Normalized DINO features of shape [1, dino_dim].
        """
        image = self._segment_with_sam(image)
        tokens = self._dino_forward(image)
        features = tokens.mean(dim=1)
        return F.normalize(features, dim=-1)