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

2
mini-nav/benchmarks/tasks/retrieval.py
View File

@@ -60,7 +60,7 @@ def _establish_eval_database(
{
"id": global_idx + j,
"label": labels_list[j],
"vector": all_features[global_idx + j].numpy(),
"vector": all_features[global_idx + j].detach().cpu().numpy(),
}
for j in range(batch_size)
]

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)

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

@@ -1,7 +1,7 @@
model:
dino_model: "facebook/dinov2-large"
compression_dim: 512
device: "auto" # auto-detect GPU
device: "cuda:3" # 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

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

@@ -18,7 +18,12 @@ class ModelConfig(BaseModel):
compression_dim: int = Field(
default=512, gt=0, description="Output feature dimension"
)
device: str = "auto"
device: str = Field(
default="auto",
description=(
"Device to use for model inference (e.g., 'cuda:1,3', 'auto', 'cpu')"
),
)
sam_model: str = Field(
default="facebook/sam2.1-hiera-large",
description="SAM model name from HuggingFace",

2
mini-nav/feature_retrieval.py
View File

@@ -102,7 +102,7 @@ class FeatureRetrieval:
{
"id": i,
"label": batch_label,
"vector": cls_tokens[i].numpy(),
"vector": cls_tokens[i].detach().cpu().numpy(),
"binary": pil_image_to_bytes(images[i]),
}
]

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

