refactor(pipeline): integrate SAM segmentation and modularize model loading

notebooks/verification.py

@@ -1,3 +1,11 @@
+# /// script
+# requires-python = ">=3.13"
+# dependencies = [
+#     "marimo>=0.21.1",
+#     "pyzmq>=27.1.0",
+# ]
+# ///
+
 import marimo
 
 __generated_with = "0.21.1"
@@ -8,11 +16,26 @@ app = marimo.App()
 def _():
     import habitat_sim
     import numpy as np
+    import polars as pl
     from habitat.utils.visualizations import maps
     from matplotlib import pyplot as plt
-    from scenegraph import RoomNode
+    from PIL import Image
+    from compressors.pipeline import HashPipeline
+    from scenegraph import ObjectNode, RoomNode, SimpleSceneGraph
+    from utils.common import get_device
+    from utils.image import extract_masked_region, segment_image
 
-    return RoomNode, habitat_sim, maps, np, plt
+    return (
+        HashPipeline,
+        Image,
+        RoomNode,
+        SimpleSceneGraph,
+        habitat_sim,
+        maps,
+        np,
+        pl,
+        plt,
+    )
 
 
 @app.cell
@@ -51,7 +74,20 @@ def _(habitat_sim):
     cfg = habitat_sim.Configuration(sim_cfg, [agent_cfg])
     sim = habitat_sim.Simulator(cfg)
     agent = sim.initialize_agent(0)
-    return agent, meters_per_pixel, num_rooms, sim, views_per_room
+    sam_max_masks = 5
+    sam_min_area = 32 * 32
+    sam_points_per_batch = 64
+    hash_bits = 512
+    return (
+        agent,
+        hash_bits,
+        meters_per_pixel,
+        num_rooms,
+        sam_max_masks,
+        sam_min_area,
+        sim,
+        views_per_room,
+    )
 
 
 @app.cell
@@ -66,7 +102,6 @@ def _(RoomNode, np, num_rooms, sim):
         )
         room_nodes.append(_room_node)
 
-    room_points = np.vstack([_node.center for _node in room_nodes])
     print("Sampled room centers:")
     for _node in room_nodes:
         print(_node.room_id, _node.center)
@@ -125,6 +160,83 @@ def _(agent, habitat_sim, plt, room_nodes, sim, views_per_room):
             _ax.axis("off")
         plt.tight_layout()
         plt.show()
+    return (all_room_views,)
+
+
+@app.cell
+def _(HashPipeline, hash_bits, sam_max_masks, sam_min_area):
+    hash_pipeline = HashPipeline(
+        dino_model="facebook/dinov2-large",
+        sam_model="facebook/sam2.1-hiera-large",
+        sam_min_mask_area=sam_min_area,
+        sam_max_masks=sam_max_masks,
+        hash_bits=hash_bits,
+    )
+    return
+
+
+@app.cell
+def _(Image, SimpleSceneGraph, all_room_views, np, room_nodes):
+    scene_graph = SimpleSceneGraph(
+        rooms={_room.room_id: _room for _room in room_nodes}, objects={}
+    )
+    total_masks = 0
+    _obj_index = 0
+
+    for _room_id, _views in all_room_views.items():
+        for _view_idx, _rgb in enumerate(_views):
+            _rgb3 = _rgb[..., :3] if _rgb.shape[-1] > 3 else _rgb
+            _image = Image.fromarray(_rgb3.astype(np.uint8))
+
+
+    print(f"Total objects created: {len(scene_graph.objects)}")
+    print(f"Total processed masks: {total_masks}")
+    return (scene_graph,)
+
+
+@app.cell
+def _(pl, scene_graph):
+    _room_rows = []
+    for _room in scene_graph.rooms.values():
+        _room_rows.append(
+            {
+                "room_id": _room.room_id,
+                "center_x": float(_room.center[0]),
+                "center_y": float(_room.center[1]),
+                "center_z": float(_room.center[2]),
+                "bbox_dx": float(_room.bbox_extent[0]),
+                "bbox_dy": float(_room.bbox_extent[1]),
+                "bbox_dz": float(_room.bbox_extent[2]),
+            }
+        )
+
+    _object_rows = []
+    for _obj in scene_graph.objects.values():
+        _object_rows.append(
+            {
+                "obj_id": _obj.obj_id,
+                "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(),
+            }
+        )
+
+    rooms_table = pl.DataFrame(_room_rows)
+    objects_table = pl.DataFrame(_object_rows)
+    return objects_table, rooms_table
+
+
+@app.cell
+def _(rooms_table):
+    rooms_table
+    return
+
+
+@app.cell
+def _(objects_table):
+    objects_table
     return