feat(pipeline): add batch processing for scene graph construction

notebooks/verification.py

@@ -19,6 +19,8 @@ def import_packages():
     import marimo as mo
     import numpy as np
     import polars as pl
+    from habitat.utils.visualizations import maps
+    from matplotlib import pyplot as plt
     from PIL import Image
 
     from compressors.pipeline import HashPipeline
@@ -33,7 +35,6 @@ def import_packages():
     from utils.image import extract_masked_region, segment_image
 
     return (
-        BytesIO,
         HabitatSimulatorConfig,
         HashPipeline,
         Image,
@@ -44,9 +45,11 @@ def import_packages():
         collect_room_views_by_room,
         create_habitat_simulator,
         extract_masked_region,
+        maps,
         mo,
         np,
         pl,
+        plt,
         render_topdown_scene_map,
         segment_image,
     )
@@ -54,7 +57,10 @@ def import_packages():
 
 @app.cell
 def setup_verification_context(
-    HabitatSimulatorConfig, RoomNode, create_habitat_simulator, np
+    HabitatSimulatorConfig,
+    RoomNode,
+    create_habitat_simulator,
+    np,
 ):
     scene_path = "data/scene_datasets/habitat-test-scenes/skokloster-castle.glb"
     image_size = 256
@@ -65,6 +71,7 @@ def setup_verification_context(
     sam_max_masks = 5
     sam_min_area = 32 * 32
     hash_bits = 512
+    pipeline_batch_size = 64
 
     sim, agent = create_habitat_simulator(
         HabitatSimulatorConfig(
@@ -91,11 +98,11 @@ def setup_verification_context(
     print("Sampled room centers:")
     for node in room_nodes:
         print(node.room_id, node.center)
-
     return (
         agent,
         hash_bits,
         meters_per_pixel,
+        pipeline_batch_size,
         room_nodes,
         sam_max_masks,
         sam_min_area,
@@ -107,7 +114,9 @@ def setup_verification_context(
 @app.cell
 def render_topdown_room_map(
     TopDownSceneElements,
+    maps,
     meters_per_pixel,
+    plt,
     render_topdown_scene_map,
     room_nodes,
     sim,
@@ -116,22 +125,25 @@ def render_topdown_room_map(
         pathfinder=sim.pathfinder,
         elements=TopDownSceneElements(room_nodes=room_nodes),
         meters_per_pixel=meters_per_pixel,
+        maps_module=maps,
+        plt_module=plt,
     )
     return
 
 
 @app.cell
 def build_scene_graph_pipeline(
-    agent,
     HashPipeline,
     Image,
     ObjectNode,
     SimpleSceneGraph,
+    agent,
     collect_room_views_by_room,
     extract_masked_region,
     hash_bits,
     mo,
     np,
+    pipeline_batch_size,
     room_nodes,
     sam_max_masks,
     sam_min_area,
@@ -161,16 +173,17 @@ def build_scene_graph_pipeline(
     total_masks = 0
     object_index = 0
 
-    view_jobs = [
+    room_view_dataset = [
         (room_id, view_idx, rgb)
         for room_id, views in all_room_views.items()
         for view_idx, rgb in enumerate(views)
     ]
+    object_dataset = []
 
     for room_id, _view_idx, rgb in mo.status.progress_bar(
-        view_jobs,
-        title="Extracting masks and hashes",
-        subtitle="Running SAM + HashPipeline",
+        room_view_dataset,
+        title="Building object dataset",
+        subtitle="Running SAM segmentation",
         show_eta=True,
         show_rate=True,
     ):
@@ -188,31 +201,56 @@ def build_scene_graph_pipeline(
 
         for mask in masks:
             masked_image = extract_masked_region(image, mask["segment"])
-            bits = hash_pipeline(masked_image)
+            object_dataset.append((room_id, mask["bbox"], masked_image))
 
-            bbox = mask["bbox"]
-            obj_center = np.array(
-                [bbox[0] + bbox[2] / 2, bbox[1] + bbox[3] / 2, 0.0],
-                dtype=np.float32,
+    if object_dataset:
+        masked_images = [item[2] for item in object_dataset]
+        batched_bits = hash_pipeline.forward_dataset(
+            masked_images,
+            batch_size=pipeline_batch_size,
+            apply_sam=False,
+        )
+        if len(batched_bits) != len(object_dataset):
+            raise RuntimeError(
+                "Batch output size mismatch between masked images and hash outputs."
+            )
+    else:
+        batched_bits = []
+
+    for ob_idx, (room_id, bbox, _) in enumerate(object_dataset):
+        bits = batched_bits[ob_idx]
+        obj_center = np.array(
+            [bbox[0] + bbox[2] / 2, bbox[1] + bbox[3] / 2, 0.0],
+            dtype=np.float32,
+        )
+
+        obj_id = f"obj_{object_index:04d}"
+        object_index += 1
+
+        bits_array = np.asarray(bits.detach().cpu().numpy()).reshape(-1)
+        if bits_array.size == 512:
+            bits_binary = (bits_array > 0).astype(np.uint8)
+            hash_bytes = np.packbits(bits_binary).tobytes()
+        elif bits_array.size == 64:
+            hash_bytes = bits_array.astype(np.uint8).tobytes()
+        else:
+            raise ValueError(
+                f"Unexpected hash length: {bits_array.size}. Expected 512 bits or 64 bytes."
             )
 
-            obj_id = f"obj_{object_index:04d}"
-            object_index += 1
-            bits_np = bits.squeeze().detach().cpu().numpy()
-
-            scene_graph.objects[obj_id] = ObjectNode(
-                obj_id=obj_id,
-                room_id=room_id,
-                position=obj_center,
-                visual_hash=bits_np,
-                semantic_hash=bits_np,
-                hit_count=1,
-                last_seen_frame=0,
-            )
+        scene_graph.objects[obj_id] = ObjectNode(
+            obj_id=obj_id,
+            room_id=room_id,
+            position=obj_center,
+            visual_hash=hash_bytes,
+            semantic_hash=hash_bytes,
+            hit_count=1,
+            last_seen_frame=0,
+        )
 
     print(f"Total objects created: {len(scene_graph.objects)}")
     print(f"Total processed masks: {total_masks}")
-    return all_room_views, hash_pipeline, scene_graph
+    return (scene_graph,)
 
 
 @app.cell
@@ -236,8 +274,8 @@ def build_room_and_object_tables(pl, scene_graph):
             "room_id": obj.room_id,
             "last_seen_frame": int(obj.last_seen_frame),
             "hit_count": int(obj.hit_count),
-            "visual_hash": obj.visual_hash.tolist(),
-            "semantic_hash": obj.semantic_hash.tolist(),
+            "visual_hash": obj.visual_hash.hex(),
+            "semantic_hash": obj.semantic_hash.hex(),
         }
         for obj in scene_graph.objects.values()
     ]
@@ -248,22 +286,25 @@ def build_room_and_object_tables(pl, scene_graph):
 
 
 @app.cell
-def upload_query_image(BytesIO, Image, mo, np):
+def upload_query_image(mo):
     file_upload = mo.ui.file(
         filetypes=["image/*"],
         kind="area",
         label="Upload a query image",
     )
     file_upload
+    return (file_upload,)
 
-    uploaded_image = None
+
+@app.cell
+def _(file_upload, mo):
+    upload_image = None
     if file_upload.value:
-        contents = file_upload.contents()
-        if contents:
-            uploaded_image = Image.open(BytesIO(contents))
-            mo.image(np.array(uploaded_image), alt="Uploaded query image")
+        upload_image = mo.image(file_upload.contents(), alt="Uploaded query image")
 
-    return file_upload, uploaded_image
+    # Build a grid.
+    upload_image if upload_image is not None else mo.md("No image uploaded yet.")
+    return
 
 
 if __name__ == "__main__":