diff --git a/mini-nav/utils/model.py b/mini-nav/utils/model.py index 711f4ee..e282292 100644 --- a/mini-nav/utils/model.py +++ b/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, ) diff --git a/notebooks/verification.py b/notebooks/verification.py index 0085ff8..a200a58 100644 --- a/notebooks/verification.py +++ b/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}")