refactor(pipeline): integrate SAM segmentation and modularize model loading

This commit is contained in:
2026-03-24 21:52:02 +08:00
parent 9e6339e580
commit 90d5a8f08a
11 changed files with 437 additions and 172 deletions

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

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:

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

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."""

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."""

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",
]

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":

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

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

View File

@@ -1,3 +1,11 @@
# /// script
# requires-python = ">=3.13"
# dependencies = [
# "marimo>=0.21.1",
# "pyzmq>=27.1.0",
# ]
# ///
import marimo
__generated_with = "0.21.1"
@@ -8,11 +16,26 @@ app = marimo.App()
def _():
import habitat_sim
import numpy as np
import polars as pl
from habitat.utils.visualizations import maps
from matplotlib import pyplot as plt
from scenegraph import RoomNode
from PIL import Image
from compressors.pipeline import HashPipeline
from scenegraph import ObjectNode, RoomNode, SimpleSceneGraph
from utils.common import get_device
from utils.image import extract_masked_region, segment_image
return RoomNode, habitat_sim, maps, np, plt
return (
HashPipeline,
Image,
RoomNode,
SimpleSceneGraph,
habitat_sim,
maps,
np,
pl,
plt,
)
@app.cell
@@ -51,7 +74,20 @@ def _(habitat_sim):
cfg = habitat_sim.Configuration(sim_cfg, [agent_cfg])
sim = habitat_sim.Simulator(cfg)
agent = sim.initialize_agent(0)
return agent, meters_per_pixel, num_rooms, sim, views_per_room
sam_max_masks = 5
sam_min_area = 32 * 32
sam_points_per_batch = 64
hash_bits = 512
return (
agent,
hash_bits,
meters_per_pixel,
num_rooms,
sam_max_masks,
sam_min_area,
sim,
views_per_room,
)
@app.cell
@@ -66,7 +102,6 @@ def _(RoomNode, np, num_rooms, sim):
)
room_nodes.append(_room_node)
room_points = np.vstack([_node.center for _node in room_nodes])
print("Sampled room centers:")
for _node in room_nodes:
print(_node.room_id, _node.center)
@@ -125,6 +160,83 @@ def _(agent, habitat_sim, plt, room_nodes, sim, views_per_room):
_ax.axis("off")
plt.tight_layout()
plt.show()
return (all_room_views,)
@app.cell
def _(HashPipeline, hash_bits, sam_max_masks, sam_min_area):
hash_pipeline = HashPipeline(
dino_model="facebook/dinov2-large",
sam_model="facebook/sam2.1-hiera-large",
sam_min_mask_area=sam_min_area,
sam_max_masks=sam_max_masks,
hash_bits=hash_bits,
)
return
@app.cell
def _(Image, SimpleSceneGraph, all_room_views, np, room_nodes):
scene_graph = SimpleSceneGraph(
rooms={_room.room_id: _room for _room in room_nodes}, objects={}
)
total_masks = 0
_obj_index = 0
for _room_id, _views in all_room_views.items():
for _view_idx, _rgb in enumerate(_views):
_rgb3 = _rgb[..., :3] if _rgb.shape[-1] > 3 else _rgb
_image = Image.fromarray(_rgb3.astype(np.uint8))
print(f"Total objects created: {len(scene_graph.objects)}")
print(f"Total processed masks: {total_masks}")
return (scene_graph,)
@app.cell
def _(pl, scene_graph):
_room_rows = []
for _room in scene_graph.rooms.values():
_room_rows.append(
{
"room_id": _room.room_id,
"center_x": float(_room.center[0]),
"center_y": float(_room.center[1]),
"center_z": float(_room.center[2]),
"bbox_dx": float(_room.bbox_extent[0]),
"bbox_dy": float(_room.bbox_extent[1]),
"bbox_dz": float(_room.bbox_extent[2]),
}
)
_object_rows = []
for _obj in scene_graph.objects.values():
_object_rows.append(
{
"obj_id": _obj.obj_id,
"room_id": _obj.room_id,
"last_seen_frame": int(_obj.last_seen_frame),
"hit_count": int(_obj.hit_count),
"visual_hash": _obj.visual_hash.tolist(),
"semantic_hash": _obj.semantic_hash.tolist(),
}
)
rooms_table = pl.DataFrame(_room_rows)
objects_table = pl.DataFrame(_object_rows)
return objects_table, rooms_table
@app.cell
def _(rooms_table):
rooms_table
return
@app.cell
def _(objects_table):
objects_table
return