refactor(verification): add progress bars and simplify device handling

notebooks/verification.py

@@ -9,12 +9,13 @@
 import marimo
 
 __generated_with = "0.21.1"
-app = marimo.App(app_title="Pipeline Verification")
+app = marimo.App(width="medium", app_title="Pipeline Verification")
 
 
 @app.cell
 def import_packages():
     import habitat_sim
+    import marimo as mo
     import numpy as np
     import polars as pl
     from habitat.utils.visualizations import maps
@@ -34,6 +35,7 @@ def import_packages():
         extract_masked_region,
         habitat_sim,
         maps,
+        mo,
         np,
         pl,
         plt,
@@ -139,16 +141,28 @@ def _(maps, meters_per_pixel, plt, room_nodes, sim):
 
 
 @app.cell
-def _(agent, habitat_sim, plt, room_nodes, sim, views_per_room):
+def _(agent, habitat_sim, mo, plt, room_nodes, sim, views_per_room):
     all_room_views = {}
 
-    for _node in room_nodes:
+    for _node in mo.status.progress_bar(
+        room_nodes,
+        title="Collecting room views",
+        subtitle="Sampling observations from Habitat",
+        show_eta=True,
+        show_rate=True,
+    ):
         _agent_state = habitat_sim.AgentState()
         _agent_state.position = _node.center.copy()
         agent.set_state(_agent_state)
 
         _room_views = []
-        for _ in range(views_per_room):
+        for _ in mo.status.progress_bar(
+            range(views_per_room),
+            title=f"Capturing {_node.room_id}",
+            subtitle="Rotating agent viewpoints",
+            show_eta=True,
+            show_rate=True,
+        ):
             _observations = sim.get_sensor_observations()
             _rgb = _observations["color_sensor"]
             _room_views.append(_rgb)
@@ -186,6 +200,7 @@ def _(
     all_room_views,
     extract_masked_region,
     hash_pipeline,
+    mo,
     np,
     room_nodes,
     segment_image,
@@ -196,44 +211,55 @@ def _(
     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))
+    _view_jobs = [
+        (_room_id, _view_idx, _rgb)
+        for _room_id, _views in all_room_views.items()
+        for _view_idx, _rgb in enumerate(_views)
+    ]
 
-            _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,
+    for _room_id, _view_idx, _rgb in mo.status.progress_bar(
+        _view_jobs,
+        title="Extracting masks and hashes",
+        subtitle="Running SAM + HashPipeline",
+        show_eta=True,
+        show_rate=True,
+    ):
+        _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,
             )
-            total_masks += len(_masks)
 
-            for _mask in _masks:
-                _masked_image = extract_masked_region(_image, _mask["segment"])
-                _bits = hash_pipeline(_masked_image)
+            _obj_id = f"obj_{_obj_index:04d}"
+            _obj_index += 1
+            _bits_np = _bits.squeeze().detach().cpu().numpy()
 
-                _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
+            _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}")