feat(notebooks): add habitat simulation and room visualization

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