mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-12 20:15:31 +08:00
refactor(simulator): extract habitat simulator and visualization modules
This commit is contained in:
@@ -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__":
|
||||
|
||||
Reference in New Issue
Block a user