refactor(simulator): extract habitat simulator and visualization modules

mini-nav/scenegraph/objectnode.py

@@ -13,12 +13,17 @@ class ObjectNode:
     position: np.ndarray  # [x, y, z] 世界坐标系下的中心点或锚点
 
     # 特征：直接存你压缩后的 512bit 结果，不搞历史缓存
-    visual_hash: np.ndarray  # 512bit 视觉特征 (用于外观检索)
+    visual_hash: bytes  # 512bit 视觉特征 (用于外观检索)
-    semantic_hash: np.ndarray  # 512bit 语义特征 (用于类别对齐)
+    semantic_hash: bytes  # 512bit 语义特征 (用于类别对齐)
-
-    # 原始图片：mask处理并裁切后的图片数据（二进制格式）
-    image_bytes: bytes | None = None  # numpy array.tobytes() 原始像素数据
 
     # Debug 必备：极简的生命周期管理，防止满屏"幽灵节点"
     hit_count: int = 1  # 被观测到的次数。太低的可以直接过滤掉
     last_seen_frame: int = 0  # 最后一次看到的帧号或时间戳
+
+    def __post_init__(self):
+        if len(self.visual_hash) != 64:
+            raise ValueError("visual_hash must be exactly 64 bytes (512 bits)")
+        if len(self.semantic_hash) != 64:
+            raise ValueError("semantic_hash must be exactly 64 bytes (512 bits)")
+        if self.position.shape != (3,):
+            raise ValueError("position must have shape (3,)")

mini-nav/simulator/__init__.py

@@ -0,0 +1,18 @@
+from .habitat import (
+    HabitatSimulatorConfig,
+    close_habitat_simulator,
+    create_habitat_simulator,
+)
+from .topdown import TopDownRenderStyle, TopDownSceneElements, render_topdown_scene_map
+from .views import RoomViewsByRoom, collect_room_views_by_room
+
+__all__ = [
+    "HabitatSimulatorConfig",
+    "TopDownRenderStyle",
+    "TopDownSceneElements",
+    "RoomViewsByRoom",
+    "close_habitat_simulator",
+    "collect_room_views_by_room",
+    "create_habitat_simulator",
+    "render_topdown_scene_map",
+]

mini-nav/simulator/habitat.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from importlib import import_module
+from typing import Any
+
+
+@dataclass(frozen=True)
+class HabitatSimulatorConfig:
+    scene_path: str
+    views_per_room: int = 6
+    image_size: int = 256
+    sensor_height: float = 1.5
+    move_forward_step: float = 0.25
+    enable_physics: bool = False
+    sensor_uuid: str = "color_sensor"
+    agent_id: int = 0
+
+
+def create_habitat_simulator(
+    config: HabitatSimulatorConfig,
+    habitat_sim_module: Any | None = None,
+) -> tuple[Any, Any]:
+    if config.views_per_room <= 0:
+        raise ValueError("views_per_room must be greater than 0")
+
+    if config.image_size <= 0:
+        raise ValueError("image_size must be greater than 0")
+
+    if config.move_forward_step <= 0:
+        raise ValueError("move_forward_step must be greater than 0")
+
+    if habitat_sim_module is None:
+        habitat_sim_module = import_module("habitat_sim")
+
+    sim_cfg = habitat_sim_module.SimulatorConfiguration()
+    sim_cfg.scene_id = config.scene_path
+    sim_cfg.enable_physics = config.enable_physics
+
+    agent_cfg = habitat_sim_module.agent.AgentConfiguration()
+    rgb_sensor_spec = habitat_sim_module.CameraSensorSpec()
+    rgb_sensor_spec.uuid = config.sensor_uuid
+    rgb_sensor_spec.sensor_type = habitat_sim_module.SensorType.COLOR
+    rgb_sensor_spec.resolution = [config.image_size, config.image_size]
+    rgb_sensor_spec.position = [0.0, config.sensor_height, 0.0]
+    agent_cfg.sensor_specifications = [rgb_sensor_spec]
+
+    turn_angle = 360.0 / config.views_per_room
+    agent_cfg.action_space = {
+        "move_forward": habitat_sim_module.agent.ActionSpec(
+            "move_forward",
+            habitat_sim_module.agent.ActuationSpec(amount=config.move_forward_step),
+        ),
+        "turn_left": habitat_sim_module.agent.ActionSpec(
+            "turn_left",
+            habitat_sim_module.agent.ActuationSpec(amount=turn_angle),
+        ),
+        "turn_right": habitat_sim_module.agent.ActionSpec(
+            "turn_right",
+            habitat_sim_module.agent.ActuationSpec(amount=turn_angle),
+        ),
+    }
+
+    simulator_cfg = habitat_sim_module.Configuration(sim_cfg, [agent_cfg])
+    simulator = habitat_sim_module.Simulator(simulator_cfg)
+    agent = simulator.initialize_agent(config.agent_id)
+    return simulator, agent
+
+
+def close_habitat_simulator(simulator: Any) -> None:
+    close = getattr(simulator, "close", None)
+    if callable(close):
+        close()

