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

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
+from utils import get_device
-to binary hash codes using HashCompressor.
+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:
+        self.device = get_device()
-            device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.device = torch.device(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.mask_generator = load_sam_model(model_name=sam_model)
-        self.processor = AutoImageProcessor.from_pretrained(dino_model)
+
-        self.dino = AutoModel.from_pretrained(dino_model).to(self.device)
+        self.processor, self.dino = load_dino_model(model_name=dino_model)
-        self.dino.eval()
 
         # 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.hash_compressor = load_hash_compressor(
-        self._init_hash_compressor(hash_bits, compressor_path)
+            input_dim=self.dino_dim,
-
+            hash_bits=hash_bits,
-    def _init_hash_compressor(
+            compressor_path=compressor_path,
-        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
         )
 
-        # 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
+        return self._encode_image(image, apply_sam=True)
-        inputs = self.processor(image, return_tensors="pt").to(self.device)
 
-        with torch.no_grad():
+    def encode_masked_region(self, image: Image.Image) -> torch.Tensor:
-            outputs = self.dino(**inputs)
+        """Encode a pre-masked region using DINO+Hash without SAM stage."""
-            tokens = outputs.last_hidden_state  # [1, N, dim]
+        return self._encode_image(image, apply_sam=False)
-
-        # Compress to hash codes
-        _, _, bits = self.hash_compressor(tokens)
-
-        return bits
 
     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

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:

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

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

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]}
+        mask1 = {
-        mask2 = {"segment": np.ones((10, 10), dtype=bool), "area": 100, "bbox": [0, 0, 10, 10]}
+            "segment": np.ones((5, 5), dtype=bool),
-        mask3 = {"segment": np.ones((3, 3), dtype=bool), "area": 9, "bbox": [0, 0, 3, 3]}
+            "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."""
 

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

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

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

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

mini-nav/utils/sam.py

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

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