From 3c9a6f6eafc463540e59ad884cc06750e96cd824 Mon Sep 17 00:00:00 2001 From: SikongJueluo Date: Sat, 28 Mar 2026 17:05:06 +0800 Subject: [PATCH] refactor(simulator): extract habitat simulator and visualization modules --- mini-nav/scenegraph/objectnode.py | 15 +- mini-nav/simulator/__init__.py | 18 ++ mini-nav/simulator/habitat.py | 73 +++++ mini-nav/simulator/topdown.py | 84 ++++++ mini-nav/simulator/views.py | 44 +++ mini-nav/tests/test_habitat_simulator.py | 124 ++++++++ mini-nav/tests/test_room_views.py | 108 +++++++ mini-nav/tests/test_topdown_map.py | 126 ++++++++ notebooks/verification.py | 359 +++++++++-------------- 9 files changed, 725 insertions(+), 226 deletions(-) create mode 100644 mini-nav/simulator/__init__.py create mode 100644 mini-nav/simulator/habitat.py create mode 100644 mini-nav/simulator/topdown.py create mode 100644 mini-nav/simulator/views.py create mode 100644 mini-nav/tests/test_habitat_simulator.py create mode 100644 mini-nav/tests/test_room_views.py create mode 100644 mini-nav/tests/test_topdown_map.py diff --git a/mini-nav/scenegraph/objectnode.py b/mini-nav/scenegraph/objectnode.py index 4af2261..b1ba4c7 100644 --- a/mini-nav/scenegraph/objectnode.py +++ b/mini-nav/scenegraph/objectnode.py @@ -13,12 +13,17 @@ class ObjectNode: position: np.ndarray # [x, y, z] 世界坐标系下的中心点或锚点 # 特征：直接存你压缩后的 512bit 结果，不搞历史缓存 - visual_hash: np.ndarray # 512bit 视觉特征 (用于外观检索) - semantic_hash: np.ndarray # 512bit 语义特征 (用于类别对齐) - - # 原始图片：mask处理并裁切后的图片数据（二进制格式） - image_bytes: bytes | None = None # numpy array.tobytes() 原始像素数据 + visual_hash: bytes # 512bit 视觉特征 (用于外观检索) + semantic_hash: bytes # 512bit 语义特征 (用于类别对齐) # 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,)") diff --git a/mini-nav/simulator/__init__.py b/mini-nav/simulator/__init__.py new file mode 100644 index 0000000..df74d74 --- /dev/null +++ b/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", +] diff --git a/mini-nav/simulator/habitat.py b/mini-nav/simulator/habitat.py new file mode 100644 index 0000000..dcf4e11 --- /dev/null +++ b/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() diff --git a/mini-nav/simulator/topdown.py b/mini-nav/simulator/topdown.py new file mode 100644 index 0000000..5dee7c5 --- /dev/null +++ b/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 diff --git a/mini-nav/simulator/views.py b/mini-nav/simulator/views.py new file mode 100644 index 0000000..47c4be2 --- /dev/null +++ b/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 diff --git a/mini-nav/tests/test_habitat_simulator.py b/mini-nav/tests/test_habitat_simulator.py new file mode 100644 index 0000000..46c99b4 --- /dev/null +++ b/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) diff --git a/mini-nav/tests/test_room_views.py b/mini-nav/tests/test_room_views.py new file mode 100644 index 0000000..f15a129 --- /dev/null +++ b/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, + ) diff --git a/mini-nav/tests/test_topdown_map.py b/mini-nav/tests/test_topdown_map.py new file mode 100644 index 0000000..931b6e9 --- /dev/null +++ b/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(), + ) diff --git a/notebooks/verification.py b/notebooks/verification.py index c4c7cca..eccbc17 100644 --- a/notebooks/verification.py +++ b/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): - room_nodes = [] - for _idx in range(num_rooms): - _point = sim.pathfinder.get_random_navigable_point() - _room_node = 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), - ) - room_nodes.append(_room_node) - - 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]), +def render_topdown_room_map( + TopDownSceneElements, + meters_per_pixel, + render_topdown_scene_map, + room_nodes, + sim, +): + render_topdown_scene_map( + pathfinder=sim.pathfinder, + elements=TopDownSceneElements(room_nodes=room_nodes), 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 = [ - (_room_id, _view_idx, _rgb) - for _room_id, _views in all_room_views.items() - for _view_idx, _rgb in enumerate(_views) + view_jobs = [ + (room_id, view_idx, rgb) + for room_id, views in all_room_views.items() + for view_idx, rgb in enumerate(views) ] - for _room_id, _view_idx, _rgb in mo.status.progress_bar( - _view_jobs, + for room_id, _view_idx, rgb in mo.status.progress_bar( + 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 - _image = Image.fromarray(_rgb3.astype(np.uint8)) + rgb3 = rgb[..., :3] if rgb.shape[-1] > 3 else rgb + 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: - _masked_image = extract_masked_region(_image, _mask["segment"]) - _bits = hash_pipeline(_masked_image) + 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], + 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_id = f"obj_{object_index:04d}" + object_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, - image_bytes=np.array(_masked_image).tobytes(), + scene_graph.objects[obj_id] = 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}") - return (scene_graph,) + return all_room_views, hash_pipeline, scene_graph @app.cell def build_room_and_object_tables(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]), - } - ) + room_rows = [ + { + "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]), + } + for room in scene_graph.rooms.values() + ] - _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(), - } - ) + object_rows = [ + { + "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(), + } + for obj in scene_graph.objects.values() + ] - rooms_table = pl.DataFrame(_room_rows) - objects_table = pl.DataFrame(_object_rows) + rooms_table = pl.DataFrame(room_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() - if _contents: - uploaded_image = Image.open(BytesIO(_contents)) - return (uploaded_image,) + contents = file_upload.contents() + if contents: + uploaded_image = Image.open(BytesIO(contents)) + mo.image(np.array(uploaded_image), alt="Uploaded query image") - -@app.cell -def display_uploaded_image(mo, np, uploaded_image): - mo.image(np.array(uploaded_image), alt="Uploaded query image") - return + return file_upload, uploaded_image if __name__ == "__main__":