"""SAM + DINO + Hash compression pipeline.""" from typing import Optional, Sequence 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, segment_image_dataset 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, points_per_batch=self.sam_points_per_batch, ) if not masks: return image return extract_masked_region(image, masks[0]["segment"]) def _segment_with_sam_dataset( self, images: Sequence[Image.Image], ) -> list[Image.Image]: image_list = list(images) masks_dataset = segment_image_dataset( self.mask_generator, image_list, min_area=self.sam_min_mask_area, max_masks=self.sam_max_masks, points_per_batch=self.sam_points_per_batch, ) return [ extract_masked_region(image, masks[0]["segment"]) if masks else image for image, masks in zip(image_list, masks_dataset) ] 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 _dino_forward_batch(self, images: Sequence[Image.Image]) -> torch.Tensor: inputs = self.processor(images=list(images), 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 forward_dataset( self, images: Sequence[Image.Image], batch_size: int = 32, apply_sam: bool = True, ) -> torch.Tensor: if batch_size <= 0: raise ValueError("batch_size must be greater than 0") image_list = list(images) if len(image_list) == 0: return torch.empty( (0, self.hash_bits), dtype=torch.int32, device=self.device ) if apply_sam: processed_images = self._segment_with_sam_dataset(image_list) else: processed_images = image_list batch_bits: list[torch.Tensor] = [] for i in range(0, len(processed_images), batch_size): batch_images = processed_images[i : i + batch_size] tokens = self._dino_forward_batch(batch_images) _, _, bits = self.hash_compressor(tokens) batch_bits.append(bits) return torch.cat(batch_bits, dim=0) 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) def extract_features_dataset( self, images: Sequence[Image.Image], batch_size: int = 32, apply_sam: bool = True, ) -> torch.Tensor: if batch_size <= 0: raise ValueError("batch_size must be greater than 0") image_list = list(images) if len(image_list) == 0: return torch.empty( (0, self.dino_dim), dtype=torch.float32, device=self.device ) if apply_sam: processed_images = self._segment_with_sam_dataset(image_list) else: processed_images = image_list all_features: list[torch.Tensor] = [] for i in range(0, len(processed_images), batch_size): batch_images = processed_images[i : i + batch_size] tokens = self._dino_forward_batch(batch_images) features = tokens.mean(dim=1) all_features.append(F.normalize(features, dim=-1)) return torch.cat(all_features, dim=0)