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

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

from dataclasses import dataclass, field
from typing import Any, Optional, Sequence

import torch
import torch.nn.functional as F
from PIL import Image
from utils import get_device
from utils.image import crop_image_by_bbox, extract_masked_region

from .filter import (
    MaskScoringConfig,
    compute_mask_features,
    score_mask,
    should_reject_mask,
)
from .model_loader import (
    get_dino_dim,
    load_dino_model,
    load_hash_compressor,
    load_owlv2_model,
    load_sam_model,
)
from .proposal import (
    detect_objects_batch,
    generate_proposals_batch,
)
from .proposal.core import DetectionResult


def create_pipeline_from_config(config) -> "HashPipeline":
    return HashPipeline(
        owlv2_model=getattr(
            config.model, "owlv2_model", "google/owlv2-base-patch16-ensemble"
        ),
        dino_model=config.model.dino_model,
        sam_model=config.model.sam_model,
        hash_bits=config.model.compression_dim,
        compressor_path=config.model.compressor_path,
        mask_scoring_config=getattr(config.model, "mask_scoring_config", None),
        score_threshold=getattr(config.model, "score_threshold", 0.25),
        postprocess_threshold=getattr(config.model, "postprocess_threshold", 0.1),
    )


@dataclass
class FramePacket:
    image: Image.Image
    boxes_xyxy: list[list[float]] = field(default_factory=list)
    scores: list[float] = field(default_factory=list)
    labels: list[str] = field(default_factory=list)
    masks: list[dict[str, Any]] = field(default_factory=list)
    selected_idx: int | None = None
    dropped_indices: list[int] = field(default_factory=list)
    fallback_reason: str | None = None
    filtered_image: Image.Image | None = None
    cropped_image: Image.Image | None = None


@dataclass
class PipelineBatchOutput:
    hash_bits: torch.Tensor
    cropped_images: list[Image.Image]
    debug_meta: list[dict[str, Any]]