mini-nav/simulator/topdown.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from importlib import import_module
+from typing import Any, Sequence
+
+
+@dataclass(frozen=True)
+class TopDownSceneElements:
+    room_nodes: Sequence[Any]
+    object_nodes: Sequence[Any] = field(default_factory=tuple)
+    edges: Sequence[tuple[str, str]] = field(default_factory=tuple)
+
+
+@dataclass(frozen=True)
+class TopDownRenderStyle:
+    room_color: str = "red"
+    room_label_color: str = "yellow"
+    room_marker_size: int = 50
+    room_label_offset: int = 2
+    figure_size: tuple[int, int] = (8, 8)
+    map_cmap: str = "gray"
+    title: str = "RoomNode Top-Down Map"
+
+
+def render_topdown_scene_map(
+    pathfinder: Any,
+    elements: TopDownSceneElements,
+    meters_per_pixel: float,
+    style: TopDownRenderStyle | None = None,
+    maps_module: Any | None = None,
+    plt_module: Any | None = None,
+) -> Any:
+    if not elements.room_nodes:
+        raise ValueError("room_nodes must not be empty")
+
+    if meters_per_pixel <= 0:
+        raise ValueError("meters_per_pixel must be greater than 0")
+
+    if elements.object_nodes:
+        raise NotImplementedError("object_nodes overlay is not implemented yet")
+
+    if elements.edges:
+        raise NotImplementedError("edge overlay is not implemented yet")
+
+    if style is None:
+        style = TopDownRenderStyle()
+
+    if maps_module is None:
+        maps_module = import_module("habitat.utils.visualizations.maps")
+
+    if plt_module is None:
+        plt_module = import_module("matplotlib.pyplot")
+
+    map_height = float(elements.room_nodes[0].center[1])
+    top_down_map = maps_module.get_topdown_map(
+        pathfinder,
+        height=map_height,
+        meters_per_pixel=meters_per_pixel,
+    )
+
+    plt_module.figure(figsize=style.figure_size)
+    plt_module.imshow(top_down_map, cmap=style.map_cmap)
+
+    for room_node in elements.room_nodes:
+        grid_y, grid_x = maps_module.to_grid(
+            float(room_node.center[2]),
+            float(room_node.center[0]),
+            top_down_map.shape,
+            pathfinder=pathfinder,
+        )
+        plt_module.scatter(grid_x, grid_y, c=style.room_color, s=style.room_marker_size)
+        plt_module.text(
+            grid_x + style.room_label_offset,
+            grid_y + style.room_label_offset,
+            room_node.room_id,
+            color=style.room_label_color,
+            fontsize=8,
+        )
+
+    plt_module.title(style.title)
+    plt_module.axis("off")
+    plt_module.show()
+    return top_down_map

mini-nav/simulator/views.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from importlib import import_module
+from typing import Any, Callable, Iterable, Sequence
+
+from rich.progress import track
+
+RoomViewsByRoom = dict[str, list[Any]]
+ProgressTrack = Callable[[Iterable[Any], str], Iterable[Any]]
+
+
+def collect_room_views_by_room(
+    agent: Any,
+    sim: Any,
+    room_nodes: Sequence[Any],
+    views_per_room: int,
+    *,
+    habitat_sim_module: Any | None = None,
+    sensor_uuid: str = "color_sensor",
+    turn_action: str = "turn_left",
+    progress_description: str = "Collecting room views",
+    progress_track: ProgressTrack = track,
+) -> RoomViewsByRoom:
+    if views_per_room <= 0:
+        raise ValueError("views_per_room must be greater than 0")
+
+    if habitat_sim_module is None:
+        habitat_sim_module = import_module("habitat_sim")
+
+    all_room_views: RoomViewsByRoom = {}
+    for room_node in progress_track(room_nodes, progress_description):
+        agent_state = habitat_sim_module.AgentState()
+        agent_state.position = room_node.center.copy()
+        agent.set_state(agent_state)
+
+        room_views = []
+        for _ in range(views_per_room):
+            observations = sim.get_sensor_observations()
+            room_views.append(observations[sensor_uuid])
+            sim.step(turn_action)
+
+        all_room_views[room_node.room_id] = room_views
+
+    return all_room_views

mini-nav/tests/test_habitat_simulator.py

