feat(scenegraph): add JSON serialization and implement top-down object/edge rendering

- 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
This commit is contained in:
2026-05-31 19:45:10 +08:00
parent 9eb52f8cef
commit e5b764520c
4 changed files with 395 additions and 81 deletions

View File

@@ -38,10 +38,13 @@ def project_imports():
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,
@@ -58,16 +61,19 @@ def project_imports():
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,
)
@@ -194,6 +200,7 @@ def build_scene_graph(
SceneGraphBuilder,
cfg_manager,
habitat_config,
load_scene_graph,
mo,
numpy_to_pil,
pipeline,
@@ -201,92 +208,145 @@ def build_scene_graph(
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"
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,
# 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,
),
)
object_images = build_artifacts.object_images
debug_meta = build_artifacts.debug_meta
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
]
# 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)
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,
)
# 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,
)
object_images = build_artifacts.object_images
debug_meta = build_artifacts.debug_meta
from collections import Counter
# Save scene graph to cache.
save_scene_graph(cache_path, scene_graph)
print(f"Saved scene graph to {cache_path}")
_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}")
# 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"Saved cropped images to: {output_dir}")
print(f"Fallback frames: {fallback_count}/{len(debug_meta)}")
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 = [
@@ -335,6 +395,40 @@ def display_tables(mo, objects_df, rooms_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(