@@ -7,7 +7,6 @@ from compressors import (
BinarySign,
HashCompressor,
HashPipeline,
SAMHashPipeline,
VideoPositiveMask,
bits_to_hash,
create_pipeline_from_config,
@@ -257,10 +256,6 @@ class TestHashPipeline:
pipeline = HashPipeline(hash_bits=256)
assert pipeline.hash_bits == 256
def test_pipeline_alias(self):
"""Verify SAMHashPipeline is alias for HashPipeline."""
assert SAMHashPipeline is HashPipeline
class TestConfigIntegration:
"""Test suite for config integration with pipeline."""

82
mini-nav/tests/test_image_utils.py Normal file
View File

@@ -0,0 +1,82 @@
from unittest.mock import Mock
import torch
from PIL import Image
from utils.image import segment_image
def test_segment_image_passes_pil_image_to_mask_generator() -> None:
mock_generator = Mock(return_value={"masks": []})
segment_image(
mock_generator,
Image.new("RGBA", (16, 16), color=(255, 0, 0, 255)),
points_per_batch=32,
)
image_arg = mock_generator.call_args.args[0]
assert isinstance(image_arg, Image.Image)
assert image_arg.mode == "RGB"
assert mock_generator.call_args.kwargs["points_per_batch"] == 32
def test_segment_image_supports_tensor_masks_output() -> None:
masks_tensor = torch.tensor(
[
[
[1, 1, 0],
[1, 1, 0],
[0, 0, 0],
],
[
[1, 1, 1],
[1, 1, 1],
[1, 1, 1],
],
],
dtype=torch.float32,
)
mock_generator = Mock(return_value={"masks": masks_tensor})
result = segment_image(
mock_generator,
Image.new("RGB", (3, 3), color=(0, 0, 0)),
min_area=3,
max_masks=5,
)
assert len(result) == 2
assert result[0]["area"] == 9
assert result[0]["bbox"] == [0, 0, 3, 3]
assert result[1]["area"] == 4
assert result[1]["bbox"] == [0, 0, 2, 2]
def test_segment_image_filters_tensor_masks_by_min_area() -> None:
masks_tensor = torch.tensor(
[
[
[1, 0, 0],
[0, 0, 0],
[0, 0, 0],
],
[
[1, 1, 0],
[1, 1, 0],
[0, 0, 0],
],
],
dtype=torch.float32,
)
mock_generator = Mock(return_value={"masks": masks_tensor})
result = segment_image(
mock_generator,
Image.new("RGB", (3, 3), color=(0, 0, 0)),
min_area=2,
max_masks=5,
)
assert len(result) == 1
assert result[0]["area"] == 4

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

@@ -9,7 +9,7 @@ def segment_image(
image: Image.Image,
min_area: int = 32 * 32,
max_masks: int = 5,
points_per_batch=64,
points_per_batch: int = 64,
) -> list[dict[str, Any]]:
"""Segment image using SAM to extract object masks.
@@ -27,26 +27,104 @@ def segment_image(
- 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"))
image_rgb = image.convert("RGB")
raw_output = mask_generator(image_rgb, points_per_batch=points_per_batch)
raw_masks = raw_output.get("masks", raw_output)
# Generate masks
masks = mask_generator(image_np, points_per_batch=points_per_batch)["masks"]
normalized_masks: list[dict[str, Any]] = []
if not masks:
if isinstance(raw_masks, list):
if raw_masks and isinstance(raw_masks[0], dict):
normalized_masks = raw_masks
else:
for mask_like in raw_masks:
mask_dict = _to_mask_dict(mask_like)
if mask_dict is not None:
normalized_masks.append(mask_dict)
else:
mask_array = _to_numpy_mask_array(raw_masks)
if mask_array is not None:
if mask_array.ndim == 2:
mask_array = np.expand_dims(mask_array, axis=0)
if mask_array.ndim == 3:
for single_mask in mask_array:
mask_dict = _to_mask_dict(single_mask)
if mask_dict is not None:
normalized_masks.append(mask_dict)
if not normalized_masks:
return []
# Filter by minimum area
filtered_masks = [m for m in masks if m["area"] >= min_area]
filtered_masks = [m for m in normalized_masks if int(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 _to_numpy_mask_array(mask_like: Any) -> np.ndarray | None:
if mask_like is None:
return None
if isinstance(mask_like, np.ndarray):
return mask_like
try:
import torch
if isinstance(mask_like, torch.Tensor):
return mask_like.detach().cpu().numpy()
except ImportError:
pass
return None
def _to_mask_dict(mask_like: Any) -> dict[str, Any] | None:
if isinstance(mask_like, dict):
if "area" in mask_like and "bbox" in mask_like and "segment" in mask_like:
return mask_like
segment = mask_like.get("segment")
if segment is None and "mask" in mask_like:
segment = mask_like["mask"]
if segment is None:
return None
mask_array = _to_numpy_mask_array(segment)
if mask_array is None:
return None
return _build_mask_dict(mask_array)
mask_array = _to_numpy_mask_array(mask_like)
if mask_array is None:
return None
return _build_mask_dict(mask_array)
def _build_mask_dict(mask_array: np.ndarray) -> dict[str, Any] | None:
if mask_array.ndim != 2:
return None
segment = mask_array.astype(bool)
area = int(segment.sum())
if area <= 0:
return None
ys, xs = np.where(segment)
min_y, max_y = int(ys.min()), int(ys.max())
min_x, max_x = int(xs.min()), int(xs.max())
bbox = [min_x, min_y, max_x - min_x + 1, max_y - min_y + 1]
return {
"segment": segment,
"area": area,
"bbox": bbox,
"predicted_iou": None,
"stability_score": None,
}
def extract_masked_region(
image: Image.Image,
mask: np.ndarray,

7
mini-nav/utils/model.py
View File

@@ -1,6 +1,6 @@
"""Model loading utilities for DINO, SAM2 and HashCompressor."""
from compressors import HashCompressor
from typing import TYPE_CHECKING
import torch
@@ -8,6 +8,9 @@ from transformers import AutoImageProcessor, AutoModel, pipeline, MaskGeneration
from .common import get_device
if TYPE_CHECKING:
from compressors.hash_compressor import HashCompressor
def load_sam_model(
model_name: str = "facebook/sam2.1-hiera-large",
@@ -44,7 +47,7 @@ def load_hash_compressor(
input_dim: int = 1024,
hash_bits: int = 512,
compressor_path: str | None = None,
) -> HashCompressor:
) -> "HashCompressor":
from compressors.hash_compressor import HashCompressor
device = get_device()