From 1af9a9caa15e6e4e27c0e8f5ef7e2c8766ba2680 Mon Sep 17 00:00:00 2001 From: SikongJueluo Date: Tue, 24 Mar 2026 21:02:42 +0800 Subject: [PATCH] feat(notebooks): add habitat simulation and room visualization --- .stignore | 2 +- notebooks/verification.py | 115 +++++++++++++++++++++++++++++++++++++- 2 files changed, 113 insertions(+), 4 deletions(-) diff --git a/.stignore b/.stignore index 7bec06c..6589e3d 100644 --- a/.stignore +++ b/.stignore @@ -3,8 +3,8 @@ .venv .pytest_cache .ruff_cache -.backup .opencode +.backup data datasets deps diff --git a/notebooks/verification.py b/notebooks/verification.py index 1e04b06..669ccc3 100644 --- a/notebooks/verification.py +++ b/notebooks/verification.py @@ -6,16 +6,125 @@ app = marimo.App() @app.cell def _(): - import marimo as mo - from scenegraph import ObjectNode, RoomNode + import habitat_sim + import numpy as np + from habitat.utils.visualizations import maps + from matplotlib import pyplot as plt + from scenegraph import RoomNode + + return RoomNode, habitat_sim, maps, np, plt +@app.cell +def _(habitat_sim): + scene_path = "data/scene_datasets/habitat-test-scenes/skokloster-castle.glb" + 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) + return agent, meters_per_pixel, num_rooms, sim, views_per_room + + +@app.cell +def _(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) + + room_points = np.vstack([_node.center for _node in room_nodes]) + print("Sampled room centers:") + for _node in room_nodes: + print(_node.room_id, _node.center) + return (room_nodes,) + + +@app.cell +def _(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 _(): +def _(agent, habitat_sim, plt, room_nodes, sim, views_per_room): + all_room_views = {} + + for _node in room_nodes: + _agent_state = habitat_sim.AgentState() + _agent_state.position = _node.center.copy() + agent.set_state(_agent_state) + + _room_views = [] + for _ in range(views_per_room): + _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