mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-13 04:25:32 +08:00
SikongJueluo
4918b654e7
- Refactor model_loader.py: improve return type annotations from tuple[Any, Any] to tuple[AutoImageProcessor, AutoModel] - Refactor proposal/core.py: add input validation for mask array dimensionality, handle 2D masks and batch dimensions gracefully - Refactor proposal_segament.ipynb: replace inline model loading with centralized load_owlv2_model() and load_sam_model() functions, use batched detect_objects_batch() and generate_proposals_batch() APIs
256 lines
7.3 KiB
Python
256 lines
7.3 KiB
Python
import marimo
|
|
|
|
__generated_with = "0.21.1"
|
|
app = marimo.App()
|
|
|
|
|
|
@app.cell
|
|
def _():
|
|
import marimo as mo
|
|
import numpy as np
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
|
|
return Image, ImageDraw, ImageFont, mo, np
|
|
|
|
|
|
@app.cell
|
|
def _(np):
|
|
from scenegraph import RoomNode
|
|
from simulator import (
|
|
HabitatSimulatorConfig,
|
|
TopDownSceneElements,
|
|
create_habitat_simulator,
|
|
render_topdown_scene_map,
|
|
)
|
|
from habitat.utils.visualizations import maps
|
|
|
|
scene_path = "data/scene_datasets/habitat-test-scenes/skokloster-castle.glb"
|
|
image_size = 768
|
|
num_rooms = 3
|
|
views_per_room = 6
|
|
meters_per_pixel = 0.05
|
|
|
|
sim, agent = create_habitat_simulator(
|
|
HabitatSimulatorConfig(
|
|
scene_path=scene_path,
|
|
views_per_room=views_per_room,
|
|
image_size=image_size,
|
|
sensor_height=1.5,
|
|
move_forward_step=0.25,
|
|
enable_physics=False,
|
|
)
|
|
)
|
|
|
|
room_nodes = []
|
|
for _idx in range(num_rooms):
|
|
point = sim.pathfinder.get_random_navigable_point()
|
|
room_nodes.append(
|
|
RoomNode(
|
|
room_id=f"room_{_idx:02d}",
|
|
center=np.asarray(point, dtype=np.float32),
|
|
bbox_extent=np.asarray([1.5, 2.0, 1.5], dtype=np.float32),
|
|
)
|
|
)
|
|
|
|
render_topdown_scene_map(
|
|
pathfinder=sim.pathfinder,
|
|
elements=TopDownSceneElements(room_nodes=room_nodes),
|
|
meters_per_pixel=meters_per_pixel,
|
|
maps_module=maps,
|
|
)
|
|
return agent, room_nodes, sim, views_per_room
|
|
|
|
|
|
@app.cell
|
|
def _(agent, room_nodes, sim, views_per_room):
|
|
from simulator import collect_room_views_by_room
|
|
|
|
all_room_views = collect_room_views_by_room(
|
|
agent=agent,
|
|
sim=sim,
|
|
room_nodes=room_nodes,
|
|
views_per_room=views_per_room,
|
|
)
|
|
return (all_room_views,)
|
|
|
|
|
|
@app.cell
|
|
def _(ImageDraw, ImageFont, all_room_views, mo):
|
|
from compressors.model_loader import load_owlv2_model
|
|
from compressors.proposal.core import detect_objects_batch
|
|
from utils.common import get_device
|
|
from utils.image import numpy_to_pil
|
|
|
|
device = get_device()
|
|
|
|
score_threshold = 0.25
|
|
postprocess_threshold = 0.1
|
|
text_labels = [
|
|
"a cat",
|
|
"a dog",
|
|
"a TV remote control",
|
|
"a chair",
|
|
"a table",
|
|
"a vase",
|
|
"a painting",
|
|
"a window",
|
|
]
|
|
|
|
owl_processor, owl_model = load_owlv2_model(
|
|
model_name="google/owlv2-base-patch16-ensemble"
|
|
)
|
|
|
|
image = numpy_to_pil(all_room_views["room_00"][1])
|
|
|
|
detection_batch = detect_objects_batch(
|
|
model=owl_model,
|
|
processor=owl_processor,
|
|
images=[image],
|
|
text_labels_per_image=[text_labels],
|
|
score_threshold=score_threshold,
|
|
postprocess_threshold=postprocess_threshold,
|
|
)
|
|
detections = detection_batch[0] if detection_batch else []
|
|
filtered_items = [(det["bbox"], det["score"], det["label"]) for det in detections]
|
|
|
|
_vis_image = image.copy()
|
|
_draw = ImageDraw.Draw(_vis_image)
|
|
|
|
try:
|
|
font = ImageFont.truetype("arial.ttf", 18)
|
|
except Exception:
|
|
font = ImageFont.load_default()
|
|
|
|
for _box, _score, _text_label in filtered_items:
|
|
_x1, _y1, _x2, _y2 = [float(v) for v in _box]
|
|
_label = f"{_text_label}: {_score:.3f}"
|
|
|
|
_draw.rectangle([_x1, _y1, _x2, _y2], outline="red", width=3)
|
|
|
|
try:
|
|
_tx1, _ty1, _tx2, _ty2 = _draw.textbbox((_x1, _y1), _label, font=font)
|
|
except Exception:
|
|
_w = _draw.textlength(_label, font=font)
|
|
_h = 20
|
|
_tx1, _ty1, _tx2, _ty2 = _x1, _y1, _x1 + _w + 6, _y1 + _h
|
|
|
|
text_bg = [_tx1, max(0, _ty1 - 2), _tx2 + 4, _ty2 + 2]
|
|
_draw.rectangle(text_bg, fill="red")
|
|
_draw.text((_x1 + 2, max(0, _y1)), _label, fill="white", font=font)
|
|
|
|
# 8. 结果文本
|
|
detection_lines = []
|
|
for _box, _score, _text_label in filtered_items:
|
|
box_rounded = [round(_v, 2) for _v in _box]
|
|
detection_lines.append(
|
|
f"- {_text_label}: score={_score:.3f}, box={box_rounded}"
|
|
)
|
|
|
|
if not detection_lines:
|
|
detection_text = f"没有检测到 score >= {score_threshold:.2f} 的目标"
|
|
else:
|
|
detection_text = "\n".join(detection_lines)
|
|
|
|
mo.vstack(
|
|
[
|
|
mo.md(
|
|
"## OWLv2 检测可视化结果"
|
|
f"\n\ndevice: `{device}`"
|
|
f"\n\n过滤阈值：`score >= {score_threshold:.2f}`"
|
|
),
|
|
mo.image(_vis_image, width=700),
|
|
mo.md(detection_text),
|
|
]
|
|
)
|
|
return device, filtered_items, image
|
|
|
|
|
|
@app.cell
|
|
def _(Image, ImageDraw, device, filtered_items, image, mo, np):
|
|
from compressors.model_loader import load_sam_model
|
|
from compressors.proposal.core import generate_proposals_batch
|
|
|
|
sam2_processor, sam2_model = load_sam_model(
|
|
model_name="facebook/sam2.1-hiera-large"
|
|
)
|
|
|
|
input_boxes = [[box for box, _score, _text_label in filtered_items]]
|
|
proposal_batch = generate_proposals_batch(
|
|
model=sam2_model,
|
|
processor=sam2_processor,
|
|
images=[image],
|
|
bboxes_per_image=input_boxes,
|
|
)
|
|
proposals = proposal_batch[0] if proposal_batch else []
|
|
|
|
base_rgba = image.convert("RGBA")
|
|
_vis_image = base_rgba.copy()
|
|
summary_lines = []
|
|
|
|
colors = [
|
|
(255, 0, 0, 90),
|
|
(0, 255, 0, 90),
|
|
(0, 0, 255, 90),
|
|
(255, 255, 0, 90),
|
|
(255, 0, 255, 90),
|
|
(0, 255, 255, 90),
|
|
(255, 128, 0, 90),
|
|
(128, 0, 255, 90),
|
|
]
|
|
|
|
for _idx, ((_box, _score, _text_label), proposal) in enumerate(
|
|
zip(filtered_items, proposals)
|
|
):
|
|
mask_np = proposal["segment"]
|
|
color = colors[_idx % len(colors)]
|
|
|
|
mask_rgba = np.zeros((image.height, image.width, 4), dtype=np.uint8)
|
|
mask_rgba[mask_np] = color
|
|
|
|
mask_img = Image.fromarray(mask_rgba, mode="RGBA")
|
|
_vis_image = Image.alpha_composite(_vis_image, mask_img)
|
|
|
|
_draw = ImageDraw.Draw(_vis_image)
|
|
for (_box, _score, _text_label), proposal in zip(filtered_items, proposals):
|
|
_x1, _y1, _x2, _y2 = [float(v) for v in _box]
|
|
mask_area = int(proposal["area"])
|
|
_label = f"{_text_label} | owl={_score:.3f} | mask_area={mask_area}"
|
|
|
|
_draw.rectangle([_x1, _y1, _x2, _y2], outline=(255, 0, 0, 255), width=3)
|
|
|
|
try:
|
|
_tx1, _ty1, _tx2, _ty2 = _draw.textbbox((_x1, _y1), _label)
|
|
except Exception:
|
|
_tx1, _ty1, _tx2, _ty2 = _x1, _y1, _x1 + 220, _y1 + 20
|
|
|
|
_draw.rectangle(
|
|
[_tx1, max(0, _ty1 - 2), _tx2 + 4, _ty2 + 2],
|
|
fill=(255, 0, 0, 220),
|
|
)
|
|
_draw.text((_x1 + 2, max(0, _y1)), _label, fill="white")
|
|
summary_lines.append(
|
|
f"- {_text_label}: owl_score={_score:.3f}, mask_area={mask_area}"
|
|
)
|
|
|
|
if not filtered_items:
|
|
summary_text = (
|
|
"没有可用于分割的检测框，请先降低 OWLv2 的 score_threshold 或检查检测结果。"
|
|
)
|
|
elif not summary_lines:
|
|
summary_text = "没有生成任何 mask"
|
|
else:
|
|
summary_text = "\n".join(summary_lines)
|
|
|
|
mo.vstack(
|
|
[
|
|
mo.md(f"## SAM2 分割可视化结果\n\ndevice: `{device}`"),
|
|
mo.image(_vis_image, width=700),
|
|
mo.md(summary_text),
|
|
]
|
|
)
|
|
return
|
|
|
|
|
|
if __name__ == "__main__":
|
|
app.run()
|