From 968819e113717994ab9ce8b3f12de10ea53d2d01 Mon Sep 17 00:00:00 2001 From: SikongJueluo Date: Thu, 26 Mar 2026 19:00:13 +0800 Subject: [PATCH] refactor(compressors): consolidate pipeline and improve mask handling --- mini-nav/benchmarks/tasks/retrieval.py | 2 +- mini-nav/compressors/__init__.py | 11 +- mini-nav/compressors/pipeline.py | 155 ++++++++++--------------- mini-nav/configs/config.yaml | 2 +- mini-nav/configs/models.py | 7 +- mini-nav/feature_retrieval.py | 2 +- mini-nav/tests/test_compressors.py | 5 - mini-nav/tests/test_image_utils.py | 82 +++++++++++++ mini-nav/utils/image.py | 96 +++++++++++++-- mini-nav/utils/model.py | 7 +- notebooks/verification.py | 54 ++++++++- 11 files changed, 302 insertions(+), 121 deletions(-) create mode 100644 mini-nav/tests/test_image_utils.py diff --git a/mini-nav/benchmarks/tasks/retrieval.py b/mini-nav/benchmarks/tasks/retrieval.py index c81c48a..435768a 100644 --- a/mini-nav/benchmarks/tasks/retrieval.py +++ b/mini-nav/benchmarks/tasks/retrieval.py @@ -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) ] diff --git a/mini-nav/compressors/__init__.py b/mini-nav/compressors/__init__.py index 5ea1831..fef16a9 100644 --- a/mini-nav/compressors/__init__.py +++ b/mini-nav/compressors/__init__.py @@ -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", diff --git a/mini-nav/compressors/pipeline.py b/mini-nav/compressors/pipeline.py index 6881105..61d1e91 100644 --- a/mini-nav/compressors/pipeline.py +++ b/mini-nav/compressors/pipeline.py @@ -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) diff --git a/mini-nav/configs/config.yaml b/mini-nav/configs/config.yaml index eef0c6a..cfe9741 100644 --- a/mini-nav/configs/config.yaml +++ b/mini-nav/configs/config.yaml @@ -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 diff --git a/mini-nav/configs/models.py b/mini-nav/configs/models.py index f36836c..8df4e04 100644 --- a/mini-nav/configs/models.py +++ b/mini-nav/configs/models.py @@ -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", diff --git a/mini-nav/feature_retrieval.py b/mini-nav/feature_retrieval.py index 1c08262..ff68aa0 100644 --- a/mini-nav/feature_retrieval.py +++ b/mini-nav/feature_retrieval.py @@ -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]), } ] diff --git a/mini-nav/tests/test_compressors.py b/mini-nav/tests/test_compressors.py index 15f06e2..7ee60c3 100644 --- a/mini-nav/tests/test_compressors.py +++ b/mini-nav/tests/test_compressors.py @@ -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.""" diff --git a/mini-nav/tests/test_image_utils.py b/mini-nav/tests/test_image_utils.py new file mode 100644 index 0000000..6f7ec25 --- /dev/null +++ b/mini-nav/tests/test_image_utils.py @@ -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 diff --git a/mini-nav/utils/image.py b/mini-nav/utils/image.py index 9c8e2eb..4bd61c3 100644 --- a/mini-nav/utils/image.py +++ b/mini-nav/utils/image.py @@ -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, diff --git a/mini-nav/utils/model.py b/mini-nav/utils/model.py index d459e7b..711f4ee 100644 --- a/mini-nav/utils/model.py +++ b/mini-nav/utils/model.py @@ -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() diff --git a/notebooks/verification.py b/notebooks/verification.py index 3ac60b2..0085ff8 100644 --- a/notebooks/verification.py +++ b/notebooks/verification.py @@ -9,11 +9,11 @@ import marimo __generated_with = "0.21.1" -app = marimo.App() +app = marimo.App(app_title="Pipeline Verification") @app.cell -def _(): +def import_packages(): import habitat_sim import numpy as np import polars as pl @@ -28,13 +28,16 @@ def _(): return ( HashPipeline, Image, + ObjectNode, RoomNode, SimpleSceneGraph, + extract_masked_region, habitat_sim, maps, np, pl, plt, + segment_image, ) @@ -172,11 +175,21 @@ def _(HashPipeline, hash_bits, sam_max_masks, sam_min_area): sam_max_masks=sam_max_masks, hash_bits=hash_bits, ) - return + return (hash_pipeline,) @app.cell -def _(Image, SimpleSceneGraph, all_room_views, np, room_nodes): +def _( + Image, + ObjectNode, + SimpleSceneGraph, + all_room_views, + extract_masked_region, + hash_pipeline, + np, + room_nodes, + segment_image, +): scene_graph = SimpleSceneGraph( rooms={_room.room_id: _room for _room in room_nodes}, objects={} ) @@ -188,6 +201,39 @@ def _(Image, SimpleSceneGraph, all_room_views, np, room_nodes): _rgb3 = _rgb[..., :3] if _rgb.shape[-1] > 3 else _rgb _image = Image.fromarray(_rgb3.astype(np.uint8)) + _masks = segment_image( + hash_pipeline.mask_generator, + _image, + min_area=hash_pipeline.sam_min_mask_area, + max_masks=hash_pipeline.sam_max_masks, + points_per_batch=hash_pipeline.sam_points_per_batch, + ) + total_masks += len(_masks) + + for _mask in _masks: + _masked_image = extract_masked_region(_image, _mask["segment"]) + _bits = hash_pipeline(_masked_image) + + _bbox = _mask["bbox"] + _obj_center = np.array( + [_bbox[0] + _bbox[2] / 2, _bbox[1] + _bbox[3] / 2, 0.0], + dtype=np.float32, + ) + + _obj_id = f"obj_{_obj_index:04d}" + _obj_index += 1 + _bits_np = _bits.squeeze().detach().cpu().numpy() + + _obj_node = ObjectNode( + obj_id=_obj_id, + room_id=_room_id, + position=_obj_center, + visual_hash=_bits_np, + semantic_hash=_bits_np, + hit_count=1, + last_seen_frame=0, + ) + scene_graph.objects[_obj_id] = _obj_node print(f"Total objects created: {len(scene_graph.objects)}") print(f"Total processed masks: {total_masks}")