@@ -0,0 +1,124 @@
+from types import SimpleNamespace
+
+import pytest
+
+from simulator import HabitatSimulatorConfig, create_habitat_simulator
+
+
+class _FakeSimulatorConfiguration:
+    def __init__(self):
+        self.scene_id = ""
+        self.enable_physics = True
+
+
+class _FakeAgentConfiguration:
+    def __init__(self):
+        self.sensor_specifications = []
+        self.action_space = {}
+
+
+class _FakeCameraSensorSpec:
+    def __init__(self):
+        self.uuid = ""
+        self.sensor_type = None
+        self.resolution = []
+        self.position = []
+
+
+class _FakeActuationSpec:
+    def __init__(self, amount):
+        self.amount = amount
+
+
+class _FakeActionSpec:
+    def __init__(self, name, actuation):
+        self.name = name
+        self.actuation = actuation
+
+
+class _FakeConfiguration:
+    def __init__(self, sim_cfg, agent_cfgs):
+        self.sim_cfg = sim_cfg
+        self.agent_cfgs = agent_cfgs
+
+
+class _FakeSimulator:
+    def __init__(self, cfg):
+        self.cfg = cfg
+        self.initialized_agent_id = None
+
+    def initialize_agent(self, agent_id):
+        self.initialized_agent_id = agent_id
+        return {"agent_id": agent_id}
+
+
+def _create_fake_habitat_module():
+    return SimpleNamespace(
+        SimulatorConfiguration=_FakeSimulatorConfiguration,
+        CameraSensorSpec=_FakeCameraSensorSpec,
+        SensorType=SimpleNamespace(COLOR="color"),
+        Configuration=_FakeConfiguration,
+        Simulator=_FakeSimulator,
+        agent=SimpleNamespace(
+            AgentConfiguration=_FakeAgentConfiguration,
+            ActionSpec=_FakeActionSpec,
+            ActuationSpec=_FakeActuationSpec,
+        ),
+    )
+
+
+def test_create_habitat_simulator_builds_expected_configuration():
+    fake_habitat = _create_fake_habitat_module()
+    config = HabitatSimulatorConfig(
+        scene_path="scene.glb",
+        views_per_room=8,
+        image_size=128,
+        sensor_height=1.25,
+        move_forward_step=0.5,
+        enable_physics=False,
+        sensor_uuid="rgb",
+        agent_id=2,
+    )
+
+    simulator, agent = create_habitat_simulator(config, habitat_sim_module=fake_habitat)
+
+    assert simulator.cfg.sim_cfg.scene_id == "scene.glb"
+    assert simulator.cfg.sim_cfg.enable_physics is False
+
+    created_agent_cfg = simulator.cfg.agent_cfgs[0]
+    sensor = created_agent_cfg.sensor_specifications[0]
+    assert sensor.uuid == "rgb"
+    assert sensor.sensor_type == "color"
+    assert sensor.resolution == [128, 128]
+    assert sensor.position == [0.0, 1.25, 0.0]
+
+    assert created_agent_cfg.action_space["move_forward"].actuation.amount == 0.5
+    assert created_agent_cfg.action_space["turn_left"].actuation.amount == 45.0
+    assert created_agent_cfg.action_space["turn_right"].actuation.amount == 45.0
+
+    assert simulator.initialized_agent_id == 2
+    assert agent == {"agent_id": 2}
+
+
+def test_create_habitat_simulator_validates_views_per_room():
+    fake_habitat = _create_fake_habitat_module()
+    config = HabitatSimulatorConfig(scene_path="scene.glb", views_per_room=0)
+
+    with pytest.raises(ValueError, match="views_per_room"):
+        create_habitat_simulator(config, habitat_sim_module=fake_habitat)
+
+
+def test_create_habitat_simulator_validates_image_size():
+    fake_habitat = _create_fake_habitat_module()
+    config = HabitatSimulatorConfig(scene_path="scene.glb", image_size=0)
+
+    with pytest.raises(ValueError, match="image_size"):
+        create_habitat_simulator(config, habitat_sim_module=fake_habitat)
+
+
+def test_create_habitat_simulator_validates_move_forward_step():
+    fake_habitat = _create_fake_habitat_module()
+    config = HabitatSimulatorConfig(scene_path="scene.glb", move_forward_step=0)
+
+    with pytest.raises(ValueError, match="move_forward_step"):
+        create_habitat_simulator(config, habitat_sim_module=fake_habitat)

mini-nav/tests/test_room_views.py

