diff --git a/mini-nav/scenegraph/__init__.py b/mini-nav/scenegraph/__init__.py index 987f667..71151b1 100644 --- a/mini-nav/scenegraph/__init__.py +++ b/mini-nav/scenegraph/__init__.py @@ -21,6 +21,7 @@ from .query import ( ) from .roomnode import RoomNode from .scenegraph import SceneGraphMatch, SimpleSceneGraph +from .serialization import load_scene_graph, save_scene_graph from .software_cam import CamMatch, SoftwareCamIndex, xnor_popcount_score __all__ = [ @@ -40,6 +41,8 @@ __all__ = [ "cam_row_to_hash_bytes", "hash_bytes_to_bits_array", "hash_bytes_to_cam_row", + "load_scene_graph", "query_image_against_scene_graph", + "save_scene_graph", "xnor_popcount_score", ] diff --git a/mini-nav/scenegraph/serialization.py b/mini-nav/scenegraph/serialization.py new file mode 100644 index 0000000..2329c98 --- /dev/null +++ b/mini-nav/scenegraph/serialization.py @@ -0,0 +1,124 @@ +"""JSON serialization for SimpleSceneGraph. + +Provides save_scene_graph() and load_scene_graph() for persisting +scene graphs to/from JSON files. Handles numpy arrays, bytes (hashes), +and optional fields. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import numpy as np + +from .objectnode import ObjectNode +from .roomnode import RoomNode +from .scenegraph import SimpleSceneGraph + + +def _ndarray_to_list(arr: np.ndarray) -> list[float]: + return arr.tolist() + + +def _bytes_to_hex(b: bytes) -> str: + return b.hex() + + +def _room_node_to_dict(node: RoomNode) -> dict[str, Any]: + return { + "room_id": node.room_id, + "center": _ndarray_to_list(node.center), + "bbox_extent": _ndarray_to_list(node.bbox_extent), + } + + +def _object_node_to_dict(node: ObjectNode) -> dict[str, Any]: + return { + "obj_id": node.obj_id, + "room_id": node.room_id, + "position": _ndarray_to_list(node.position), + "visual_hash": _bytes_to_hex(node.visual_hash), + "semantic_hash": _bytes_to_hex(node.semantic_hash), + "hit_count": node.hit_count, + "last_seen_frame": node.last_seen_frame, + "label": node.label, + "confidence": node.confidence, + "bbox_xyxy": list(node.bbox_xyxy) if node.bbox_xyxy is not None else None, + "source_view_id": node.source_view_id, + "position_confidence": node.position_confidence, + } + + +def _room_node_from_dict(d: dict[str, Any]) -> RoomNode: + return RoomNode( + room_id=d["room_id"], + center=np.asarray(d["center"], dtype=np.float32), + bbox_extent=np.asarray(d["bbox_extent"], dtype=np.float32), + ) + + +def _object_node_from_dict(d: dict[str, Any]) -> ObjectNode: + bbox = d.get("bbox_xyxy") + return ObjectNode( + obj_id=d["obj_id"], + room_id=d["room_id"], + position=np.asarray(d["position"], dtype=np.float32), + visual_hash=bytes.fromhex(d["visual_hash"]), + semantic_hash=bytes.fromhex(d["semantic_hash"]), + hit_count=d["hit_count"], + last_seen_frame=d["last_seen_frame"], + label=d.get("label"), + confidence=d.get("confidence"), + bbox_xyxy=tuple(bbox) if bbox is not None else None, + source_view_id=d.get("source_view_id"), + position_confidence=d.get("position_confidence"), + ) + + +def save_scene_graph(path: str | Path, graph: SimpleSceneGraph) -> None: + """Save a SimpleSceneGraph to a JSON file. + + Args: + path: Output file path. + graph: The scene graph to serialize. + """ + data = { + "version": 1, + "rooms": { + room_id: _room_node_to_dict(node) + for room_id, node in graph.rooms.items() + }, + "objects": { + obj_id: _object_node_to_dict(node) + for obj_id, node in graph.objects.items() + }, + } + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, "w", encoding="utf-8") as f: + json.dump(data, f, indent=2, ensure_ascii=False) + + +def load_scene_graph(path: str | Path) -> SimpleSceneGraph: + """Load a SimpleSceneGraph from a JSON file. + + Args: + path: Input file path. + + Returns: + Deserialized SimpleSceneGraph. + """ + with open(path, "r", encoding="utf-8") as f: + data = json.load(f) + + rooms = { + room_id: _room_node_from_dict(node_dict) + for room_id, node_dict in data["rooms"].items() + } + objects = { + obj_id: _object_node_from_dict(node_dict) + for obj_id, node_dict in data["objects"].items() + } + return SimpleSceneGraph(rooms=rooms, objects=objects) diff --git a/mini-nav/simulator/topdown.py b/mini-nav/simulator/topdown.py index f0a08bb..50e1503 100644 --- a/mini-nav/simulator/topdown.py +++ b/mini-nav/simulator/topdown.py @@ -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() diff --git a/notebooks/verification.py b/notebooks/verification.py index b601734..8314a8e 100644 --- a/notebooks/verification.py +++ b/notebooks/verification.py @@ -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(