mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-12 20:15:31 +08:00
244 lines
7.5 KiB
Python
244 lines
7.5 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 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)
|