@@ -0,0 +1,108 @@
+from types import SimpleNamespace
+
+import pytest
+
+from simulator import collect_room_views_by_room
+
+
+class _FakeAgent:
+    def __init__(self):
+        self.positions = []
+
+    def set_state(self, state):
+        self.positions.append(state.position)
+
+
+class _FakeSimulator:
+    def __init__(self):
+        self._frame_index = 0
+        self.actions = []
+
+    def get_sensor_observations(self):
+        observations = {
+            "color_sensor": f"frame_{self._frame_index}",
+            "depth_sensor": f"depth_{self._frame_index}",
+        }
+        self._frame_index += 1
+        return observations
+
+    def step(self, action_name):
+        self.actions.append(action_name)
+
+
+class _FakeAgentState:
+    def __init__(self):
+        self.position = None
+
+
+def test_collect_room_views_by_room_collects_grouped_frames_with_single_outer_progress():
+    track_calls = []
+
+    def fake_track(iterable, description):
+        track_calls.append(description)
+        return iterable
+
+    agent = _FakeAgent()
+    sim = _FakeSimulator()
+    room_nodes = [
+        SimpleNamespace(room_id="room_00", center=[1.0, 2.0, 3.0]),
+        SimpleNamespace(room_id="room_01", center=[4.0, 5.0, 6.0]),
+    ]
+    fake_habitat = SimpleNamespace(AgentState=_FakeAgentState)
+
+    room_views = collect_room_views_by_room(
+        agent=agent,
+        sim=sim,
+        room_nodes=room_nodes,
+        views_per_room=3,
+        habitat_sim_module=fake_habitat,
+        progress_track=fake_track,
+    )
+
+    assert track_calls == ["Collecting room views"]
+    assert room_views == {
+        "room_00": ["frame_0", "frame_1", "frame_2"],
+        "room_01": ["frame_3", "frame_4", "frame_5"],
+    }
+    assert sim.actions == [
+        "turn_left",
+        "turn_left",
+        "turn_left",
+        "turn_left",
+        "turn_left",
+        "turn_left",
+    ]
+    assert agent.positions == [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]
+
+
+def test_collect_room_views_by_room_uses_custom_sensor_and_turn_action():
+    agent = _FakeAgent()
+    sim = _FakeSimulator()
+    room_nodes = [SimpleNamespace(room_id="room_00", center=[0.0, 1.0, 0.0])]
+    fake_habitat = SimpleNamespace(AgentState=_FakeAgentState)
+
+    room_views = collect_room_views_by_room(
+        agent=agent,
+        sim=sim,
+        room_nodes=room_nodes,
+        views_per_room=2,
+        habitat_sim_module=fake_habitat,
+        sensor_uuid="depth_sensor",
+        turn_action="turn_right",
+        progress_track=lambda iterable, description: iterable,
+    )
+
+    assert room_views == {"room_00": ["depth_0", "depth_1"]}
+    assert sim.actions == ["turn_right", "turn_right"]
+
+
+def test_collect_room_views_by_room_validates_views_per_room():
+    with pytest.raises(ValueError, match="views_per_room"):
+        collect_room_views_by_room(
+            agent=_FakeAgent(),
+            sim=_FakeSimulator(),
+            room_nodes=[SimpleNamespace(room_id="room_00", center=[0.0, 1.0, 0.0])],
+            views_per_room=0,
+            habitat_sim_module=SimpleNamespace(AgentState=_FakeAgentState),
+            progress_track=lambda iterable, description: iterable,
+        )

mini-nav/tests/test_topdown_map.py

