From 90d5a8f08aa7a67f61f91dc39080b21a4a4ed2cb Mon Sep 17 00:00:00 2001 From: SikongJueluo Date: Tue, 24 Mar 2026 21:52:02 +0800 Subject: [PATCH] refactor(pipeline): integrate SAM segmentation and modularize model loading --- mini-nav/compressors/pipeline.py | 133 ++++++++++++++++------------- mini-nav/configs/config.yaml | 3 +- mini-nav/configs/models.py | 12 ++- mini-nav/tests/test_compressors.py | 62 +++++++++++++- mini-nav/tests/test_sam.py | 43 +++++++++- mini-nav/utils/__init__.py | 8 ++ mini-nav/utils/common.py | 3 +- mini-nav/utils/image.py | 68 +++++++++++++++ mini-nav/utils/model.py | 57 +++++++++++++ mini-nav/utils/sam.py | 100 ---------------------- notebooks/verification.py | 120 +++++++++++++++++++++++++- 11 files changed, 437 insertions(+), 172 deletions(-) create mode 100644 mini-nav/utils/model.py delete mode 100644 mini-nav/utils/sam.py diff --git a/mini-nav/compressors/pipeline.py b/mini-nav/compressors/pipeline.py index 0451a68..6881105 100644 --- a/mini-nav/compressors/pipeline.py +++ b/mini-nav/compressors/pipeline.py @@ -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 diff --git a/mini-nav/configs/config.yaml b/mini-nav/configs/config.yaml index f3535c3..eef0c6a 100644 --- a/mini-nav/configs/config.yaml +++ b/mini-nav/configs/config.yaml @@ -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: diff --git a/mini-nav/configs/models.py b/mini-nav/configs/models.py index 7109178..f36836c 100644 --- a/mini-nav/configs/models.py +++ b/mini-nav/configs/models.py @@ -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" ) diff --git a/mini-nav/tests/test_compressors.py b/mini-nav/tests/test_compressors.py index 99677f1..15f06e2 100644 --- a/mini-nav/tests/test_compressors.py +++ b/mini-nav/tests/test_compressors.py @@ -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.""" diff --git a/mini-nav/tests/test_sam.py b/mini-nav/tests/test_sam.py index 3d9ba2c..918c0a4 100644 --- a/mini-nav/tests/test_sam.py +++ b/mini-nav/tests/test_sam.py @@ -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.""" diff --git a/mini-nav/utils/__init__.py b/mini-nav/utils/__init__.py index 1591fce..835da16 100644 --- a/mini-nav/utils/__init__.py +++ b/mini-nav/utils/__init__.py @@ -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", ] diff --git a/mini-nav/utils/common.py b/mini-nav/utils/common.py index 0041a50..1a08583 100644 --- a/mini-nav/utils/common.py +++ b/mini-nav/utils/common.py @@ -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": diff --git a/mini-nav/utils/image.py b/mini-nav/utils/image.py index e69de29..9c8e2eb 100644 --- a/mini-nav/utils/image.py +++ b/mini-nav/utils/image.py @@ -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)) diff --git a/mini-nav/utils/model.py b/mini-nav/utils/model.py new file mode 100644 index 0000000..d459e7b --- /dev/null +++ b/mini-nav/utils/model.py @@ -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 diff --git a/mini-nav/utils/sam.py b/mini-nav/utils/sam.py deleted file mode 100644 index a896d90..0000000 --- a/mini-nav/utils/sam.py +++ /dev/null @@ -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)) diff --git a/notebooks/verification.py b/notebooks/verification.py index 669ccc3..3ac60b2 100644 --- a/notebooks/verification.py +++ b/notebooks/verification.py @@ -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