mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-13 04:25:32 +08:00
SikongJueluo
42acb3ee1b
- 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
286 lines
8.8 KiB
Python
286 lines
8.8 KiB
Python
"""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 .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":
|
|
"""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,
|
|
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,
|
|
hash_bits: int = 512,
|
|
compressor_path: Optional[str] = None,
|
|
):
|
|
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
|
|
|
|
# Load models.
|
|
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)
|
|
|
|
# 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, 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 best object, or original if no masks.
|
|
"""
|
|
masks = generate_proposals(
|
|
self.sam_model,
|
|
self.sam_processor,
|
|
image,
|
|
bboxes,
|
|
)
|
|
|
|
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(
|
|
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.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.
|
|
|
|
Args:
|
|
image: Input PIL Image.
|
|
|
|
Returns:
|
|
Last hidden state tokens of shape [1, N, dim].
|
|
"""
|
|
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.dino_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, 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, bboxes)
|
|
tokens = self._dino_forward(image)
|
|
_, _, bits = self.hash_compressor(tokens)
|
|
return bits
|
|
|
|
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:
|
|
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, bboxes_per_image
|
|
)
|
|
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, 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)
|
|
|
|
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:
|
|
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, bboxes_per_image
|
|
)
|
|
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)
|
|
|
|
|
|
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]
|