mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-13 04:25:32 +08:00
refactor(compressors): consolidate pipeline and improve mask handling
This commit is contained in:
@@ -1,14 +1,5 @@
|
||||
"""SAM + DINO + Hash compression pipeline."""
|
||||
|
||||
from utils import get_device
|
||||
from utils.model import (
|
||||
get_dino_dim,
|
||||
load_dino_model,
|
||||
load_hash_compressor,
|
||||
load_sam_model,
|
||||
)
|
||||
from utils.image import extract_masked_region, segment_image
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
@@ -16,15 +7,24 @@ 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
|
||||
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
|
||||
config: Configuration object with model settings.
|
||||
|
||||
Returns:
|
||||
Initialized HashPipeline
|
||||
Initialized HashPipeline.
|
||||
"""
|
||||
return HashPipeline(
|
||||
dino_model=config.model.dino_model,
|
||||
@@ -38,21 +38,15 @@ def create_pipeline_from_config(config) -> "HashPipeline":
|
||||
|
||||
|
||||
class HashPipeline(nn.Module):
|
||||
"""Pipeline: SAM segmentation + DINO features + Hash compression.
|
||||
"""Pipeline for SAM segmentation + DINO features + Hash compression.
|
||||
|
||||
Pipeline flow:
|
||||
PIL Image -> SAM (largest object mask) -> DINO (features) -> Hash (binary codes)
|
||||
|
||||
Usage:
|
||||
# Initialize with config
|
||||
pipeline = HashPipeline(
|
||||
dino_model="facebook/dinov2-large",
|
||||
hash_bits=512,
|
||||
)
|
||||
|
||||
# Process image
|
||||
Example:
|
||||
pipeline = HashPipeline(dino_model="facebook/dinov2-large", hash_bits=512)
|
||||
image = Image.open("path/to/image.jpg")
|
||||
hash_bits = pipeline(image) # [1, 512] binary bits
|
||||
hash_bits = pipeline(image) # Returns [1, 512] binary bits
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -65,38 +59,25 @@ class HashPipeline(nn.Module):
|
||||
hash_bits: int = 512,
|
||||
compressor_path: Optional[str] = None,
|
||||
):
|
||||
"""Initialize the pipeline.
|
||||
|
||||
Args:
|
||||
dino_model: DINOv2 model name from HuggingFace
|
||||
sam_model: SAM2.1 model name from HuggingFace
|
||||
sam_min_mask_area: Minimum area threshold for valid SAM masks
|
||||
sam_max_masks: Maximum number of SAM masks to keep
|
||||
sam_points_per_batch: Prompt points batch size for SAM2 mask generation
|
||||
sam_checkpoint_dir: Optional local cache directory for SAM2 weights
|
||||
hash_bits: Number of bits in hash code
|
||||
compressor_path: Optional path to trained HashCompressor weights
|
||||
device: Device to run models on
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# Auto detect device
|
||||
# Device for model placement.
|
||||
self.device = get_device()
|
||||
|
||||
self.dino_model = dino_model
|
||||
# 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)
|
||||
|
||||
# Determine DINO feature dimension
|
||||
# DINO feature dimension based on model size.
|
||||
self.dino_dim = get_dino_dim(dino_model)
|
||||
|
||||
# Initialize HashCompressor
|
||||
# Hash compressor for binarizing DINO features.
|
||||
self.hash_compressor = load_hash_compressor(
|
||||
input_dim=self.dino_dim,
|
||||
hash_bits=hash_bits,
|
||||
@@ -104,18 +85,21 @@ class HashPipeline(nn.Module):
|
||||
)
|
||||
|
||||
@property
|
||||
def hash_bits(self):
|
||||
"""Return the number of hash bits."""
|
||||
def hash_bits(self) -> int:
|
||||
"""Number of bits in the hash code."""
|
||||
return self.hash_compressor.hash_bits
|
||||
|
||||
def _prepare_image_for_encoding(
|
||||
self,
|
||||
image: Image.Image,
|
||||
apply_sam: bool,
|
||||
) -> Image.Image:
|
||||
if not apply_sam:
|
||||
return image
|
||||
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,
|
||||
@@ -128,62 +112,45 @@ class HashPipeline(nn.Module):
|
||||
|
||||
return extract_masked_region(image, masks[0]["segment"])
|
||||
|
||||
def _encode_image(self, image: Image.Image, apply_sam: bool) -> torch.Tensor:
|
||||
image_for_encoding = self._prepare_image_for_encoding(
|
||||
image, apply_sam=apply_sam
|
||||
)
|
||||
inputs = self.processor(image_for_encoding, return_tensors="pt").to(self.device)
|
||||
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)
|
||||
tokens = outputs.last_hidden_state
|
||||
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(self, image: Image.Image) -> torch.Tensor:
|
||||
"""Process a single image through the pipeline.
|
||||
|
||||
Args:
|
||||
image: Input PIL Image
|
||||
|
||||
Returns:
|
||||
Binary hash codes [1, hash_bits] as int32
|
||||
"""
|
||||
return self._encode_image(image, apply_sam=True)
|
||||
|
||||
def encode_masked_region(self, image: Image.Image) -> torch.Tensor:
|
||||
"""Encode a pre-masked region using DINO+Hash without SAM stage."""
|
||||
return self._encode_image(image, apply_sam=False)
|
||||
|
||||
def encode(self, image: Image.Image) -> torch.Tensor:
|
||||
"""Encode an image to binary hash bits.
|
||||
|
||||
Alias for forward().
|
||||
|
||||
Args:
|
||||
image: Input PIL Image
|
||||
|
||||
Returns:
|
||||
Binary hash codes [1, hash_bits] as int32
|
||||
"""
|
||||
return self.forward(image)
|
||||
|
||||
def extract_features(self, image: Image.Image) -> torch.Tensor:
|
||||
"""Extract DINO features from an image.
|
||||
"""Extract normalized DINO features from an image.
|
||||
|
||||
Args:
|
||||
image: Input PIL Image
|
||||
image: Input PIL Image.
|
||||
|
||||
Returns:
|
||||
DINO features [1, dino_dim], normalized
|
||||
Normalized DINO features of shape [1, dino_dim].
|
||||
"""
|
||||
image_for_encoding = self._prepare_image_for_encoding(image, apply_sam=True)
|
||||
inputs = self.processor(image_for_encoding, return_tensors="pt").to(self.device)
|
||||
|
||||
with torch.no_grad():
|
||||
outputs = self.dino(**inputs)
|
||||
features = outputs.last_hidden_state.mean(dim=1) # [1, dim]
|
||||
features = F.normalize(features, dim=-1)
|
||||
|
||||
return features
|
||||
image = self._segment_with_sam(image)
|
||||
tokens = self._dino_forward(image)
|
||||
features = tokens.mean(dim=1)
|
||||
return F.normalize(features, dim=-1)
|
||||
|
||||
Reference in New Issue
Block a user