@@ -0,0 +1,126 @@
+from types import SimpleNamespace
+
+import pytest
+
+from simulator import (
+    TopDownRenderStyle,
+    TopDownSceneElements,
+    render_topdown_scene_map,
+)
+
+
+class _FakeMaps:
+    def __init__(self):
+        self.to_grid_calls: list[tuple[float, float]] = []
+
+    def get_topdown_map(self, pathfinder, height, meters_per_pixel):
+        return [[0, 0], [0, 0]]
+
+    def to_grid(self, z, x, shape, pathfinder):
+        self.to_grid_calls.append((z, x))
+        return (int(z), int(x))
+
+
+class _FakePlt:
+    def __init__(self):
+        self.scatter_calls: list[tuple[int, int]] = []
+        self.text_calls: list[str] = []
+        self.shown = False
+
+    def figure(self, figsize):
+        return None
+
+    def imshow(self, image, cmap):
+        return None
+
+    def scatter(self, x, y, c, s):
+        self.scatter_calls.append((x, y))
+
+    def text(self, x, y, text, color, fontsize):
+        self.text_calls.append(text)
+
+    def title(self, title):
+        return None
+
+    def axis(self, mode):
+        return None
+
+    def show(self):
+        self.shown = True
+
+
+def test_render_topdown_scene_map_renders_room_nodes_only():
+    fake_maps = _FakeMaps()
+    fake_plt = _FakePlt()
+    room_nodes = [
+        SimpleNamespace(room_id="room_00", center=[1.0, 2.0, 3.0]),
+        SimpleNamespace(room_id="room_01", center=[4.0, 2.0, 5.0]),
+    ]
+    elements = TopDownSceneElements(room_nodes=room_nodes)
+
+    top_down_map = render_topdown_scene_map(
+        pathfinder=SimpleNamespace(),
+        elements=elements,
+        meters_per_pixel=0.05,
+        style=TopDownRenderStyle(),
+        maps_module=fake_maps,
+        plt_module=fake_plt,
+    )
+
+    assert top_down_map == [[0, 0], [0, 0]]
+    assert fake_maps.to_grid_calls == [(3.0, 1.0), (5.0, 4.0)]
+    assert fake_plt.scatter_calls == [(1, 3), (4, 5)]
+    assert fake_plt.text_calls == ["room_00", "room_01"]
+    assert fake_plt.shown is True
+
+
+def test_render_topdown_scene_map_validates_room_nodes():
+    with pytest.raises(ValueError, match="room_nodes"):
+        render_topdown_scene_map(
+            pathfinder=SimpleNamespace(),
+            elements=TopDownSceneElements(room_nodes=[]),
+            meters_per_pixel=0.05,
+            maps_module=_FakeMaps(),
+            plt_module=_FakePlt(),
+        )
+
+
+def test_render_topdown_scene_map_validates_meters_per_pixel():
+    with pytest.raises(ValueError, match="meters_per_pixel"):
+        render_topdown_scene_map(
+            pathfinder=SimpleNamespace(),
+            elements=TopDownSceneElements(
+                room_nodes=[SimpleNamespace(room_id="room_00", center=[0.0, 1.0, 0.0])]
+            ),
+            meters_per_pixel=0,
+            maps_module=_FakeMaps(),
+            plt_module=_FakePlt(),
+        )
+
+
+def test_render_topdown_scene_map_rejects_object_nodes_before_implementation():
+    with pytest.raises(NotImplementedError, match="object_nodes"):
+        render_topdown_scene_map(
+            pathfinder=SimpleNamespace(),
+            elements=TopDownSceneElements(
+                room_nodes=[SimpleNamespace(room_id="room_00", center=[0.0, 1.0, 0.0])],
+                object_nodes=[SimpleNamespace(obj_id="obj_00")],
+            ),
+            meters_per_pixel=0.05,
+            maps_module=_FakeMaps(),
+            plt_module=_FakePlt(),
+        )
+
+
+def test_render_topdown_scene_map_rejects_edges_before_implementation():
+    with pytest.raises(NotImplementedError, match="edge"):
+        render_topdown_scene_map(
+            pathfinder=SimpleNamespace(),
+            elements=TopDownSceneElements(
+                room_nodes=[SimpleNamespace(room_id="room_00", center=[0.0, 1.0, 0.0])],
+                edges=[("room_00", "obj_00")],
+            ),
+            meters_per_pixel=0.05,
+            maps_module=_FakeMaps(),
+            plt_module=_FakePlt(),
+        )

notebooks/verification.py

@@ -8,7 +8,7 @@
 
 import marimo
 
-__generated_with = "0.20.4"
+__generated_with = "0.21.1"
 app = marimo.App(width="medium", app_title="Pipeline Verification")
 
 
@@ -16,81 +16,87 @@ app = marimo.App(width="medium", app_title="Pipeline Verification")
 def import_packages():
     from io import BytesIO
 
-    import habitat_sim
     import marimo as mo
     import numpy as np
     import polars as pl
-    from habitat.utils.visualizations import maps
-    from matplotlib import pyplot as plt
     from PIL import Image
+
     from compressors.pipeline import HashPipeline
     from scenegraph import ObjectNode, RoomNode, SimpleSceneGraph
-    from utils.common import get_device
+    from simulator import (
+        HabitatSimulatorConfig,
+        TopDownSceneElements,
+        collect_room_views_by_room,
+        create_habitat_simulator,
+        render_topdown_scene_map,
+    )
     from utils.image import extract_masked_region, segment_image
 
     return (
         BytesIO,
+        HabitatSimulatorConfig,
         HashPipeline,
         Image,
         ObjectNode,
         RoomNode,
         SimpleSceneGraph,
+        TopDownSceneElements,
+        collect_room_views_by_room,
+        create_habitat_simulator,
         extract_masked_region,
-        habitat_sim,
-        maps,
         mo,
         np,
         pl,
-        plt,
+        render_topdown_scene_map,
         segment_image,
     )
 
 
 @app.cell
-def setup_habitat_simulator(habitat_sim):
+def setup_verification_context(
+    HabitatSimulatorConfig, RoomNode, create_habitat_simulator, np
+):
     scene_path = "data/scene_datasets/habitat-test-scenes/skokloster-castle.glb"
