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

mini-nav/utils/model.py

@@ -15,13 +15,12 @@ if TYPE_CHECKING:
 def load_sam_model(
     model_name: str = "facebook/sam2.1-hiera-large",
 ) -> MaskGenerationPipeline:
-    device = str(get_device())
+    device = get_device()
-    device_id = 0 if device.startswith("cuda") else -1
 
     return pipeline(
         task="mask-generation",
         model=model_name,
-        device=device_id,
+        device=device,
     )
 
 

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():
+    _view_jobs = [
-        for _view_idx, _rgb in enumerate(_views):
+        (_room_id, _view_idx, _rgb)
-            _rgb3 = _rgb[..., :3] if _rgb.shape[-1] > 3 else _rgb
+        for _room_id, _views in all_room_views.items()
-            _image = Image.fromarray(_rgb3.astype(np.uint8))
+        for _view_idx, _rgb in enumerate(_views)
+    ]
 
-            _masks = segment_image(
+    for _room_id, _view_idx, _rgb in mo.status.progress_bar(
-                hash_pipeline.mask_generator,
+        _view_jobs,
-                _image,
+        title="Extracting masks and hashes",
-                min_area=hash_pipeline.sam_min_mask_area,
+        subtitle="Running SAM + HashPipeline",
-                max_masks=hash_pipeline.sam_max_masks,
+        show_eta=True,
-                points_per_batch=hash_pipeline.sam_points_per_batch,
+        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:
+            _obj_id = f"obj_{_obj_index:04d}"
-                _masked_image = extract_masked_region(_image, _mask["segment"])
+            _obj_index += 1
-                _bits = hash_pipeline(_masked_image)
+            _bits_np = _bits.squeeze().detach().cpu().numpy()
 
-                _bbox = _mask["bbox"]
+            _obj_node = ObjectNode(
-                _obj_center = np.array(
+                obj_id=_obj_id,
-                    [_bbox[0] + _bbox[2] / 2, _bbox[1] + _bbox[3] / 2, 0.0],
+                room_id=_room_id,
-                    dtype=np.float32,
+                position=_obj_center,
-                )
+                visual_hash=_bits_np,
-
+                semantic_hash=_bits_np,
-                _obj_id = f"obj_{_obj_index:04d}"
+                hit_count=1,
-                _obj_index += 1
+                last_seen_frame=0,
-                _bits_np = _bits.squeeze().detach().cpu().numpy()
+            )
-
+            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}")