refactor(compressors): consolidate pipeline and improve mask handling

notebooks/verification.py

@@ -9,11 +9,11 @@
 import marimo
 
 __generated_with = "0.21.1"
-app = marimo.App()
+app = marimo.App(app_title="Pipeline Verification")
 
 
 @app.cell
-def _():
+def import_packages():
     import habitat_sim
     import numpy as np
     import polars as pl
@@ -28,13 +28,16 @@ def _():
     return (
         HashPipeline,
         Image,
+        ObjectNode,
         RoomNode,
         SimpleSceneGraph,
+        extract_masked_region,
         habitat_sim,
         maps,
         np,
         pl,
         plt,
+        segment_image,
     )
 
 
@@ -172,11 +175,21 @@ def _(HashPipeline, hash_bits, sam_max_masks, sam_min_area):
         sam_max_masks=sam_max_masks,
         hash_bits=hash_bits,
     )
-    return
+    return (hash_pipeline,)
 
 
 @app.cell
-def _(Image, SimpleSceneGraph, all_room_views, np, room_nodes):
+def _(
+    Image,
+    ObjectNode,
+    SimpleSceneGraph,
+    all_room_views,
+    extract_masked_region,
+    hash_pipeline,
+    np,
+    room_nodes,
+    segment_image,
+):
     scene_graph = SimpleSceneGraph(
         rooms={_room.room_id: _room for _room in room_nodes}, objects={}
     )
@@ -188,6 +201,39 @@ def _(Image, SimpleSceneGraph, all_room_views, np, room_nodes):
             _rgb3 = _rgb[..., :3] if _rgb.shape[-1] > 3 else _rgb
             _image = Image.fromarray(_rgb3.astype(np.uint8))
 
+            _masks = segment_image(
+                hash_pipeline.mask_generator,
+                _image,
+                min_area=hash_pipeline.sam_min_mask_area,
+                max_masks=hash_pipeline.sam_max_masks,
+                points_per_batch=hash_pipeline.sam_points_per_batch,
+            )
+            total_masks += len(_masks)
+
+            for _mask in _masks:
+                _masked_image = extract_masked_region(_image, _mask["segment"])
+                _bits = hash_pipeline(_masked_image)
+
+                _bbox = _mask["bbox"]
+                _obj_center = np.array(
+                    [_bbox[0] + _bbox[2] / 2, _bbox[1] + _bbox[3] / 2, 0.0],
+                    dtype=np.float32,
+                )
+
+                _obj_id = f"obj_{_obj_index:04d}"
+                _obj_index += 1
+                _bits_np = _bits.squeeze().detach().cpu().numpy()
+
+                _obj_node = 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] = _obj_node
 
     print(f"Total objects created: {len(scene_graph.objects)}")
     print(f"Total processed masks: {total_masks}")