"""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)