class HashPipeline:
    """Pipeline for OWLv2 detection + SAM segmentation + DINO features + Hash compression.

    Pipeline flow:
        Images + Text Labels -> OWLv2 (detections) -> SAM (masks) -> Filter (best mask) ->
        Crop (matched OWLv2 box) -> DINO (features) -> Hash (binary codes)

    Example:
        pipeline = HashPipeline(dino_model="facebook/dinov2-large", hash_bits=512)
        images = [Image.open("path/to/image.jpg")]
        text_labels = ["object"]
        output = pipeline.process_batch(images, text_labels)
        hash_bits = output.hash_bits  # [N, 512]
    """

    def __init__(
        self,
        owlv2_model: str = "google/owlv2-base-patch16-ensemble",
        dino_model: str = "facebook/dinov2-large",
        sam_model: str = "facebook/sam2.1-hiera-large",
        hash_bits: int = 512,
        compressor_path: Optional[str] = None,
        mask_scoring_config: Optional["MaskScoringConfig"] = None,
        score_threshold: float = 0.25,
        postprocess_threshold: float = 0.1,
    ):
        super().__init__()

        self.device = get_device()
        self.owlv2_processor, self.owlv2_model = load_owlv2_model(
            model_name=owlv2_model
        )
        self.score_threshold = score_threshold
        self.postprocess_threshold = postprocess_threshold
        self.mask_scoring_config = mask_scoring_config
        self.sam_processor, self.sam_model = load_sam_model(model_name=sam_model)
        self.dino_processor, self.dino = load_dino_model(model_name=dino_model)
        self.dino_dim = get_dino_dim(dino_model)
        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:
        return self.hash_compressor.hash_bits

    def _detect_batch(
        self,
        images: Sequence[Image.Image],
        text_labels: list[str],
    ) -> list[list[DetectionResult]]:
        image_list = list(images)
        if not image_list:
            return []

        text_labels_per_image = [text_labels] * len(image_list)
        return detect_objects_batch(
            self.owlv2_model,
            self.owlv2_processor,
            image_list,
            text_labels_per_image,
            score_threshold=self.score_threshold,
            postprocess_threshold=self.postprocess_threshold,
        )

    def _segment_batch(
        self,
        images: Sequence[Image.Image],
        bboxes_per_image: list[list[list[float]]],
    ) -> list[list[dict[str, Any]]]:
        image_list = list(images)
        if not image_list:
            return []

        return generate_proposals_batch(
            self.sam_model,
            self.sam_processor,
            image_list,
            bboxes_per_image,
        )

    def _build_frame_packets(self, images: Sequence[Image.Image]) -> list[FramePacket]:
        return [FramePacket(image=image) for image in images]

    def _attach_detections(
        self,
        packets: list[FramePacket],
        text_labels: list[str],
    ) -> list[list[DetectionResult]]:
        if not packets:
            return []

        image_list = [packet.image for packet in packets]
        detections_per_image = self._detect_batch(image_list, text_labels)
        for index, packet in enumerate(packets):
            detections = (
                detections_per_image[index] if index < len(detections_per_image) else []
            )
            packet.boxes_xyxy = [list(det["bbox"]) for det in detections]
            packet.scores = [float(det["score"]) for det in detections]
            packet.labels = [str(det["label"]) for det in detections]
            if not packet.boxes_xyxy:
                packet.fallback_reason = "no_detection"
        return detections_per_image

    def _attach_masks(self, packets: list[FramePacket]) -> None:
        if not packets:
            return

        image_list = [packet.image for packet in packets]
        boxes_per_image = [packet.boxes_xyxy for packet in packets]
        masks_per_image = self._segment_batch(image_list, boxes_per_image)

        for index, packet in enumerate(packets):
            packet.masks = (
                masks_per_image[index] if index < len(masks_per_image) else []
            )
            if (
                packet.boxes_xyxy
                and not packet.masks
                and packet.fallback_reason is None
            ):
                packet.fallback_reason = "no_mask"

    def _select_candidates(self, packets: list[FramePacket]) -> None:
        config = self.mask_scoring_config or MaskScoringConfig()

        for packet in packets:
            if not packet.masks:
                packet.selected_idx = None
                packet.dropped_indices = []
                if packet.fallback_reason is None:
                    packet.fallback_reason = "no_mask"
                continue

            kept: list[tuple[float, int, int]] = []
            dropped: list[int] = []
            for index, mask in enumerate(packet.masks):
                features = compute_mask_features(
                    mask, image_shape=(packet.image.height, packet.image.width)
                )
                if should_reject_mask(features, config):
                    dropped.append(index)
                    continue

                mask_score = score_mask(
                    mask,
                    image_shape=(packet.image.height, packet.image.width),
                    config=config,
                )
                area = int(mask.get("area", 0))
                kept.append((float(mask_score), area, index))

            if kept:
                kept.sort(reverse=True)
                packet.selected_idx = kept[0][2]
                packet.dropped_indices = dropped
                continue

            fallback_index = max(
                range(len(packet.masks)),
                key=lambda idx: int(packet.masks[idx].get("area", 0)),
            )
            packet.selected_idx = fallback_index
            packet.dropped_indices = [
                index for index in range(len(packet.masks)) if index != fallback_index
            ]
            if packet.fallback_reason is None:
                packet.fallback_reason = "all_masks_rejected_fallback_area"

    def _render_filtered_images(self, packets: list[FramePacket]) -> None:
        for packet in packets:
            if packet.selected_idx is None:
                packet.filtered_image = packet.image
                continue

            if packet.selected_idx >= len(packet.masks):
                packet.filtered_image = packet.image
                packet.fallback_reason = (
                    packet.fallback_reason or "selected_index_out_of_mask_range"
                )
                packet.selected_idx = None
                continue

            selected_mask = packet.masks[packet.selected_idx]
            packet.filtered_image = extract_masked_region(
                packet.image, selected_mask["segment"]
            )

    def _render_cropped_images(
        self,
        packets: list[FramePacket],
        detections_per_image: list[list[DetectionResult]],
    ) -> None:
        for index, packet in enumerate(packets):
            base_image = (
                packet.filtered_image if packet.filtered_image else packet.image
            )
            detections = (
                detections_per_image[index] if index < len(detections_per_image) else []
            )

            if packet.selected_idx is None:
                packet.cropped_image = base_image
                if packet.fallback_reason is None:
                    packet.fallback_reason = "no_selected_candidate"
                continue

            if packet.selected_idx >= len(detections):
                packet.cropped_image = base_image
                packet.fallback_reason = (
                    packet.fallback_reason or "selected_index_out_of_detection_range"
                )
                continue

            selected_detection = detections[packet.selected_idx]
            cropped = crop_image_by_bbox(base_image, selected_detection["bbox"])
            packet.cropped_image = cropped
            if cropped.size == base_image.size:
                packet.fallback_reason = (
                    packet.fallback_reason or "invalid_or_full_bbox"
                )

    def _build_debug_meta(
        self,
        packets: list[FramePacket],
        return_debug_details: bool,
    ) -> list[dict[str, Any]]:
        debug_meta: list[dict[str, Any]] = []
        for packet in packets:
            item: dict[str, Any] = {
                "selected_idx": packet.selected_idx,
                "dropped_indices": list(packet.dropped_indices),
                "fallback_reason": packet.fallback_reason,
                "num_boxes": len(packet.boxes_xyxy),
                "num_masks": len(packet.masks),
            }
            if return_debug_details:
                item["boxes_xyxy"] = [list(box) for box in packet.boxes_xyxy]
                item["scores"] = [float(score) for score in packet.scores]
                item["labels"] = [str(label) for label in packet.labels]
                item["masks"] = packet.masks
            debug_meta.append(item)
        return debug_meta

    def extract_dino_batch(self, images: Sequence[Image.Image]) -> torch.Tensor:
        """Extract DINO tokens from a batch of images.

        Args:
            images: Sequence of PIL Images.

        Returns:
            Last hidden state tokens of shape [B, N, dim].
        """
        image_list = list(images)
        if not image_list:
            return torch.empty(
                (0, 1, self.dino_dim), dtype=torch.float32, device=self.device
            )

        inputs = self.dino_processor(images=image_list, return_tensors="pt").to(
            self.device
        )

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

    def compress_batch(self, tokens: torch.Tensor) -> torch.Tensor:
        """Compress DINO tokens to binary hash codes.

        Args:
            tokens: DINO tokens of shape [B, N, dim].

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

    def process_batch(
        self,
        images: Sequence[Image.Image],
        text_labels: list[str],
        batch_size: int = 32,
        return_debug_details: bool = False,
    ) -> PipelineBatchOutput:
        """Run full pipeline and return cropped images + hashes + debug metadata.

        Args:
            images: Sequence of PIL Images.
            text_labels: Text labels for detection (same for all images).
            batch_size: Batch size for DINO feature extraction.
            return_debug_details: Include boxes/scores/labels/masks in debug output.

        Returns:
            PipelineBatchOutput with final cropped images, binary hash bits,
            and per-image debug metadata.
        """
        if batch_size <= 0:
            raise ValueError("batch_size must be greater than 0")

        image_list = list(images)
        if not image_list:
            return PipelineBatchOutput(
                hash_bits=torch.empty(
                    (0, self.hash_bits), dtype=torch.int32, device=self.device
                ),
                cropped_images=[],
                debug_meta=[],
            )

        packets = self._build_frame_packets(image_list)
        detections_per_image = self._attach_detections(packets, text_labels)
        self._attach_masks(packets)
        self._select_candidates(packets)
        self._render_filtered_images(packets)
        self._render_cropped_images(packets, detections_per_image)

        cropped_images = [
            packet.cropped_image if packet.cropped_image is not None else packet.image
            for packet in packets
        ]

        all_bits: list[torch.Tensor] = []
        for index in range(0, len(cropped_images), batch_size):
            sub_batch = cropped_images[index : index + batch_size]
            tokens = self.extract_dino_batch(sub_batch)
            bits = self.compress_batch(tokens)
            all_bits.append(bits)

        hash_bits = torch.cat(all_bits, dim=0)
        debug_meta = self._build_debug_meta(packets, return_debug_details)
        return PipelineBatchOutput(
            hash_bits=hash_bits,
            cropped_images=cropped_images,
            debug_meta=debug_meta,
        )

    def forward_batch(
        self,
        images: Sequence[Image.Image],
        text_labels: list[str],
        batch_size: int = 32,
    ) -> torch.Tensor:
        """Process a batch of images through the full pipeline.

        Args:
            images: Sequence of PIL Images.
            text_labels: Text labels for detection (same for all images).
            batch_size: Batch size for DINO feature extraction.

        Returns:
            Binary hash codes of shape [N, hash_bits] as int32.
        """
        return self.process_batch(
            images=images,
            text_labels=text_labels,
            batch_size=batch_size,
        ).hash_bits

    def extract_features_dataset(
        self,
        images: Sequence[Image.Image],
        text_labels: list[str],
        batch_size: int = 32,
    ) -> torch.Tensor:
        """Extract normalized DINO features from a batch of images.

        Args:
            images: Sequence of PIL Images.
            text_labels: Text labels for detection (same for all images).
            batch_size: Batch size for DINO feature extraction.

        Returns:
            Normalized DINO features of shape [N, dino_dim].
        """
        if batch_size <= 0:
            raise ValueError("batch_size must be greater than 0")

        image_list = list(images)
        if not image_list:
            return torch.empty(
                (0, self.dino_dim), dtype=torch.float32, device=self.device
            )

        processed = self.process_batch(
            images=image_list,
            text_labels=text_labels,
            batch_size=batch_size,
        ).cropped_images

        all_features: list[torch.Tensor] = []
        for index in range(0, len(processed), batch_size):
            sub_batch = processed[index : index + batch_size]
            tokens = self.extract_dino_batch(sub_batch)
            features = tokens.mean(dim=1)
            all_features.append(F.normalize(features, dim=-1))

        return torch.cat(all_features, dim=0)