SikongJueluo/Mini-Nav
1
0
Fork 0
mirror of https://github.com/SikongJueluo/Mini-Nav.git synced 2026-07-12 20:15:31 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(compressors): consolidate pipeline and improve mask handling

Browse Source
This commit is contained in:
SikongJueluo
2026-03-26 19:00:13 +08:00
parent 90d5a8f08a
commit 968819e113
11 changed files with 302 additions and 121 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
mini-nav/compressors/__init__.py
View File

@@ -1,6 +1,12 @@
from .common import BinarySign, bits_to_hash, hamming_distance, hamming_similarity, hash_to_bits
from .common import (
BinarySign,
bits_to_hash,
hamming_distance,
hamming_similarity,
hash_to_bits,
)
from .hash_compressor import HashCompressor, HashLoss, VideoPositiveMask
from .pipeline import HashPipeline, SAMHashPipeline, create_pipeline_from_config
from .pipeline import HashPipeline, create_pipeline_from_config
from .train import train
__all__ = [
@@ -9,7 +15,6 @@ __all__ = [
"HashLoss",
"VideoPositiveMask",
"HashPipeline",
"SAMHashPipeline", # Backward compatibility alias
"create_pipeline_from_config",
"BinarySign",
"hamming_distance",

155
mini-nav/compressors/pipeline.py
View File

@@ -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)