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

359
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_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__":