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

@@ -16,6 +16,23 @@ class TopDownSceneElements:
edges: Sequence[tuple[str, str]] = field(default_factory=tuple)
# Default color palette for object labels (cycled if more labels than colors).
_DEFAULT_LABEL_COLORS: dict[str, str] = {
"chair": "#e6194b",
"table": "#3cb44b",
"sofa": "#ffe119",
"cabinet": "#4363d8",
"shelf": "#f58231",
"lamp": "#911eb4",
"picture": "#42d4f4",
"window": "#f032e6",
"door": "#bfef45",
"plant": "#fabed4",
}
_FALLBACK_COLOR = "#808080"
@dataclass(frozen=True)
class TopDownRenderStyle:
room_color: str = "red"
@@ -25,6 +42,18 @@ class TopDownRenderStyle:
figure_size: tuple[int, int] = (8, 8)
map_cmap: str = "gray"
title: str = "RoomNode Top-Down Map"
# Object node style.
object_marker_size: int = 25
object_label_colors: dict[str, str] = field(
default_factory=lambda: dict(_DEFAULT_LABEL_COLORS)
)
object_fallback_color: str = _FALLBACK_COLOR
object_show_label: bool = True
object_label_fontsize: int = 6
# Edge style.
edge_color: str = "cyan"
edge_alpha: float = 0.4
edge_arrow_size: float = 8.0
def render_topdown_scene_map(
@@ -35,13 +64,14 @@ def render_topdown_scene_map(
) -> Image.Image:
"""Render a top-down scene map as a PIL Image.
Uses matplotlib Agg backend (non-interactive) to draw room markers and
labels on the habitat top-down occupancy map, then returns the result
as a PIL Image.
Uses matplotlib Agg backend (non-interactive) to draw room markers,
object markers (colored by label), and edges (arrows from room to
object) on the habitat top-down occupancy map, then returns the
result as a PIL Image.
Args:
pathfinder: Habitat pathfinder instance.
elements: Scene elements containing room nodes to annotate.
elements: Scene elements containing room/object nodes and edges.
meters_per_pixel: Resolution for the top-down map.
style: Visual style configuration.
@@ -50,7 +80,6 @@ def render_topdown_scene_map(
Raises:
ValueError: If room_nodes is empty or meters_per_pixel <= 0.
NotImplementedError: If object_nodes or edges are provided.
"""
if not elements.room_nodes:
raise ValueError("room_nodes must not be empty")
@@ -58,12 +87,6 @@ def render_topdown_scene_map(
if meters_per_pixel <= 0:
raise ValueError("meters_per_pixel must be greater than 0")
if elements.object_nodes:
raise NotImplementedError("object_nodes overlay is not implemented yet")
if elements.edges:
raise NotImplementedError("edge overlay is not implemented yet")
if style is None:
style = TopDownRenderStyle()
@@ -84,7 +107,8 @@ def render_topdown_scene_map(
ax = fig.add_subplot(111)
ax.imshow(top_down_map, cmap=style.map_cmap)
# Annotate room positions and labels.
# --- Room nodes ---
room_grid_positions: dict[str, tuple[int, int]] = {}
for room_node in elements.room_nodes:
grid_y, grid_x = maps.to_grid(
float(room_node.center[2]),
@@ -92,6 +116,7 @@ def render_topdown_scene_map(
top_down_map.shape,
pathfinder=pathfinder,
)
room_grid_positions[room_node.room_id] = (grid_x, grid_y)
ax.scatter(grid_x, grid_y, c=style.room_color, s=style.room_marker_size)
ax.text(
grid_x + style.room_label_offset,
@@ -101,6 +126,74 @@ def render_topdown_scene_map(
fontsize=8,
)
# --- Object nodes (colored by label) ---
obj_grid_positions: dict[str, tuple[int, int]] = {}
drawn_labels: set[str] = set()
for obj_node in elements.object_nodes:
if obj_node.position is None:
continue
grid_y, grid_x = maps.to_grid(
float(obj_node.position[2]),
float(obj_node.position[0]),
top_down_map.shape,
pathfinder=pathfinder,
)
obj_grid_positions[obj_node.obj_id] = (grid_x, grid_y)
label = obj_node.label or "unknown"
color = style.object_label_colors.get(label, style.object_fallback_color)
# Only add to legend once per label.
legend_label = label if label not in drawn_labels else None
drawn_labels.add(label)
ax.scatter(
grid_x,
grid_y,
c=color,
s=style.object_marker_size,
marker="^",
label=legend_label,
zorder=3,
)
if style.object_show_label:
ax.text(
grid_x + 1,
grid_y + 1,
label,
color=color,
fontsize=style.object_label_fontsize,
)
# --- Edges (arrows from room → object) ---
for room_id, obj_id in elements.edges:
room_pos = room_grid_positions.get(room_id)
obj_pos = obj_grid_positions.get(obj_id)
if room_pos is None or obj_pos is None:
continue
ax.annotate(
"",
xy=obj_pos,
xytext=room_pos,
arrowprops=dict(
arrowstyle="->",
color=style.edge_color,
alpha=style.edge_alpha,
lw=0.8,
),
)
# --- Legend (only if object labels were drawn) ---
if drawn_labels:
ax.legend(
loc="upper right",
fontsize=7,
markerscale=0.8,
framealpha=0.7,
)
ax.set_title(style.title)
ax.axis("off")
fig.canvas.draw()