From 42acb3ee1bd3bf8971c2a86b4c563752f5375dd8 Mon Sep 17 00:00:00 2001 From: SikongJueluo Date: Thu, 2 Apr 2026 16:42:59 +0800 Subject: [PATCH] refactor(compressors): switch SAM from automatic mask generation to bbox-prompted segmentation - Replace SAM2AutomaticMaskGenerator pipeline with Sam2Processor+Sam2Model - Freeze SAM model parameters at load time, removing torch.no_grad() at call sites - Rewrite proposal/core.py to use bbox prompts instead of automatic point sampling - Add bboxes parameter to all HashPipeline public methods (forward, forward_dataset, extract_features, extract_features_dataset) - Extract mask filtering logic (_filter_masks) from proposal into pipeline - Rename object_score/ to filter/ - Add load_owlv2_model to model_loader - Rename notebooks/test.py to habitat_sim_setup.py --- .../{object_score => filter}/__init__.py | 0 .../{object_score => filter}/config.py | 0 .../{object_score => filter}/features.py | 0 .../{object_score => filter}/scorer.py | 0 .../{object_score => filter}/selector.py | 0 mini-nav/compressors/model_loader.py | 41 +++- mini-nav/compressors/pipeline.py | 95 +++++---- mini-nav/compressors/proposal/core.py | 192 +++++++----------- notebooks/{test.py => habitat_sim_setup.py} | 0 9 files changed, 161 insertions(+), 167 deletions(-) rename mini-nav/compressors/{object_score => filter}/__init__.py (100%) rename mini-nav/compressors/{object_score => filter}/config.py (100%) rename mini-nav/compressors/{object_score => filter}/features.py (100%) rename mini-nav/compressors/{object_score => filter}/scorer.py (100%) rename mini-nav/compressors/{object_score => filter}/selector.py (100%) rename notebooks/{test.py => habitat_sim_setup.py} (100%) diff --git a/mini-nav/compressors/object_score/__init__.py b/mini-nav/compressors/filter/__init__.py similarity index 100% rename from mini-nav/compressors/object_score/__init__.py rename to mini-nav/compressors/filter/__init__.py diff --git a/mini-nav/compressors/object_score/config.py b/mini-nav/compressors/filter/config.py similarity index 100% rename from mini-nav/compressors/object_score/config.py rename to mini-nav/compressors/filter/config.py diff --git a/mini-nav/compressors/object_score/features.py b/mini-nav/compressors/filter/features.py similarity index 100% rename from mini-nav/compressors/object_score/features.py rename to mini-nav/compressors/filter/features.py diff --git a/mini-nav/compressors/object_score/scorer.py b/mini-nav/compressors/filter/scorer.py similarity index 100% rename from mini-nav/compressors/object_score/scorer.py rename to mini-nav/compressors/filter/scorer.py diff --git a/mini-nav/compressors/object_score/selector.py b/mini-nav/compressors/filter/selector.py similarity index 100% rename from mini-nav/compressors/object_score/selector.py rename to mini-nav/compressors/filter/selector.py diff --git a/mini-nav/compressors/model_loader.py b/mini-nav/compressors/model_loader.py index f778905..15aae84 100644 --- a/mini-nav/compressors/model_loader.py +++ b/mini-nav/compressors/model_loader.py @@ -1,11 +1,17 @@ -"""Model loading utilities for DINO, SAM2 and HashCompressor.""" +"""Model loading utilities for DINO, SAM2, OWLv2 and HashCompressor.""" from typing import TYPE_CHECKING, Any - import torch -from transformers import AutoImageProcessor, AutoModel, pipeline, MaskGenerationPipeline - +from transformers import ( + AutoImageProcessor, + AutoModel, + Owlv2ForObjectDetection, + Owlv2Model, + Owlv2Processor, + Sam2Model, + Sam2Processor, +) from utils import get_device if TYPE_CHECKING: @@ -14,14 +20,17 @@ if TYPE_CHECKING: def load_sam_model( model_name: str = "facebook/sam2.1-hiera-large", -) -> MaskGenerationPipeline: +) -> tuple[Sam2Processor, Sam2Model]: + """Load SAM2 processor and model with frozen parameters.""" device = get_device() - return pipeline( - task="mask-generation", - model=model_name, - device=device, - ) + processor = Sam2Processor.from_pretrained(model_name) + model = Sam2Model.from_pretrained(model_name).to(device) + model.eval() + for param in model.parameters(): + param.requires_grad = False + + return processor, model def load_dino_model( @@ -36,6 +45,18 @@ def load_dino_model( return processor, dino +def load_owlv2_model( + model_name: str = "google/owlv2-base-patch16-ensemble", +) -> tuple[Owlv2Processor, Owlv2Model]: + device = get_device() + + processor = Owlv2Processor.from_pretrained(model_name) + model = Owlv2ForObjectDetection.from_pretrained(model_name).to(device) + model.eval() + + return processor, model + + def get_dino_dim(model_name: str) -> int: if "large" in model_name.lower(): return 1024 diff --git a/mini-nav/compressors/pipeline.py b/mini-nav/compressors/pipeline.py index 9901e43..9a59c6d 100644 --- a/mini-nav/compressors/pipeline.py +++ b/mini-nav/compressors/pipeline.py @@ -6,20 +6,20 @@ import torch import torch.nn as nn import torch.nn.functional as F from PIL import Image - -from .object_score import select_best_mask -from .proposal import ( - extract_masked_region, - generate_proposals, - generate_proposals_batch, -) from utils import get_device + +from .filter import select_best_mask from .model_loader import ( get_dino_dim, load_dino_model, load_hash_compressor, load_sam_model, ) +from .proposal import ( + extract_masked_region, + generate_proposals, + generate_proposals_batch, +) def create_pipeline_from_config(config) -> "HashPipeline": @@ -36,7 +36,6 @@ def create_pipeline_from_config(config) -> "HashPipeline": 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, ) @@ -60,7 +59,6 @@ class HashPipeline(nn.Module): 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, ): @@ -69,15 +67,13 @@ class HashPipeline(nn.Module): # Device for model placement. self.device = get_device() - # SAM2 settings. - self.sam_model_name = sam_model + # SAM2 filter settings. 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) + 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. self.dino_dim = get_dino_dim(dino_model) @@ -94,25 +90,27 @@ 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) -> Image.Image: - """Segment image with SAM and extract the largest object mask. - - If no valid masks are found, returns the original image. + def _segment_with_sam( + self, image: Image.Image, bboxes: list[list[float]] + ) -> Image.Image: + """Segment image with SAM and extract the best object mask. Args: image: Input PIL Image. + bboxes: Bounding boxes from object detector as [[x1,y1,x2,y2], ...]. Returns: - Masked image containing only the largest object, or original if no masks. + Masked image containing only the best object, or original if no masks. """ masks = generate_proposals( - self.mask_generator, + self.sam_model, + self.sam_processor, image, - min_area=self.sam_min_mask_area, - max_masks=self.sam_max_masks, - points_per_batch=self.sam_points_per_batch, + bboxes, ) + masks = _filter_masks(masks, self.sam_min_mask_area, self.sam_max_masks) + if not masks: return image @@ -124,15 +122,19 @@ class HashPipeline(nn.Module): def _segment_with_sam_dataset( self, images: Sequence[Image.Image], + bboxes_per_image: list[list[list[float]]], ) -> list[Image.Image]: image_list = list(images) masks_dataset = generate_proposals_batch( - self.mask_generator, + self.sam_model, + self.sam_processor, image_list, - min_area=self.sam_min_mask_area, - max_masks=self.sam_max_masks, - points_per_batch=self.sam_points_per_batch, + 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: @@ -156,14 +158,14 @@ class HashPipeline(nn.Module): Returns: Last hidden state tokens of shape [1, N, dim]. """ - inputs = self.processor(image, return_tensors="pt").to(self.device) + inputs = self.dino_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( + inputs = self.dino_processor(images=list(images), return_tensors="pt").to( self.device ) @@ -171,16 +173,17 @@ class HashPipeline(nn.Module): outputs = self.dino(**inputs) return outputs.last_hidden_state - def forward(self, image: Image.Image) -> torch.Tensor: + def forward(self, image: Image.Image, bboxes: list[list[float]]) -> torch.Tensor: """Process a single image through the full pipeline. Args: image: Input PIL Image. + bboxes: Bounding boxes from object detector as [[x1,y1,x2,y2], ...]. Returns: Binary hash codes of shape [1, hash_bits] as int32. """ - image = self._segment_with_sam(image) + image = self._segment_with_sam(image, bboxes) tokens = self._dino_forward(image) _, _, bits = self.hash_compressor(tokens) return bits @@ -188,6 +191,7 @@ class HashPipeline(nn.Module): def forward_dataset( self, images: Sequence[Image.Image], + bboxes_per_image: list[list[list[float]]], batch_size: int = 32, apply_sam: bool = True, ) -> torch.Tensor: @@ -202,7 +206,9 @@ class HashPipeline(nn.Module): ) if apply_sam: - processed_images = self._segment_with_sam_dataset(image_list) + processed_images = self._segment_with_sam_dataset( + image_list, bboxes_per_image + ) else: processed_images = image_list @@ -215,16 +221,19 @@ class HashPipeline(nn.Module): return torch.cat(batch_bits, dim=0) - def extract_features(self, image: Image.Image) -> torch.Tensor: + 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) + image = self._segment_with_sam(image, bboxes) tokens = self._dino_forward(image) features = tokens.mean(dim=1) return F.normalize(features, dim=-1) @@ -232,6 +241,7 @@ class HashPipeline(nn.Module): def extract_features_dataset( self, images: Sequence[Image.Image], + bboxes_per_image: list[list[list[float]]], batch_size: int = 32, apply_sam: bool = True, ) -> torch.Tensor: @@ -246,7 +256,9 @@ class HashPipeline(nn.Module): ) if apply_sam: - processed_images = self._segment_with_sam_dataset(image_list) + processed_images = self._segment_with_sam_dataset( + image_list, bboxes_per_image + ) else: processed_images = image_list @@ -258,3 +270,16 @@ class HashPipeline(nn.Module): 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] diff --git a/mini-nav/compressors/proposal/core.py b/mini-nav/compressors/proposal/core.py index a2a9c0f..486f348 100644 --- a/mini-nav/compressors/proposal/core.py +++ b/mini-nav/compressors/proposal/core.py @@ -1,27 +1,26 @@ -"""SAM mask proposal generation.""" +"""SAM mask proposal generation via bounding box prompts.""" from typing import Any, Sequence -import torch import numpy as np from PIL import Image +from transformers import Sam2Model, Sam2Processor +from utils import get_device def generate_proposals( - mask_generator: Any, + model: Sam2Model, + processor: Sam2Processor, image: Image.Image, - min_area: int = 32 * 32, - max_masks: int = 5, - points_per_batch: int = 64, + bboxes: list[list[float]], ) -> list[dict[str, Any]]: - """Segment image using SAM to extract object masks. + """Segment regions in image using SAM2 with bounding box prompts. Args: - mask_generator: SAM2 mask generator. + model: Sam2Model instance. + processor: Sam2Processor instance. image: PIL Image to segment. - min_area: Minimum mask area threshold in pixels. - max_masks: Maximum number of masks to return. - points_per_batch: Number of prompt points to process in each batch. + bboxes: Bounding boxes as [[x1, y1, x2, y2], ...]. Returns: List of mask dictionaries with keys: @@ -31,28 +30,40 @@ def generate_proposals( - predicted_iou: Model's confidence in the mask - stability_score: Stability score for the mask """ + if not bboxes: + return [] + + device = get_device() image_rgb = image.convert("RGB") - raw_output = mask_generator(image_rgb, points_per_batch=points_per_batch) - return _normalize_and_filter_masks( - raw_output, min_area=min_area, max_masks=max_masks - ) + input_boxes = [bboxes] + inputs = processor( + images=image_rgb, + input_boxes=input_boxes, + return_tensors="pt", + ).to(device) + + outputs = model(**inputs, multimask_output=False) + masks = processor.post_process_masks( + outputs.pred_masks.cpu(), + inputs["original_sizes"], + )[0] + + return _masks_to_proposals(masks) def generate_proposals_batch( - mask_generator: Any, + model: Sam2Model, + processor: Sam2Processor, images: Sequence[Image.Image], - min_area: int = 32 * 32, - max_masks: int = 5, - points_per_batch: int = 64, + bboxes_per_image: list[list[list[float]]], ) -> list[list[dict[str, Any]]]: - """Segment a batch of images using SAM. + """Segment a batch of images using SAM2 with bounding box prompts. Args: - mask_generator: SAM2 mask generator. + model: Sam2Model instance. + processor: Sam2Processor instance. images: Sequence of PIL Images to segment. - min_area: Minimum mask area threshold in pixels. - max_masks: Maximum number of masks to return per image. - points_per_batch: Number of prompt points to process in each batch. + bboxes_per_image: Bounding boxes per image, outer list matches images length. Returns: List of lists of mask dictionaries, one inner list per image. @@ -61,90 +72,48 @@ def generate_proposals_batch( if not image_list: return [] - image_rgb_list = [image.convert("RGB") for image in image_list] - raw_batch_output = mask_generator( - image_rgb_list, - points_per_batch=points_per_batch, - ) - batch_items = _split_batch_output(raw_batch_output, expected_size=len(image_list)) - if batch_items is not None: - return [ - _normalize_and_filter_masks( - batch_item, - min_area=min_area, - max_masks=max_masks, - ) - for batch_item in batch_items - ] + device = get_device() + image_rgb_list = [img.convert("RGB") for img in image_list] - return [ - _normalize_and_filter_masks( - mask_generator(image_rgb, points_per_batch=points_per_batch), - min_area=min_area, - max_masks=max_masks, - ) - for image_rgb in image_rgb_list - ] + inputs = processor( + images=image_rgb_list, + input_boxes=bboxes_per_image, + return_tensors="pt", + ).to(device) - -def _split_batch_output(raw_output: Any, expected_size: int) -> list[Any] | None: - """Attempt to split raw batch output into per-image results.""" - if isinstance(raw_output, list): - if len(raw_output) == expected_size: - return raw_output - return None - - if isinstance(raw_output, dict): - raw_masks = raw_output.get("masks", raw_output) - if isinstance(raw_masks, list) and len(raw_masks) == expected_size: - return raw_masks - - return None - - -def _normalize_and_filter_masks( - raw_output: Any, - min_area: int, - max_masks: int, -) -> list[dict[str, Any]]: - """Normalize raw SAM output into mask dicts and filter by area/count.""" - raw_masks = ( - raw_output.get("masks", raw_output) - if isinstance(raw_output, dict) - else raw_output + outputs = model(**inputs, multimask_output=False) + all_masks = processor.post_process_masks( + outputs.pred_masks.cpu(), + inputs["original_sizes"], ) - normalized_masks: list[dict[str, Any]] = [] - if isinstance(raw_masks, list): - if raw_masks and isinstance(raw_masks[0], dict): - normalized_masks = raw_masks - else: - for mask_like in raw_masks: - mask_dict = _to_mask_dict(mask_like) - if mask_dict is not None: - normalized_masks.append(mask_dict) - else: - mask_array = _to_numpy_mask_array(raw_masks) - if mask_array is not None: - if mask_array.ndim == 2: - mask_array = np.expand_dims(mask_array, axis=0) - if mask_array.ndim == 3: - for single_mask in mask_array: - mask_dict = _to_mask_dict(single_mask) - if mask_dict is not None: - normalized_masks.append(mask_dict) + return [_masks_to_proposals(image_masks) for image_masks in all_masks] - if not normalized_masks: + +def _masks_to_proposals(masks: Any) -> list[dict[str, Any]]: + """Convert model output masks to list of mask dicts.""" + mask_array = _to_numpy_mask_array(masks) + if mask_array is None: return [] - filtered_masks = [ - mask for mask in normalized_masks if int(mask["area"]) >= min_area - ] - if not filtered_masks: + # Ensure 3D: [num_masks, H, W] + if mask_array.ndim == 2: + mask_array = np.expand_dims(mask_array, axis=0) + + if mask_array.ndim != 3: return [] - sorted_masks = sorted(filtered_masks, key=lambda mask: mask["area"], reverse=True) - return sorted_masks[:max_masks] + # Remove batch dim if present: [1, num_masks, H, W] → [num_masks, H, W] + if mask_array.ndim == 3 and mask_array.shape[0] == 1: + mask_array = mask_array[0] + + proposals: list[dict[str, Any]] = [] + for single_mask in mask_array: + mask_dict = _build_mask_dict(single_mask) + if mask_dict is not None: + proposals.append(mask_dict) + + return proposals def _to_numpy_mask_array(mask_like: Any) -> np.ndarray | None: @@ -154,35 +123,14 @@ def _to_numpy_mask_array(mask_like: Any) -> np.ndarray | None: if isinstance(mask_like, np.ndarray): return mask_like + import torch + if isinstance(mask_like, torch.Tensor): return mask_like.detach().cpu().numpy() return None -def _to_mask_dict(mask_like: Any) -> dict[str, Any] | None: - """Convert a single mask-like object to a standardized mask dict.""" - if isinstance(mask_like, dict): - if "area" in mask_like and "bbox" in mask_like and "segment" in mask_like: - return mask_like - - segment = mask_like.get("segment") - if segment is None and "mask" in mask_like: - segment = mask_like["mask"] - if segment is None: - return None - - mask_array = _to_numpy_mask_array(segment) - if mask_array is None: - return None - return _build_mask_dict(mask_array) - - mask_array = _to_numpy_mask_array(mask_like) - if mask_array is None: - return None - return _build_mask_dict(mask_array) - - def _build_mask_dict(mask_array: np.ndarray) -> dict[str, Any] | None: """Build a mask dictionary from a 2D boolean numpy array.""" if mask_array.ndim != 2: diff --git a/notebooks/test.py b/notebooks/habitat_sim_setup.py similarity index 100% rename from notebooks/test.py rename to notebooks/habitat_sim_setup.py