feat(scenegraph): add image storage and verification UI

notebooks/verification.py

@@ -8,12 +8,14 @@
 
 import marimo
 
-__generated_with = "0.21.1"
+__generated_with = "0.20.4"
 app = marimo.App(width="medium", app_title="Pipeline Verification")
 
 
 @app.cell
 def import_packages():
+    from io import BytesIO
+
     import habitat_sim
     import marimo as mo
     import numpy as np
@@ -27,6 +29,7 @@ def import_packages():
     from utils.image import extract_masked_region, segment_image
 
     return (
+        BytesIO,
         HashPipeline,
         Image,
         ObjectNode,
@@ -44,7 +47,7 @@ def import_packages():
 
 
 @app.cell
-def _(habitat_sim):
+def setup_habitat_simulator(habitat_sim):
     scene_path = "data/scene_datasets/habitat-test-scenes/skokloster-castle.glb"
     num_rooms = 4
     views_per_room = 6
@@ -96,7 +99,7 @@ def _(habitat_sim):
 
 
 @app.cell
-def _(RoomNode, np, num_rooms, sim):
+def sample_room_nodes(RoomNode, np, num_rooms, sim):
     room_nodes = []
     for _idx in range(num_rooms):
         _point = sim.pathfinder.get_random_navigable_point()
@@ -114,7 +117,7 @@ def _(RoomNode, np, num_rooms, sim):
 
 
 @app.cell
-def _(maps, meters_per_pixel, plt, room_nodes, sim):
+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]),
@@ -141,7 +144,15 @@ def _(maps, meters_per_pixel, plt, room_nodes, sim):
 
 
 @app.cell
-def _(agent, habitat_sim, mo, plt, room_nodes, sim, views_per_room):
+def collect_room_views(
+    agent,
+    habitat_sim,
+    mo,
+    plt,
+    room_nodes,
+    sim,
+    views_per_room,
+):
     all_room_views = {}
 
     for _node in mo.status.progress_bar(
@@ -181,7 +192,7 @@ def _(agent, habitat_sim, mo, plt, room_nodes, sim, views_per_room):
 
 
 @app.cell
-def _(HashPipeline, hash_bits, sam_max_masks, sam_min_area):
+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",
@@ -193,7 +204,7 @@ def _(HashPipeline, hash_bits, sam_max_masks, sam_min_area):
 
 
 @app.cell
-def _(
+def build_scene_graph_from_views(
     Image,
     ObjectNode,
     SimpleSceneGraph,
@@ -256,6 +267,7 @@ def _(
                 position=_obj_center,
                 visual_hash=_bits_np,
                 semantic_hash=_bits_np,
+                image_bytes=np.array(_masked_image).tobytes(),
                 hit_count=1,
                 last_seen_frame=0,
             )
@@ -267,7 +279,7 @@ def _(
 
 
 @app.cell
-def _(pl, scene_graph):
+def build_room_and_object_tables(pl, scene_graph):
     _room_rows = []
     for _room in scene_graph.rooms.values():
         _room_rows.append(
@@ -300,17 +312,42 @@ def _(pl, scene_graph):
     return objects_table, rooms_table
 
 
-@app.cell
-def _(rooms_table):
+@app.cell(disabled=True)
+def display_rooms_table(rooms_table):
     rooms_table
     return
 
 
-@app.cell
-def _(objects_table):
+@app.cell(disabled=True)
+def display_objects_table(objects_table):
     objects_table
     return
 
 
+@app.cell
+def create_file_upload(mo):
+    file_upload = mo.ui.file(
+        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,)
+
+
+@app.cell
+def display_uploaded_image(mo, np, uploaded_image):
+    mo.image(np.array(uploaded_image), alt="Uploaded query image")
+    return
+
+
 if __name__ == "__main__":
     app.run()