+    image_size = 256
     num_rooms = 4
     views_per_room = 6
-    image_size = 256
     meters_per_pixel = 0.05
 
-    sim_cfg = habitat_sim.SimulatorConfiguration()
-    sim_cfg.scene_id = scene_path
-    sim_cfg.enable_physics = False
-
-    agent_cfg = habitat_sim.agent.AgentConfiguration()
-    rgb_sensor_spec = habitat_sim.CameraSensorSpec()
-    rgb_sensor_spec.uuid = "color_sensor"
-    rgb_sensor_spec.sensor_type = habitat_sim.SensorType.COLOR
-    rgb_sensor_spec.resolution = [image_size, image_size]
-    rgb_sensor_spec.position = [0.0, 1.5, 0.0]
-    agent_cfg.sensor_specifications = [rgb_sensor_spec]
-
-    turn_angle = 360.0 / views_per_room
-    agent_cfg.action_space = {
-        "move_forward": habitat_sim.agent.ActionSpec(
-            "move_forward", habitat_sim.agent.ActuationSpec(amount=0.25)
-        ),
-        "turn_left": habitat_sim.agent.ActionSpec(
-            "turn_left", habitat_sim.agent.ActuationSpec(amount=turn_angle)
-        ),
-        "turn_right": habitat_sim.agent.ActionSpec(
-            "turn_right", habitat_sim.agent.ActuationSpec(amount=turn_angle)
-        ),
-    }
-
-    cfg = habitat_sim.Configuration(sim_cfg, [agent_cfg])
-    sim = habitat_sim.Simulator(cfg)
-    agent = sim.initialize_agent(0)
     sam_max_masks = 5
     sam_min_area = 32 * 32
-    sam_points_per_batch = 64
     hash_bits = 512
