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(pipeline): integrate SAM segmentation and modularize model loading

Browse Source
This commit is contained in:
SikongJueluo
2026-03-24 21:52:02 +08:00
parent 9e6339e580
commit 90d5a8f08a
11 changed files with 437 additions and 172 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1,8 +1,13 @@
"""Hash compression pipeline with DINO feature extraction.
"""SAM + DINO + Hash compression pipeline."""
This pipeline extracts features using DINOv2 and compresses them
to binary hash codes using HashCompressor.
"""
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
@@ -10,7 +15,6 @@ import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from transformers import AutoImageProcessor, AutoModel
def create_pipeline_from_config(config) -> "HashPipeline":
@@ -24,17 +28,20 @@ def create_pipeline_from_config(config) -> "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,
device=config.model.device if config.model.device != "auto" else None,
)
class HashPipeline(nn.Module):
"""Pipeline: DINO features + Hash compression.
"""Pipeline: SAM segmentation + DINO features + Hash compression.
Pipeline flow:
PIL Image -> DINO (features) -> Hash (binary codes)
PIL Image -> SAM (largest object mask) -> DINO (features) -> Hash (binary codes)
Usage:
# Initialize with config
@@ -51,14 +58,22 @@ class HashPipeline(nn.Module):
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,
device: 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
@@ -66,53 +81,66 @@ class HashPipeline(nn.Module):
super().__init__()
# Auto detect device
if device is None:
device = "cuda" if torch.cuda.is_available() else "cpu"
self.device = torch.device(device)
self.device = get_device()
self.dino_model = dino_model
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
# Initialize DINO processor and model
self.processor = AutoImageProcessor.from_pretrained(dino_model)
self.dino = AutoModel.from_pretrained(dino_model).to(self.device)
self.dino.eval()
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
self.dino_dim = 1024 if "large" in dino_model else 768
self.dino_dim = get_dino_dim(dino_model)
# Initialize HashCompressor
self.hash_compressor = nn.Module() # Placeholder, will be replaced
self._init_hash_compressor(hash_bits, compressor_path)
def _init_hash_compressor(
self, hash_bits: int, compressor_path: Optional[str] = None
):
"""Initialize the hash compressor module.
This is called during __init__ but we need to replace it properly.
"""
# Import here to avoid circular imports
from .hash_compressor import HashCompressor
compressor = HashCompressor(input_dim=self.dino_dim, hash_bits=hash_bits).to(
self.device
self.hash_compressor = load_hash_compressor(
input_dim=self.dino_dim,
hash_bits=hash_bits,
compressor_path=compressor_path,
)
# Load pretrained compressor if provided
if compressor_path is not None:
compressor.load_state_dict(
torch.load(compressor_path, map_location=self.device)
)
print(f"[OK] Loaded HashCompressor from {compressor_path}")
# Replace the placeholder
self.hash_compressor = compressor
@property
def hash_bits(self):
"""Return the number of hash bits."""
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
masks = segment_image(
self.mask_generator,
image,
min_area=self.sam_min_mask_area,
max_masks=self.sam_max_masks,
)
if not masks:
return image
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)
with torch.no_grad():
outputs = self.dino(**inputs)
tokens = outputs.last_hidden_state
_, _, bits = self.hash_compressor(tokens)
return bits
def forward(self, image: Image.Image) -> torch.Tensor:
"""Process a single image through the pipeline.
@@ -122,17 +150,11 @@ class HashPipeline(nn.Module):
Returns:
Binary hash codes [1, hash_bits] as int32
"""
# Extract DINO features
inputs = self.processor(image, return_tensors="pt").to(self.device)
return self._encode_image(image, apply_sam=True)
with torch.no_grad():
outputs = self.dino(**inputs)
tokens = outputs.last_hidden_state # [1, N, dim]
# Compress to hash codes
_, _, bits = self.hash_compressor(tokens)
return bits
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.
@@ -156,7 +178,8 @@ class HashPipeline(nn.Module):
Returns:
DINO features [1, dino_dim], normalized
"""
inputs = self.processor(image, return_tensors="pt").to(self.device)
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)
@@ -164,7 +187,3 @@ class HashPipeline(nn.Module):
features = F.normalize(features, dim=-1)
return features
# Backward compatibility alias
SAMHashPipeline = HashPipeline

3
mini-nav/configs/config.yaml
View File

@@ -1,10 +1,11 @@
model:
name: "facebook/dinov2-large"
dino_model: "facebook/dinov2-large"
compression_dim: 512
device: "auto" # auto-detect GPU
sam_model: "facebook/sam2.1-hiera-large" # SAM model name
sam_min_mask_area: 100 # Minimum mask area threshold
sam_max_masks: 10 # Maximum number of masks to keep
sam_points_per_batch: 64
compressor_path: null # Path to trained HashCompressor weights (optional)
output:

12
mini-nav/configs/models.py
View File

@@ -3,7 +3,7 @@
from pathlib import Path
from typing import Literal, Optional
from pydantic import BaseModel, ConfigDict, Field, field_validator
from pydantic import AliasChoices, BaseModel, ConfigDict, Field, field_validator
class ModelConfig(BaseModel):
@@ -11,7 +11,10 @@ class ModelConfig(BaseModel):
model_config = ConfigDict(extra="ignore")
dino_model: str = "facebook/dinov2-large"
dino_model: str = Field(
default="facebook/dinov2-large",
validation_alias=AliasChoices("dino_model", "name"),
)
compression_dim: int = Field(
default=512, gt=0, description="Output feature dimension"
)
@@ -26,6 +29,11 @@ class ModelConfig(BaseModel):
sam_max_masks: int = Field(
default=10, gt=0, description="Maximum number of masks to keep"
)
sam_points_per_batch: int = Field(
default=64,
gt=0,
description="SAM2 mask generation batch size for prompt points",
)
compressor_path: Optional[str] = Field(
default=None, description="Path to trained HashCompressor weights"
)

62
mini-nav/tests/test_compressors.py
View File

@@ -2,6 +2,7 @@
import pytest
import torch
from unittest.mock import Mock, patch
from compressors import (
BinarySign,
HashCompressor,
@@ -205,18 +206,54 @@ class TestVideoPositiveMask:
class TestHashPipeline:
"""Test suite for HashPipeline."""
def test_pipeline_init(self):
@patch("compressors.pipeline.load_sam_model")
@patch("compressors.pipeline.AutoModel.from_pretrained")
@patch("compressors.pipeline.AutoImageProcessor.from_pretrained")
def test_pipeline_init(
self,
mock_processor_from_pretrained,
mock_model_from_pretrained,
mock_load_sam_model,
):
"""Verify pipeline initializes all components."""
mock_processor_from_pretrained.return_value = Mock()
mock_model = Mock()
mock_model.to.return_value = mock_model
mock_model.eval.return_value = None
mock_model_from_pretrained.return_value = mock_model
mock_load_sam_model.return_value = (Mock(), Mock())
pipeline = HashPipeline(
dino_model="facebook/dinov2-large",
hash_bits=512,
)
assert pipeline.dino_model == "facebook/dinov2-large"
assert pipeline.sam_model_name == "facebook/sam2.1-hiera-large"
assert pipeline.dino_dim == 1024
mock_load_sam_model.assert_called_once()
def test_pipeline_hash_bits(self):
@patch("compressors.pipeline.load_sam_model")
@patch("compressors.pipeline.AutoModel.from_pretrained")
@patch("compressors.pipeline.AutoImageProcessor.from_pretrained")
def test_pipeline_hash_bits(
self,
mock_processor_from_pretrained,
mock_model_from_pretrained,
mock_load_sam_model,
):
"""Verify pipeline uses correct hash bits."""
mock_processor_from_pretrained.return_value = Mock()
mock_model = Mock()
mock_model.to.return_value = mock_model
mock_model.eval.return_value = None
mock_model_from_pretrained.return_value = mock_model
mock_load_sam_model.return_value = (Mock(), Mock())
pipeline = HashPipeline(hash_bits=256)
assert pipeline.hash_bits == 256
@@ -228,14 +265,33 @@ class TestHashPipeline:
class TestConfigIntegration:
"""Test suite for config integration with pipeline."""
def test_create_pipeline_from_config(self):
@patch("compressors.pipeline.load_sam_model")
@patch("compressors.pipeline.AutoModel.from_pretrained")
@patch("compressors.pipeline.AutoImageProcessor.from_pretrained")
def test_create_pipeline_from_config(
self,
mock_processor_from_pretrained,
mock_model_from_pretrained,
mock_load_sam_model,
):
"""Verify pipeline can be created from config."""
mock_processor_from_pretrained.return_value = Mock()
mock_model = Mock()
mock_model.to.return_value = mock_model
mock_model.eval.return_value = None
mock_model_from_pretrained.return_value = mock_model
mock_load_sam_model.return_value = (Mock(), Mock())
config = cfg_manager.load()
pipeline = create_pipeline_from_config(config)
assert isinstance(pipeline, HashPipeline)
assert pipeline.hash_bits == config.model.compression_dim
assert pipeline.sam_max_masks == config.model.sam_max_masks
assert pipeline.sam_min_mask_area == config.model.sam_min_mask_area
def test_config_settings(self):
"""Verify config contains required settings."""

43
mini-nav/tests/test_sam.py
View File

@@ -97,9 +97,21 @@ class TestSAMSegmentation:
from utils.sam import segment_image
# Create masks with known areas (unordered)
mask1 = {"segment": np.ones((5, 5), dtype=bool), "area": 25, "bbox": [0, 0, 5, 5]}
mask2 = {"segment": np.ones((10, 10), dtype=bool), "area": 100, "bbox": [0, 0, 10, 10]}
mask3 = {"segment": np.ones((3, 3), dtype=bool), "area": 9, "bbox": [0, 0, 3, 3]}
mask1 = {
"segment": np.ones((5, 5), dtype=bool),
"area": 25,
"bbox": [0, 0, 5, 5],
}
mask2 = {
"segment": np.ones((10, 10), dtype=bool),
"area": 100,
"bbox": [0, 0, 10, 10],
}
mask3 = {
"segment": np.ones((3, 3), dtype=bool),
"area": 9,
"bbox": [0, 0, 3, 3],
}
mock_generator = Mock()
mock_generator.generate.return_value = [mask1, mask2, mask3]
@@ -117,6 +129,31 @@ class TestSAMSegmentation:
assert result[2]["area"] == 9
class TestSAMLoading:
@patch("utils.sam.pipeline")
def test_load_sam_model_uses_transformers_pipeline(self, mock_pipeline):
from utils.sam import Sam2MaskGenerator, load_sam_model
mock_pipe_obj = Mock()
mock_pipe_obj.model = Mock()
mock_pipeline.return_value = mock_pipe_obj
sam_model, mask_generator = load_sam_model(
model_name="facebook/sam2.1-hiera-large",
device="cpu",
points_per_batch=16,
)
assert sam_model is mock_pipe_obj.model
assert isinstance(mask_generator, Sam2MaskGenerator)
assert mask_generator.points_per_batch == 16
_, kwargs = mock_pipeline.call_args
assert kwargs["task"] == "mask-generation"
assert kwargs["model"] == "facebook/sam2.1-hiera-large"
assert kwargs["device"] == -1
class TestExtractMaskedRegion:
"""Test suite for extracting masked regions from images."""

8
mini-nav/utils/__init__.py
View File

@@ -4,6 +4,8 @@ from .feature_extractor import (
extract_single_image_feature,
infer_vector_dim,
)
from .image import segment_image, extract_masked_region
from .model import get_dino_dim, load_dino_model, load_hash_compressor, load_sam_model
__all__ = [
"get_device",
@@ -11,4 +13,10 @@ __all__ = [
"infer_vector_dim",
"extract_single_image_feature",
"extract_batch_features",
"segment_image",
"extract_masked_region",
"load_dino_model",
"load_sam_model",
"get_dino_dim",
"load_hash_compressor",
]

3
mini-nav/utils/common.py
View File

@@ -3,11 +3,10 @@ from pathlib import Path
import torch
from configs import cfg_manager
from torch.types import Device
@lru_cache(maxsize=1)
def get_device() -> Device:
def get_device() -> torch.device:
config = cfg_manager.get()
device = config.model.device
if device == "auto":

68
mini-nav/utils/image.py
View File

@@ -0,0 +1,68 @@
from typing import Any
import numpy as np
from PIL import Image
def segment_image(
mask_generator: Any,
image: Image.Image,
min_area: int = 32 * 32,
max_masks: int = 5,
points_per_batch=64,
) -> list[dict[str, Any]]:
"""Segment image using SAM to extract object masks.
Args:
mask_generator: SAM2 mask generator.
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.
Returns:
List of mask dictionaries with keys:
- segment: Binary mask (numpy array)
- area: Mask area in pixels
- bbox: Bounding box [x, y, width, height]
- predicted_iou: Model's confidence in the mask
- stability_score: Stability score for the mask
"""
# Convert PIL Image to numpy array
image_np = np.array(image.convert("RGB"))
# Generate masks
masks = mask_generator(image_np, points_per_batch=points_per_batch)["masks"]
if not masks:
return []
# Filter by minimum area
filtered_masks = [m for m in masks if m["area"] >= min_area]
if not filtered_masks:
return []
# Sort by area (largest first) and limit to max_masks
sorted_masks = sorted(filtered_masks, key=lambda x: x["area"], reverse=True)
return sorted_masks[:max_masks]
def extract_masked_region(
image: Image.Image,
mask: np.ndarray,
) -> Image.Image:
"""Extract masked region from image.
Args:
image: Original PIL Image.
mask: Binary mask as numpy array (True = keep).
Returns:
PIL Image with only the masked region visible.
"""
image_np = np.array(image.convert("RGB"))
# Apply mask
masked_np = image_np * mask[:, :, np.newaxis]
return Image.fromarray(masked_np.astype(np.uint8))

57
mini-nav/utils/model.py Normal file
View File

@@ -0,0 +1,57 @@
"""Model loading utilities for DINO, SAM2 and HashCompressor."""
from compressors import HashCompressor
import torch
from transformers import AutoImageProcessor, AutoModel, pipeline, MaskGenerationPipeline
from .common import get_device
def load_sam_model(
model_name: str = "facebook/sam2.1-hiera-large",
) -> MaskGenerationPipeline:
device = str(get_device())
device_id = 0 if device.startswith("cuda") else -1
return pipeline(
task="mask-generation",
model=model_name,
device=device_id,
)
def load_dino_model(
model_name: str = "facebook/dinov2-large",
) -> tuple[AutoImageProcessor, AutoModel]:
device = get_device()
processor = AutoImageProcessor.from_pretrained(model_name)
dino = AutoModel.from_pretrained(model_name).to(device)
dino.eval()
return processor, dino
def get_dino_dim(model_name: str) -> int:
if "large" in model_name.lower():
return 1024
return 768
def load_hash_compressor(
input_dim: int = 1024,
hash_bits: int = 512,
compressor_path: str | None = None,
) -> HashCompressor:
from compressors.hash_compressor import HashCompressor
device = get_device()
compressor = HashCompressor(input_dim=input_dim, hash_bits=hash_bits).to(device)
if compressor_path is not None:
compressor.load_state_dict(torch.load(compressor_path, map_location=device))
print(f"[OK] Loaded HashCompressor from {compressor_path}")
return compressor

100
mini-nav/utils/sam.py
View File

@@ -1,100 +0,0 @@
"""SAM (Segment Anything Model) utilities for object segmentation."""
from pathlib import Path
from typing import Any
import numpy as np
import torch
from PIL import Image
from sam2.build_sam import build_sam2
from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator
def load_sam_model(
model_name: str = "facebook/sam2.1-hiera-large",
device: str = "cuda",
checkpoint_dir: Path | None = None,
) -> tuple[Any, Any]:
"""Load SAM 2.1 model and mask generator.
Args:
model_name: SAM model name (currently supports facebook/sam2.1-hiera-*).
device: Device to load model on (cuda or cpu).
checkpoint_dir: Optional directory for model checkpoint cache.
Returns:
Tuple of (sam_model, mask_generator).
"""
if device == "cuda" and not torch.cuda.is_available():
device = "cpu"
# Build SAM2 model
sam_model = build_sam2(model_name, device=device)
# Create automatic mask generator
mask_generator = SAM2AutomaticMaskGenerator(sam_model)
return sam_model, mask_generator
def segment_image(
mask_generator: Any,
image: Image.Image,
min_area: int = 32 * 32,
max_masks: int = 5,
) -> list[dict[str, Any]]:
"""Segment image using SAM to extract object masks.
Args:
mask_generator: SAM2AutomaticMaskGenerator instance.
image: PIL Image to segment.
min_area: Minimum mask area threshold in pixels.
max_masks: Maximum number of masks to return.
Returns:
List of mask dictionaries with keys:
- segment: Binary mask (numpy array)
- area: Mask area in pixels
- bbox: Bounding box [x, y, width, height]
- predicted_iou: Model's confidence in the mask
- stability_score: Stability score for the mask
"""
# Convert PIL Image to numpy array
image_np = np.array(image.convert("RGB"))
# Generate masks
masks = mask_generator.generate(image_np)
if not masks:
return []
# Filter by minimum area
filtered_masks = [m for m in masks if m["area"] >= min_area]
if not filtered_masks:
return []
# Sort by area (largest first) and limit to max_masks
sorted_masks = sorted(filtered_masks, key=lambda x: x["area"], reverse=True)
return sorted_masks[:max_masks]
def extract_masked_region(
image: Image.Image,
mask: np.ndarray,
) -> Image.Image:
"""Extract masked region from image.
Args:
image: Original PIL Image.
mask: Binary mask as numpy array (True = keep).
Returns:
PIL Image with only the masked region visible.
"""
image_np = np.array(image.convert("RGB"))
# Apply mask
masked_np = image_np * mask[:, :, np.newaxis]
return Image.fromarray(masked_np.astype(np.uint8))