mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-12 20:15:31 +08:00
SikongJueluo
e5b764520c
- Add save_scene_graph() and load_scene_graph() for persisting scene graphs to JSON - Implement object node overlay (colored by label) in render_topdown_scene_map - Implement edge arrows (room → object) in render_topdown_scene_map - Add TopDownRenderStyle fields for object/edge visual configuration - Add scene graph caching to verification notebook to skip rebuilds - Add render_scene_graph_birdseye cell for full scene graph visualization - Add display_objects_by_room and display_room_summary cells
537 lines
14 KiB
Python
537 lines
14 KiB
Python
# /// script
|
|
# requires-python = ">=3.13"
|
|
# dependencies = [
|
|
# "marimo>=0.21.1",
|
|
# "pyzmq>=27.1.0",
|
|
# ]
|
|
# ///
|
|
|
|
import marimo
|
|
|
|
__generated_with = "0.21.1"
|
|
app = marimo.App(width="medium", app_title="Pipeline Verification")
|
|
|
|
|
|
@app.cell
|
|
def _():
|
|
import marimo as mo
|
|
|
|
return (mo,)
|
|
|
|
|
|
@app.cell
|
|
def base_dependencies():
|
|
"""Basic dependencies for data processing."""
|
|
import numpy as np
|
|
import polars as pl
|
|
from PIL import Image
|
|
|
|
return Image, np, pl
|
|
|
|
|
|
@app.cell
|
|
def project_imports():
|
|
"""Project module imports using new architecture."""
|
|
from compressors import HashPipeline
|
|
from configs import cfg_manager
|
|
from scenegraph import (
|
|
RoomNode,
|
|
SceneGraphBuildConfig,
|
|
SceneGraphBuilder,
|
|
load_scene_graph,
|
|
query_image_against_scene_graph,
|
|
save_scene_graph,
|
|
)
|
|
from simulator import (
|
|
HabitatSimulatorConfig,
|
|
TopDownRenderStyle,
|
|
TopDownSceneElements,
|
|
collect_room_views_by_room,
|
|
create_habitat_simulator,
|
|
flatten_room_views,
|
|
render_topdown_scene_map,
|
|
save_object_image,
|
|
save_room_view,
|
|
)
|
|
from utils.image import numpy_to_pil
|
|
|
|
return (
|
|
HashPipeline,
|
|
HabitatSimulatorConfig,
|
|
RoomNode,
|
|
SceneGraphBuildConfig,
|
|
SceneGraphBuilder,
|
|
TopDownRenderStyle,
|
|
TopDownSceneElements,
|
|
cfg_manager,
|
|
collect_room_views_by_room,
|
|
create_habitat_simulator,
|
|
flatten_room_views,
|
|
load_scene_graph,
|
|
numpy_to_pil,
|
|
query_image_against_scene_graph,
|
|
render_topdown_scene_map,
|
|
save_object_image,
|
|
save_room_view,
|
|
save_scene_graph,
|
|
)
|
|
|
|
|
|
@app.cell
|
|
def text_labels():
|
|
"""Shared text labels for detection during graph build and query."""
|
|
text_labels = [
|
|
"a chair",
|
|
"a table",
|
|
"a sofa",
|
|
"a cabinet",
|
|
"a shelf",
|
|
"a lamp",
|
|
"a picture",
|
|
"a window",
|
|
"a door",
|
|
"a plant",
|
|
]
|
|
return (text_labels,)
|
|
|
|
|
|
@app.cell
|
|
def habitat_setup(HabitatSimulatorConfig, RoomNode, create_habitat_simulator, np):
|
|
"""Initialize Habitat simulator and sample room nodes."""
|
|
_scene_path = "data/scene_datasets/habitat-test-scenes/skokloster-castle.glb"
|
|
_image_size = 768
|
|
_num_rooms = 5
|
|
views_per_room = 12
|
|
meters_per_pixel = 0.05
|
|
|
|
habitat_config = 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,
|
|
)
|
|
sim, agent = create_habitat_simulator(habitat_config)
|
|
|
|
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(f" {_node.room_id}: {_node.center}")
|
|
|
|
return agent, habitat_config, meters_per_pixel, room_nodes, sim, views_per_room
|
|
|
|
|
|
@app.cell
|
|
def render_topdown(
|
|
TopDownSceneElements,
|
|
meters_per_pixel,
|
|
mo,
|
|
render_topdown_scene_map,
|
|
room_nodes,
|
|
sim,
|
|
):
|
|
image = render_topdown_scene_map(
|
|
pathfinder=sim.pathfinder,
|
|
elements=TopDownSceneElements(room_nodes=room_nodes),
|
|
meters_per_pixel=meters_per_pixel,
|
|
)
|
|
mo.image(image)
|
|
return (image,)
|
|
|
|
|
|
@app.cell
|
|
def pipeline_init(HashPipeline):
|
|
pipeline = HashPipeline(
|
|
dino_model="facebook/dinov2-large",
|
|
sam_model="facebook/sam2.1-hiera-large",
|
|
hash_bits=512,
|
|
score_threshold=0.10,
|
|
postprocess_threshold=0.05,
|
|
)
|
|
print(f"Pipeline initialized: {pipeline.hash_bits} bits")
|
|
return (pipeline,)
|
|
|
|
|
|
@app.cell
|
|
def collect_views(
|
|
agent,
|
|
collect_room_views_by_room,
|
|
flatten_room_views,
|
|
habitat_config,
|
|
numpy_to_pil,
|
|
room_nodes,
|
|
sim,
|
|
views_per_room,
|
|
):
|
|
all_room_views = collect_room_views_by_room(
|
|
agent=agent,
|
|
sim=sim,
|
|
room_nodes=room_nodes,
|
|
views_per_room=views_per_room,
|
|
depth_sensor_uuid=habitat_config.depth_sensor_uuid,
|
|
)
|
|
|
|
room_views = flatten_room_views(all_room_views)
|
|
|
|
# Build dataset of (room_id, view_idx, PIL.Image) tuples.
|
|
room_view_dataset = [
|
|
(_view.room_id, _view.view_idx, numpy_to_pil(_view.rgb))
|
|
for _view in room_views
|
|
]
|
|
|
|
print(f"Collected {len(room_view_dataset)} room views")
|
|
return all_room_views, room_view_dataset, room_views
|
|
|
|
|
|
@app.cell
|
|
def build_scene_graph(
|
|
SceneGraphBuildConfig,
|
|
SceneGraphBuilder,
|
|
cfg_manager,
|
|
habitat_config,
|
|
load_scene_graph,
|
|
mo,
|
|
numpy_to_pil,
|
|
pipeline,
|
|
room_nodes,
|
|
room_views,
|
|
save_object_image,
|
|
save_room_view,
|
|
save_scene_graph,
|
|
text_labels,
|
|
):
|
|
output_dir = cfg_manager.get().output.directory / "verification"
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
cache_path = output_dir / "scene_graph.json"
|
|
|
|
# Try loading cached scene graph first.
|
|
if cache_path.exists():
|
|
scene_graph = load_scene_graph(cache_path)
|
|
print(f"Loaded cached scene graph from {cache_path}")
|
|
print(f" {len(scene_graph.rooms)} rooms, {len(scene_graph.objects)} objects")
|
|
object_images = {}
|
|
build_artifacts = None
|
|
else:
|
|
builder = SceneGraphBuilder(
|
|
pipeline=pipeline,
|
|
config=SceneGraphBuildConfig(
|
|
inference_batch_size=4,
|
|
position_strategy="bbox_depth_center",
|
|
camera_hfov_degrees=habitat_config.hfov_degrees,
|
|
),
|
|
)
|
|
|
|
pil_room_views = [
|
|
type(_view)(
|
|
room_id=_view.room_id,
|
|
view_idx=_view.view_idx,
|
|
rgb=numpy_to_pil(_view.rgb),
|
|
depth=_view.depth,
|
|
agent_position=_view.agent_position,
|
|
agent_rotation=_view.agent_rotation,
|
|
camera_position=_view.camera_position,
|
|
camera_rotation=_view.camera_rotation,
|
|
)
|
|
for _view in room_views
|
|
]
|
|
|
|
with mo.status.spinner(title="Building scene graph from room views"):
|
|
scene_graph, build_artifacts = builder.build_from_room_views(
|
|
room_nodes=room_nodes,
|
|
room_views=pil_room_views,
|
|
text_labels=text_labels,
|
|
)
|
|
|
|
object_images = build_artifacts.object_images
|
|
debug_meta = build_artifacts.debug_meta
|
|
|
|
# Save scene graph to cache.
|
|
save_scene_graph(cache_path, scene_graph)
|
|
print(f"Saved scene graph to {cache_path}")
|
|
|
|
# Save original room views.
|
|
for _room_view in mo.status.progress_bar(
|
|
pil_room_views,
|
|
title="Saving room-view snapshots",
|
|
subtitle="Writing original room images to disk",
|
|
completion_title="Room-view snapshots saved",
|
|
completion_subtitle=f"Saved {len(pil_room_views)} room views",
|
|
show_eta=True,
|
|
show_rate=True,
|
|
remove_on_exit=False,
|
|
):
|
|
save_room_view(output_dir, _room_view.room_id, _room_view.view_idx, _room_view.rgb)
|
|
|
|
# Save object crops.
|
|
for _obj_id, _cropped in mo.status.progress_bar(
|
|
object_images.items(),
|
|
title="Saving object crops",
|
|
subtitle="Writing cropped object images to disk",
|
|
completion_title="Object crops saved",
|
|
completion_subtitle=f"Saved {len(object_images)} object crops",
|
|
show_eta=True,
|
|
show_rate=True,
|
|
remove_on_exit=False,
|
|
):
|
|
_node = scene_graph.objects[_obj_id]
|
|
save_object_image(
|
|
output_dir,
|
|
_node.room_id,
|
|
_obj_id,
|
|
_node.last_seen_frame,
|
|
0, # mask idx 0 ok for M0
|
|
_cropped,
|
|
)
|
|
|
|
from collections import Counter
|
|
|
|
_meta_fallbacks = [_meta.get("fallback_reason") for _meta in debug_meta]
|
|
fallback_count = sum(1 for f in _meta_fallbacks if f is not None)
|
|
_reasons = Counter(f or "ok" for f in _meta_fallbacks)
|
|
print(f"Fallback breakdown: {_reasons}")
|
|
print(f"Fallback frames: {fallback_count}/{len(debug_meta)}")
|
|
|
|
print(f"Created {len(scene_graph.objects)} objects")
|
|
print(f"Output directory: {output_dir}")
|
|
|
|
return build_artifacts, object_images, output_dir, scene_graph
|
|
|
|
|
|
@app.cell
|
|
def render_scene_graph_birdseye(
|
|
TopDownRenderStyle,
|
|
TopDownSceneElements,
|
|
meters_per_pixel,
|
|
mo,
|
|
render_topdown_scene_map,
|
|
room_nodes,
|
|
scene_graph,
|
|
sim,
|
|
):
|
|
"""Bird's-eye view of the full scene graph: rooms, objects, and edges."""
|
|
_object_nodes = list(scene_graph.objects.values())
|
|
_edges = [
|
|
(_obj.room_id, _obj.obj_id) for _obj in _object_nodes
|
|
]
|
|
|
|
_style = TopDownRenderStyle(
|
|
title="Scene Graph — Bird's-Eye View",
|
|
figure_size=(10, 10),
|
|
)
|
|
|
|
_elements = TopDownSceneElements(
|
|
room_nodes=room_nodes,
|
|
object_nodes=_object_nodes,
|
|
edges=_edges,
|
|
)
|
|
|
|
_birdseye_image = render_topdown_scene_map(
|
|
pathfinder=sim.pathfinder,
|
|
elements=_elements,
|
|
meters_per_pixel=meters_per_pixel,
|
|
style=_style,
|
|
)
|
|
|
|
mo.md("## Scene Graph Bird's-Eye View"), mo.image(_birdseye_image)
|
|
return (_birdseye_image,)
|
|
|
|
|
|
@app.cell
|
|
def build_tables(pl, scene_graph):
|
|
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]),
|
|
}
|
|
for _room in scene_graph.rooms.values()
|
|
]
|
|
|
|
object_rows = [
|
|
{
|
|
"obj_id": obj.obj_id,
|
|
"room_id": obj.room_id,
|
|
"visual_hash": obj.visual_hash.hex()[:16] + "...",
|
|
"label": obj.label,
|
|
"confidence": obj.confidence,
|
|
"source_view_id": obj.source_view_id,
|
|
"position_x": float(obj.position[0]) if obj.position is not None else None,
|
|
"position_y": float(obj.position[1]) if obj.position is not None else None,
|
|
"position_z": float(obj.position[2]) if obj.position is not None else None,
|
|
"bbox_xyxy": str(obj.bbox_xyxy) if obj.bbox_xyxy is not None else None,
|
|
"position_confidence": obj.position_confidence,
|
|
}
|
|
for obj in scene_graph.objects.values()
|
|
]
|
|
|
|
rooms_df = pl.DataFrame(room_rows)
|
|
objects_df = pl.DataFrame(object_rows)
|
|
|
|
return objects_df, rooms_df
|
|
|
|
|
|
@app.cell
|
|
def display_tables(mo, objects_df, rooms_df):
|
|
mo.vstack(
|
|
[
|
|
mo.md("## Rooms"),
|
|
mo.ui.table(rooms_df),
|
|
mo.md("## Objects"),
|
|
mo.ui.table(objects_df),
|
|
]
|
|
)
|
|
return
|
|
|
|
|
|
@app.cell
|
|
def display_objects_by_room(mo, objects_df, pl):
|
|
"""Table of all objects with their room assignments."""
|
|
_obj_room_df = objects_df.select(
|
|
["obj_id", "room_id", "label", "confidence", "position_confidence"]
|
|
).sort(["room_id", "label"])
|
|
|
|
mo.vstack([
|
|
mo.md("## Objects by Room"),
|
|
mo.ui.table(_obj_room_df),
|
|
])
|
|
return
|
|
|
|
|
|
@app.cell
|
|
def display_room_summary(mo, objects_df, pl, rooms_df):
|
|
"""Table of all rooms with their object counts."""
|
|
_counts = (
|
|
objects_df.group_by("room_id")
|
|
.agg(pl.col("obj_id").count().alias("object_count"))
|
|
)
|
|
_room_summary = rooms_df.select(["room_id"]).join(
|
|
_counts, on="room_id", how="left"
|
|
).with_columns(
|
|
pl.col("object_count").fill_null(0).cast(pl.Int64)
|
|
).sort("room_id")
|
|
|
|
mo.vstack([
|
|
mo.md("## Room Summary"),
|
|
mo.ui.table(_room_summary),
|
|
])
|
|
return
|
|
|
|
|
|
@app.cell
|
|
def upload_query(mo):
|
|
file_upload = mo.ui.file(
|
|
filetypes=["image/*"],
|
|
kind="area",
|
|
label="Upload a query image to find matching objects",
|
|
)
|
|
file_upload
|
|
return (file_upload,)
|
|
|
|
|
|
@app.cell
|
|
def query_matching(
|
|
Image,
|
|
file_upload,
|
|
mo,
|
|
pipeline,
|
|
query_image_against_scene_graph,
|
|
scene_graph,
|
|
text_labels,
|
|
):
|
|
from io import BytesIO
|
|
|
|
query_result = None
|
|
query_cropped = None
|
|
top_matches = []
|
|
|
|
_file_contents = file_upload.contents()
|
|
mo.stop(not _file_contents, mo.md("请先上传文件"))
|
|
|
|
_query_image = Image.open(BytesIO(_file_contents)).convert("RGB")
|
|
|
|
_query_results = query_image_against_scene_graph(
|
|
image=_query_image,
|
|
pipeline=pipeline,
|
|
scene_graph=scene_graph,
|
|
text_labels=text_labels,
|
|
top_k=5,
|
|
batch_size=1,
|
|
)
|
|
|
|
if _query_results:
|
|
_best_result = max(
|
|
_query_results,
|
|
key=lambda result: result.matches[0].score if result.matches else -1,
|
|
)
|
|
query_cropped = _best_result.query_crop
|
|
top_matches = [
|
|
{
|
|
"obj_id": match.obj_id,
|
|
"distance": int(pipeline.hash_bits - match.score),
|
|
"similarity": match.similarity,
|
|
"hash_hex": match.hash_bytes.hex(),
|
|
}
|
|
for match in _best_result.matches
|
|
]
|
|
|
|
query_result = {
|
|
"query_cropped": query_cropped,
|
|
"query_hash_hex": _best_result.query_hash.hex(),
|
|
"top_matches": top_matches,
|
|
"num_query_results": len(_query_results),
|
|
}
|
|
|
|
return query_cropped, query_result, top_matches
|
|
|
|
|
|
@app.cell
|
|
def display_results(mo, object_images, query_cropped, query_result, top_matches):
|
|
mo.stop(not query_result, mo.md("No query results yet. Upload an image above."))
|
|
|
|
_result_items = [
|
|
mo.vstack(
|
|
[
|
|
mo.md("**Query (cropped)**"),
|
|
mo.image(query_cropped),
|
|
],
|
|
align="center",
|
|
)
|
|
]
|
|
|
|
for _match in top_matches:
|
|
_obj_id = _match["obj_id"]
|
|
_obj_img = object_images.get(_obj_id)
|
|
|
|
if _obj_img is not None:
|
|
_result_items.append(
|
|
mo.vstack(
|
|
[
|
|
mo.md(f"**{_obj_id}**"),
|
|
mo.image(_obj_img),
|
|
mo.md(f"Distance: {_match['distance']}"),
|
|
mo.md(f"Similarity: {_match['similarity']:.2%}"),
|
|
],
|
|
align="center",
|
|
)
|
|
)
|
|
|
|
mo.vstack(_result_items, justify="center", gap=2)
|
|
|
|
return
|
|
|
|
|
|
if __name__ == "__main__":
|
|
app.run()
|