SikongJueluo/Mini-Nav
1
0
Fork 0
mirror of https://github.com/SikongJueluo/Mini-Nav.git synced 2026-07-13 04:25:32 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
SikongJueluo
2026-03-26 20:08:05 +08:00
parent 968819e113
commit 1c9752cf9e
2 changed files with 66 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
mini-nav/utils/model.py
View File

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

38
notebooks/verification.py
View File

@@ -9,12 +9,13 @@
import marimo import marimo
__generated_with = "0.21.1" __generated_with = "0.21.1"
app = marimo.App(app_title="Pipeline Verification") app = marimo.App(width="medium", app_title="Pipeline Verification")
@app.cell @app.cell
def import_packages(): def import_packages():
import habitat_sim import habitat_sim
import marimo as mo
import numpy as np import numpy as np
import polars as pl import polars as pl
from habitat.utils.visualizations import maps from habitat.utils.visualizations import maps
@@ -34,6 +35,7 @@ def import_packages():
extract_masked_region, extract_masked_region,
habitat_sim, habitat_sim,
maps, maps,
mo,
np, np,
pl, pl,
plt, plt,
@@ -139,16 +141,28 @@ def _(maps, meters_per_pixel, plt, room_nodes, sim):
@app.cell @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 = {} 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 = habitat_sim.AgentState()
_agent_state.position = _node.center.copy() _agent_state.position = _node.center.copy()
agent.set_state(_agent_state) agent.set_state(_agent_state)
_room_views = [] _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() _observations = sim.get_sensor_observations()
_rgb = _observations["color_sensor"] _rgb = _observations["color_sensor"]
_room_views.append(_rgb) _room_views.append(_rgb)
@@ -186,6 +200,7 @@ def _(
all_room_views, all_room_views,
extract_masked_region, extract_masked_region,
hash_pipeline, hash_pipeline,
mo,
np, np,
room_nodes, room_nodes,
segment_image, segment_image,
@@ -196,8 +211,19 @@ def _(
total_masks = 0 total_masks = 0
_obj_index = 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)
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 _rgb3 = _rgb[..., :3] if _rgb.shape[-1] > 3 else _rgb
_image = Image.fromarray(_rgb3.astype(np.uint8)) _image = Image.fromarray(_rgb3.astype(np.uint8))