SikongJueluo/Mini-Nav
1
0
Fork 0
mirror of https://github.com/SikongJueluo/Mini-Nav.git synced 2026-07-12 20:15:31 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(simulator): extract habitat simulator and visualization modules

Browse Source
This commit is contained in:
SikongJueluo
2026-03-28 17:05:06 +08:00
parent 817f45b935
commit 3c9a6f6eaf
9 changed files with 725 additions and 226 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
mini-nav/scenegraph/objectnode.py
View File

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

18
mini-nav/simulator/__init__.py Normal file
View File

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

73
mini-nav/simulator/habitat.py Normal file
View File

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

84
mini-nav/simulator/topdown.py Normal file
View File

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

44
mini-nav/simulator/views.py Normal file
View File

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

124
mini-nav/tests/test_habitat_simulator.py Normal file
View File

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

108
mini-nav/tests/test_room_views.py Normal file
View File

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

126
mini-nav/tests/test_topdown_map.py Normal file
View File

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

351
notebooks/verification.py
View File

@@ -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_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]),
"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(
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(),
"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,)
@app.cell
def display_uploaded_image(mo, np, uploaded_image):
contents = file_upload.contents()
if contents:
uploaded_image = Image.open(BytesIO(contents))
mo.image(np.array(uploaded_image), alt="Uploaded query image")
return
return file_upload, uploaded_image
if __name__ == "__main__":