From 4e16e38f32bc73e33b94a49f4601322f91ff72d5 Mon Sep 17 00:00:00 2001 From: SikongJueluo Date: Fri, 3 Apr 2026 15:49:23 +0800 Subject: [PATCH] refactor(compressors): migrate pipeline to OWLv2-based detection with text labels - Replace bbox-prompted segmentation with OWLv2 text-guided object detection - Refactor HashPipeline from nn.Module to plain class with modular stage methods - Add detect_batch, segment_batch, filter_batch for explicit pipeline stages - Rename forward to forward_batch with text_labels API instead of bboxes - Add mask_scoring_config, score_threshold, postprocess_threshold configuration - Update model_loader to expose Dinov2Model type annotation --- mini-nav/compressors/__init__.py | 2 +- mini-nav/compressors/model_loader.py | 3 +- mini-nav/compressors/pipeline.py | 302 +++++++++++++++------------ 3 files changed, 166 insertions(+), 141 deletions(-) diff --git a/mini-nav/compressors/__init__.py b/mini-nav/compressors/__init__.py index 3e55583..2ca6126 100644 --- a/mini-nav/compressors/__init__.py +++ b/mini-nav/compressors/__init__.py @@ -6,7 +6,7 @@ from .common import ( hash_to_bits, ) from .hash_compressor import HashCompressor, HashLoss, VideoPositiveMask -from .object_score import ( +from .filter import ( MaskFeatures, MaskScoringConfig, compute_mask_features, diff --git a/mini-nav/compressors/model_loader.py b/mini-nav/compressors/model_loader.py index e8b4b59..6100cd8 100644 --- a/mini-nav/compressors/model_loader.py +++ b/mini-nav/compressors/model_loader.py @@ -6,6 +6,7 @@ import torch from transformers import ( AutoImageProcessor, AutoModel, + Dinov2Model, Owlv2ForObjectDetection, Owlv2Processor, Sam2Model, @@ -34,7 +35,7 @@ def load_sam_model( def load_dino_model( model_name: str = "facebook/dinov2-large", -) -> tuple[AutoImageProcessor, AutoModel]: +) -> tuple[AutoImageProcessor, Dinov2Model]: device = get_device() processor = AutoImageProcessor.from_pretrained(model_name) diff --git a/mini-nav/compressors/pipeline.py b/mini-nav/compressors/pipeline.py index 9a59c6d..db099de 100644 --- a/mini-nav/compressors/pipeline.py +++ b/mini-nav/compressors/pipeline.py @@ -1,25 +1,26 @@ -"""SAM + DINO + Hash compression pipeline.""" +"""OWLv2 + 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 .filter import select_best_mask +from .filter import MaskScoringConfig, select_best_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, extract_masked_region, - generate_proposals, generate_proposals_batch, ) +from .proposal.core import DetectionResult def create_pipeline_from_config(config) -> "HashPipeline": @@ -32,50 +33,64 @@ def create_pipeline_from_config(config) -> "HashPipeline": Initialized 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, - sam_min_mask_area=config.model.sam_min_mask_area, - sam_max_masks=config.model.sam_max_masks, 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), ) -class HashPipeline(nn.Module): - """Pipeline for SAM segmentation + DINO features + Hash compression. +class HashPipeline: + """Pipeline for OWLv2 detection + SAM segmentation + DINO features + Hash compression. Pipeline flow: - PIL Image -> SAM (largest object mask) -> DINO (features) -> Hash (binary codes) + Images + Text Labels -> OWLv2 (detections) -> SAM (masks) -> Filter (best 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 + images = [Image.open("path/to/image.jpg")] + text_labels = ["object"] + hash_bits = pipeline.forward_batch(images, text_labels) # Returns [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", - sam_min_mask_area: int = 100, - sam_max_masks: int = 10, 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__() # Device for model placement. self.device = get_device() - # SAM2 filter settings. - self.sam_min_mask_area = sam_min_mask_area - self.sam_max_masks = sam_max_masks + # OWLv2 detection settings. + self.owlv2_processor, self.owlv2_model = load_owlv2_model( + model_name=owlv2_model + ) + self.score_threshold = score_threshold + self.postprocess_threshold = postprocess_threshold - # Load models. + # Mask scoring config for filter step. + self.mask_scoring_config = mask_scoring_config + + # SAM2 model for segmentation. 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) - # DINO feature dimension based on model size. + # DINO model for feature extraction. + self.dino_processor, self.dino = load_dino_model(model_name=dino_model) self.dino_dim = get_dino_dim(dino_model) # Hash compressor for binarizing DINO features. @@ -90,82 +105,112 @@ class HashPipeline(nn.Module): """Number of bits in the hash code.""" return self.hash_compressor.hash_bits - def _segment_with_sam( - self, image: Image.Image, bboxes: list[list[float]] - ) -> Image.Image: - """Segment image with SAM and extract the best object mask. + def detect_batch( + self, + images: Sequence[Image.Image], + text_labels: list[str], + ) -> list[list[DetectionResult]]: + """Detect objects in a batch of images using OWLv2. Args: - image: Input PIL Image. - bboxes: Bounding boxes from object detector as [[x1,y1,x2,y2], ...]. + images: Sequence of PIL Images. + text_labels: Text labels for all images (same labels used for each image). Returns: - Masked image containing only the best object, or original if no masks. + List of lists of DetectionResult dicts, one inner list per image. """ - masks = generate_proposals( - self.sam_model, - self.sam_processor, - image, - bboxes, + 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, ) - masks = _filter_masks(masks, self.sam_min_mask_area, self.sam_max_masks) - - if not masks: - return image - - best_mask = select_best_mask(masks, image_shape=(image.height, image.width)) - if best_mask is None: - return image - return extract_masked_region(image, best_mask["segment"]) - - def _segment_with_sam_dataset( + def segment_batch( self, images: Sequence[Image.Image], bboxes_per_image: list[list[list[float]]], - ) -> list[Image.Image]: + ) -> list[list[dict]]: + """Segment objects in images using SAM2 with bounding box prompts. + + Args: + images: Sequence of PIL Images. + bboxes_per_image: Bounding boxes per image as [[[x1,y1,x2,y2], ...], ...]. + + Returns: + List of lists of mask dictionaries, one inner list per image. + """ image_list = list(images) - masks_dataset = generate_proposals_batch( + if not image_list: + return [] + + return generate_proposals_batch( self.sam_model, self.sam_processor, image_list, bboxes_per_image, ) - masks_dataset = [ - _filter_masks(masks, self.sam_min_mask_area, self.sam_max_masks) - for masks in masks_dataset - ] - selected_images: list[Image.Image] = [] - for image, masks in zip(image_list, masks_dataset): - if not masks: - selected_images.append(image) - continue - best_mask = select_best_mask(masks, image_shape=(image.height, image.width)) - if best_mask is None: - selected_images.append(image) - continue - selected_images.append(extract_masked_region(image, best_mask["segment"])) - - return selected_images - - def _dino_forward(self, image: Image.Image) -> torch.Tensor: - """Extract DINO tokens from an image. + def filter_batch( + self, + images: Sequence[Image.Image], + masks_per_image: list[list[dict]], + ) -> list[Image.Image]: + """Filter masks and extract best masked region for each image. Args: - image: Input PIL Image. + images: Sequence of PIL Images. + masks_per_image: Masks per image from segment_batch. Returns: - Last hidden state tokens of shape [1, N, dim]. + List of PIL Images, one per input image (original if no valid masks). """ - inputs = self.dino_processor(image, return_tensors="pt").to(self.device) + image_list = list(images) + if not image_list: + return [] - with torch.no_grad(): - outputs = self.dino(**inputs) - return outputs.last_hidden_state + filtered_images: list[Image.Image] = [] + for image, masks in zip(image_list, masks_per_image): + if not masks: + filtered_images.append(image) + continue - def _dino_forward_batch(self, images: Sequence[Image.Image]) -> torch.Tensor: - inputs = self.dino_processor(images=list(images), return_tensors="pt").to( + best_mask = select_best_mask( + masks, + image_shape=(image.height, image.width), + config=self.mask_scoring_config, + ) + if best_mask is None: + filtered_images.append(image) + continue + + filtered_images.append(extract_masked_region(image, best_mask["segment"])) + + return filtered_images + + 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 ) @@ -173,113 +218,92 @@ class HashPipeline(nn.Module): outputs = self.dino(**inputs) return outputs.last_hidden_state - def forward(self, image: Image.Image, bboxes: list[list[float]]) -> torch.Tensor: - """Process a single image through the full pipeline. + def compress_batch(self, tokens: torch.Tensor) -> torch.Tensor: + """Compress DINO tokens to binary hash codes. Args: - image: Input PIL Image. - bboxes: Bounding boxes from object detector as [[x1,y1,x2,y2], ...]. + tokens: DINO tokens of shape [B, N, dim]. Returns: - Binary hash codes of shape [1, hash_bits] as int32. + Binary hash codes of shape [B, hash_bits] as int32. """ - image = self._segment_with_sam(image, bboxes) - tokens = self._dino_forward(image) _, _, bits = self.hash_compressor(tokens) return bits - def forward_dataset( + def forward_batch( self, images: Sequence[Image.Image], - bboxes_per_image: list[list[list[float]]], + text_labels: list[str], batch_size: int = 32, - apply_sam: bool = True, ) -> 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. + """ if batch_size <= 0: raise ValueError("batch_size must be greater than 0") image_list = list(images) - - if len(image_list) == 0: + if not image_list: 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, bboxes_per_image - ) - else: - processed_images = image_list + detections = self.detect_batch(image_list, text_labels) + bboxes = [[d["bbox"] for d in dets] for dets in detections] + masks = self.segment_batch(image_list, bboxes) + processed = self.filter_batch(image_list, masks) - 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) + all_bits: list[torch.Tensor] = [] + for i in range(0, len(processed), batch_size): + sub_batch = processed[i : i + batch_size] + tokens = self.extract_dino_batch(sub_batch) + bits = self.compress_batch(tokens) + all_bits.append(bits) - return torch.cat(batch_bits, dim=0) - - def extract_features( - self, image: Image.Image, bboxes: list[list[float]] - ) -> torch.Tensor: - """Extract normalized DINO features from an image. - - Args: - image: Input PIL Image. - bboxes: Bounding boxes from object detector as [[x1,y1,x2,y2], ...]. - - Returns: - Normalized DINO features of shape [1, dino_dim]. - """ - image = self._segment_with_sam(image, bboxes) - tokens = self._dino_forward(image) - features = tokens.mean(dim=1) - return F.normalize(features, dim=-1) + return torch.cat(all_bits, dim=0) def extract_features_dataset( self, images: Sequence[Image.Image], - bboxes_per_image: list[list[list[float]]], + text_labels: list[str], batch_size: int = 32, - apply_sam: bool = True, ) -> 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 len(image_list) == 0: + if not image_list: 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, bboxes_per_image - ) - else: - processed_images = image_list + detections = self.detect_batch(image_list, text_labels) + bboxes = [[d["bbox"] for d in dets] for dets in detections] + masks = self.segment_batch(image_list, bboxes) + processed = self.filter_batch(image_list, masks) 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) + for i in range(0, len(processed), batch_size): + sub_batch = processed[i : i + 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) - - -def _filter_masks( - masks: list[dict], - min_area: int, - max_masks: int, -) -> list[dict]: - """Filter masks by area and keep top-N largest.""" - filtered = [m for m in masks if int(m["area"]) >= min_area] - if not filtered: - return [] - sorted_masks = sorted(filtered, key=lambda m: m["area"], reverse=True) - return sorted_masks[:max_masks]