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_id = 0 if device.startswith("cuda") else -1
+    device = get_device()
 
     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,8 +211,19 @@ def _(
     total_masks = 0
     _obj_index = 0
 
-    for _room_id, _views in all_room_views.items():
-        for _view_idx, _rgb in enumerate(_views):
+    _view_jobs = [
+        (_room_id, _view_idx, _rgb)
+        for _room_id, _views in all_room_views.items()
+        for _view_idx, _rgb in enumerate(_views)
+    ]
+
+    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))