+
+    sim, agent = create_habitat_simulator(
+        HabitatSimulatorConfig(
+            scene_path=scene_path,
+            views_per_room=views_per_room,
+            image_size=image_size,
+            sensor_height=1.5,
+            move_forward_step=0.25,
+            enable_physics=False,
+        )
+    )
+
+    room_nodes = []
+    for idx in range(num_rooms):
+        point = sim.pathfinder.get_random_navigable_point()
+        room_nodes.append(
+            RoomNode(
+                room_id=f"room_{idx:02d}",
+                center=np.asarray(point, dtype=np.float32),
+                bbox_extent=np.asarray([1.5, 2.0, 1.5], dtype=np.float32),
+            )
+        )
+
+    print("Sampled room centers:")
+    for node in room_nodes:
+        print(node.room_id, node.center)
+
     return (
         agent,
         hash_bits,
         meters_per_pixel,
-        num_rooms,
+        room_nodes,
         sam_max_masks,
         sam_min_area,
         sim,
@@ -99,100 +105,47 @@ def setup_habitat_simulator(habitat_sim):
 
 
 @app.cell
-def sample_room_nodes(RoomNode, np, num_rooms, sim):
+def render_topdown_room_map(
-    room_nodes = []
+    TopDownSceneElements,
-    for _idx in range(num_rooms):
+    meters_per_pixel,
-        _point = sim.pathfinder.get_random_navigable_point()
+    render_topdown_scene_map,
-        _room_node = RoomNode(
+    room_nodes,
-            room_id=f"room_{_idx:02d}",
+    sim,
-            center=np.asarray(_point, dtype=np.float32),
+):
-            bbox_extent=np.asarray([1.5, 2.0, 1.5], dtype=np.float32),
+    render_topdown_scene_map(
-        )
+        pathfinder=sim.pathfinder,
-        room_nodes.append(_room_node)
+        elements=TopDownSceneElements(room_nodes=room_nodes),
-
-    print("Sampled room centers:")
-    for _node in room_nodes:
-        print(_node.room_id, _node.center)
-    return (room_nodes,)
-
-
-@app.cell
-def render_topdown_room_map(maps, meters_per_pixel, plt, room_nodes, sim):
-    top_down_map = maps.get_topdown_map(
-        sim.pathfinder,
-        height=float(room_nodes[0].center[1]),
         meters_per_pixel=meters_per_pixel,
     )
-
-    plt.figure(figsize=(8, 8))
-    plt.imshow(top_down_map, cmap="gray")
-
-    for _node in room_nodes:
-        _gy, _gx = maps.to_grid(
-            float(_node.center[2]),
-            float(_node.center[0]),
-            top_down_map.shape,
-            pathfinder=sim.pathfinder,
-        )
-        plt.scatter(_gx, _gy, c="red", s=50)
-        plt.text(_gx + 2, _gy + 2, _node.room_id, color="yellow", fontsize=8)
-
-    plt.title("RoomNode Top-Down Map")
-    plt.axis("off")
-    plt.show()
     return
 
 
 @app.cell
-def collect_room_views(
+def build_scene_graph_pipeline(
     agent,
-    habitat_sim,
+    HashPipeline,
+    Image,
+    ObjectNode,
+    SimpleSceneGraph,
+    collect_room_views_by_room,
+    extract_masked_region,
+    hash_bits,
     mo,
-    plt,
+    np,
     room_nodes,
+    sam_max_masks,
+    sam_min_area,
+    segment_image,
     sim,
     views_per_room,
 ):
-    all_room_views = {}
+    all_room_views = collect_room_views_by_room(
+        agent=agent,
+        sim=sim,
+        room_nodes=room_nodes,
+        views_per_room=views_per_room,
+    )
 
-    for _node in mo.status.progress_bar(
-        room_nodes,
-        title="Collecting room views",
-        subtitle="Sampling observations from Habitat",
-        show_eta=True,
-        show_rate=True,
-    ):
-        _agent_state = habitat_sim.AgentState()
-        _agent_state.position = _node.center.copy()
-        agent.set_state(_agent_state)
-
-        _room_views = []
-        for _ in mo.status.progress_bar(
-            range(views_per_room),
-            title=f"Capturing {_node.room_id}",
-            subtitle="Rotating agent viewpoints",
-            show_eta=True,
-            show_rate=True,
-        ):
-            _observations = sim.get_sensor_observations()
-            _rgb = _observations["color_sensor"]
-            _room_views.append(_rgb)
-            sim.step("turn_left")
-
-        all_room_views[_node.room_id] = _room_views
-
-        _fig, _axes = plt.subplots(2, 3, figsize=(10, 6))
-        for _view_idx, _ax in enumerate(_axes.flatten()):
-            _ax.imshow(_room_views[_view_idx])
-            _ax.set_title(f"{_node.room_id} - view {_view_idx + 1}")
-            _ax.axis("off")
-        plt.tight_layout()
-        plt.show()
-    return (all_room_views,)
-
-
-@app.cell
-def build_hash_pipeline(HashPipeline, hash_bits, sam_max_masks, sam_min_area):
     hash_pipeline = HashPipeline(
         dino_model="facebook/dinov2-large",
         sam_model="facebook/sam2.1-hiera-large",
@@ -200,153 +153,117 @@ def build_hash_pipeline(HashPipeline, hash_bits, sam_max_masks, sam_min_area):
         sam_max_masks=sam_max_masks,
         hash_bits=hash_bits,
     )
-    return (hash_pipeline,)
 
-
-@app.cell
-def build_scene_graph_from_views(
-    Image,
-    ObjectNode,
-    SimpleSceneGraph,
-    all_room_views,
-    extract_masked_region,
-    hash_pipeline,
-    mo,
-    np,
-    room_nodes,
-    segment_image,
-):
     scene_graph = SimpleSceneGraph(
-        rooms={_room.room_id: _room for _room in room_nodes}, objects={}
+        rooms={room.room_id: room for room in room_nodes},
+        objects={},
     )
     total_masks = 0
-    _obj_index = 0
+    object_index = 0
 
-    _view_jobs = [
+    view_jobs = [
-        (_room_id, _view_idx, _rgb)
+        (room_id, view_idx, rgb)
-        for _room_id, _views in all_room_views.items()
+        for room_id, views in all_room_views.items()
-        for _view_idx, _rgb in enumerate(_views)
+        for view_idx, rgb in enumerate(views)
     ]
 
-    for _room_id, _view_idx, _rgb in mo.status.progress_bar(
+    for room_id, _view_idx, rgb in mo.status.progress_bar(
-        _view_jobs,
+        view_jobs,
         title="Extracting masks and hashes",
         subtitle="Running SAM + HashPipeline",
         show_eta=True,
         show_rate=True,
     ):
-        _rgb3 = _rgb[..., :3] if _rgb.shape[-1] > 3 else _rgb
+        rgb3 = rgb[..., :3] if rgb.shape[-1] > 3 else rgb
-        _image = Image.fromarray(_rgb3.astype(np.uint8))
+        image = Image.fromarray(rgb3.astype(np.uint8))
 
-        _masks = segment_image(
+        masks = segment_image(
             hash_pipeline.mask_generator,
-            _image,
+            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)
+        total_masks += len(masks)
 
-        for _mask in _masks:
+        for mask in masks:
-            _masked_image = extract_masked_region(_image, _mask["segment"])
+            masked_image = extract_masked_region(image, mask["segment"])
-            _bits = hash_pipeline(_masked_image)
+            bits = hash_pipeline(masked_image)
 
-            _bbox = _mask["bbox"]
+            bbox = mask["bbox"]
-            _obj_center = np.array(
+            obj_center = np.array(
-                [_bbox[0] + _bbox[2] / 2, _bbox[1] + _bbox[3] / 2, 0.0],
+                [bbox[0] + bbox[2] / 2, bbox[1] + bbox[3] / 2, 0.0],
                 dtype=np.float32,
             )
 
-            _obj_id = f"obj_{_obj_index:04d}"
+            obj_id = f"obj_{object_index:04d}"
-            _obj_index += 1
+            object_index += 1
-            _bits_np = _bits.squeeze().detach().cpu().numpy()
+            bits_np = bits.squeeze().detach().cpu().numpy()
 
-            _obj_node = ObjectNode(
+            scene_graph.objects[obj_id] = ObjectNode(
-                obj_id=_obj_id,
+                obj_id=obj_id,
-                room_id=_room_id,
+                room_id=room_id,
-                position=_obj_center,
+                position=obj_center,
-                visual_hash=_bits_np,
+                visual_hash=bits_np,
-                semantic_hash=_bits_np,
+                semantic_hash=bits_np,
-                image_bytes=np.array(_masked_image).tobytes(),
                 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}")
-    return (scene_graph,)
+    return all_room_views, hash_pipeline, scene_graph
 
 
 @app.cell
 def build_room_and_object_tables(pl, scene_graph):
-    _room_rows = []
+    room_rows = [
-    for _room in scene_graph.rooms.values():
+        {
-        _room_rows.append(
+            "room_id": room.room_id,
-            {
+            "center_x": float(room.center[0]),
-                "room_id": _room.room_id,
+            "center_y": float(room.center[1]),
-                "center_x": float(_room.center[0]),
+            "center_z": float(room.center[2]),
-                "center_y": float(_room.center[1]),
+            "bbox_dx": float(room.bbox_extent[0]),
-                "center_z": float(_room.center[2]),
+            "bbox_dy": float(room.bbox_extent[1]),
-                "bbox_dx": float(_room.bbox_extent[0]),
+            "bbox_dz": float(room.bbox_extent[2]),
-                "bbox_dy": float(_room.bbox_extent[1]),
+        }
-                "bbox_dz": float(_room.bbox_extent[2]),
+        for room in scene_graph.rooms.values()
-            }
+    ]
-        )
 
-    _object_rows = []
+    object_rows = [
-    for _obj in scene_graph.objects.values():
+        {
-        _object_rows.append(
+            "obj_id": obj.obj_id,
-            {
+            "room_id": obj.room_id,
-                "obj_id": _obj.obj_id,
+            "last_seen_frame": int(obj.last_seen_frame),
-                "room_id": _obj.room_id,
+            "hit_count": int(obj.hit_count),
-                "last_seen_frame": int(_obj.last_seen_frame),
+            "visual_hash": obj.visual_hash.tolist(),
-                "hit_count": int(_obj.hit_count),
+            "semantic_hash": obj.semantic_hash.tolist(),
-                "visual_hash": _obj.visual_hash.tolist(),
+        }
-                "semantic_hash": _obj.semantic_hash.tolist(),
+        for obj in scene_graph.objects.values()
-            }
+    ]
-        )
 
-    rooms_table = pl.DataFrame(_room_rows)
+    rooms_table = pl.DataFrame(room_rows)
-    objects_table = pl.DataFrame(_object_rows)
+    objects_table = pl.DataFrame(object_rows)
     return objects_table, rooms_table
 
 
-@app.cell(disabled=True)
-def display_rooms_table(rooms_table):
-    rooms_table
-    return
-
-
-@app.cell(disabled=True)
-def display_objects_table(objects_table):
-    objects_table
-    return
-
-
 @app.cell
-def create_file_upload(mo):
+def upload_query_image(BytesIO, Image, mo, np):
     file_upload = mo.ui.file(
-        filetypes=["image/*"], kind="area", label="Upload a query image"
+        filetypes=["image/*"],
+        kind="area",
+        label="Upload a query image",
     )
     file_upload
-    return (file_upload,)
 
-
-@app.cell
-def load_uploaded_image(BytesIO, Image, file_upload):
     uploaded_image = None
     if file_upload.value:
-        _contents = file_upload.contents()
+        contents = file_upload.contents()
-        if _contents:
+        if contents:
-            uploaded_image = Image.open(BytesIO(_contents))
+            uploaded_image = Image.open(BytesIO(contents))
-    return (uploaded_image,)
+            mo.image(np.array(uploaded_image), alt="Uploaded query image")
 
-
+    return file_upload, uploaded_image
-@app.cell
-def display_uploaded_image(mo, np, uploaded_image):
-    mo.image(np.array(uploaded_image), alt="Uploaded query image")
-    return
 
 
 if __name__ == "__main__":