mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-12 20:15:31 +08:00
404 lines
14 KiB
Python
404 lines
14 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,
|
||
)
|
||
|
||
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),
|
||
)
|
||
)
|
||
|
||
topdown_image = render_topdown_scene_map(
|
||
pathfinder=sim.pathfinder,
|
||
elements=TopDownSceneElements(room_nodes=room_nodes),
|
||
meters_per_pixel=meters_per_pixel,
|
||
)
|
||
return agent, room_nodes, sim, topdown_image, 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, topdown_image):
|
||
from compressors import HashPipeline
|
||
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",
|
||
]
|
||
|
||
pipeline = HashPipeline(
|
||
owlv2_model="google/owlv2-base-patch16-ensemble",
|
||
score_threshold=score_threshold,
|
||
postprocess_threshold=postprocess_threshold,
|
||
)
|
||
|
||
image = numpy_to_pil(all_room_views["room_00"][1])
|
||
|
||
output = pipeline.process_batch(
|
||
images=[image],
|
||
text_labels=text_labels,
|
||
batch_size=1,
|
||
return_debug_details=True,
|
||
)
|
||
meta = output.debug_meta[0] if output.debug_meta else {}
|
||
boxes = meta.get("boxes_xyxy", [])
|
||
scores = meta.get("scores", [])
|
||
labels = meta.get("labels", [])
|
||
masks = meta.get("masks", [])
|
||
pipeline_selected_indices = meta.get("selected_indices", [])
|
||
dropped_indices = meta.get("dropped_indices", [])
|
||
|
||
proposal_items = []
|
||
for idx, proposal_data in enumerate(masks):
|
||
detection_box = boxes[idx] if idx < len(boxes) else [0.0, 0.0, 0.0, 0.0]
|
||
detection_score = scores[idx] if idx < len(scores) else 0.0
|
||
detection_label = labels[idx] if idx < len(labels) else f"obj_{idx}"
|
||
proposal_items.append(
|
||
{
|
||
"idx": idx,
|
||
"proposal": proposal_data,
|
||
"bbox": detection_box,
|
||
"score": detection_score,
|
||
"label": detection_label,
|
||
"is_selected": idx in pipeline_selected_indices,
|
||
"is_dropped": idx in dropped_indices,
|
||
}
|
||
)
|
||
|
||
detected_items = [
|
||
{
|
||
"idx": idx,
|
||
"bbox": box,
|
||
"score": score,
|
||
"label": label,
|
||
"has_mask": idx < len(masks),
|
||
"is_selected": idx in pipeline_selected_indices,
|
||
}
|
||
for idx, (box, score, label) in enumerate(
|
||
zip(boxes, scores, labels, strict=False)
|
||
)
|
||
]
|
||
|
||
_vis_image = image.copy()
|
||
_draw = ImageDraw.Draw(_vis_image)
|
||
|
||
try:
|
||
font = ImageFont.truetype("arial.ttf", 18)
|
||
except Exception:
|
||
font = ImageFont.load_default()
|
||
|
||
for _item in detected_items:
|
||
_x1, _y1, _x2, _y2 = [float(v) for v in _item["bbox"]]
|
||
_label = f"{_item['label']}: {_item['score']:.3f}"
|
||
_outline_color = "green" if _item["is_selected"] else "red"
|
||
|
||
_draw.rectangle([_x1, _y1, _x2, _y2], outline=_outline_color, 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=_outline_color)
|
||
_draw.text((_x1 + 2, max(0, _y1)), _label, fill="white", font=font)
|
||
|
||
detection_lines = []
|
||
for _item in detected_items:
|
||
box_rounded = [round(_v, 2) for _v in _item["bbox"]]
|
||
status = "selected" if _item["is_selected"] else "dropped"
|
||
detection_lines.append(
|
||
f"- {_item['label']}: score={_item['score']:.3f}, box={box_rounded}, status={status}"
|
||
)
|
||
|
||
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}`"
|
||
f"\n\n绿色框表示最终被 pipeline 保留，红色框表示未被保留"
|
||
),
|
||
mo.md("## Top-down 房间采样图"),
|
||
mo.image(topdown_image, width=520),
|
||
mo.image(_vis_image, width=700),
|
||
mo.md(detection_text),
|
||
]
|
||
)
|
||
return detected_items, device, image, meta, proposal_items
|
||
|
||
|
||
@app.cell
|
||
def _(Image, ImageDraw, image, meta, mo, np, proposal_items):
|
||
from compressors.filter import (
|
||
MaskScoringConfig,
|
||
compute_mask_features,
|
||
score_mask,
|
||
)
|
||
|
||
image_shape = (image.height, image.width)
|
||
config = MaskScoringConfig()
|
||
fallback_reason = meta.get("fallback_reason")
|
||
selected_index_set = set(meta.get("selected_indices", []))
|
||
|
||
_kept = []
|
||
_rejected = []
|
||
for _item in proposal_items:
|
||
_idx = _item["idx"]
|
||
current_proposal = _item["proposal"]
|
||
_feat = compute_mask_features(current_proposal, image_shape)
|
||
_score = score_mask(current_proposal, image_shape, config)
|
||
_owl_label = _item["label"]
|
||
_owl_score = _item["score"]
|
||
_owl_bbox = _item["bbox"]
|
||
_is_selected = _idx in selected_index_set
|
||
|
||
_entry = {
|
||
"idx": _idx,
|
||
"proposal": current_proposal,
|
||
"features": _feat,
|
||
"mask_score": _score,
|
||
"owl_label": _owl_label,
|
||
"owl_score": _owl_score,
|
||
"owl_bbox": _owl_bbox,
|
||
"is_selected": _is_selected,
|
||
}
|
||
if _is_selected:
|
||
_kept.append(_entry)
|
||
else:
|
||
_rejected.append(_entry)
|
||
|
||
_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),
|
||
]
|
||
|
||
def _overlay_masks(base_img, entries, border_color, show_score=False):
|
||
"""Overlay masks on image and draw bounding boxes with labels."""
|
||
_rgba = base_img.convert("RGBA").copy()
|
||
for _e in entries:
|
||
_p = _e["proposal"]
|
||
_c = _colors[_e["idx"] % len(_colors)]
|
||
_mask_rgba = np.zeros((base_img.height, base_img.width, 4), dtype=np.uint8)
|
||
_mask_rgba[_p["segment"]] = _c
|
||
_rgba = Image.alpha_composite(
|
||
_rgba, Image.fromarray(_mask_rgba, mode="RGBA")
|
||
)
|
||
|
||
_draw = ImageDraw.Draw(_rgba)
|
||
for _e in entries:
|
||
_x1, _y1, _x2, _y2 = [float(v) for v in _e["owl_bbox"]]
|
||
_mask_area = int(_e["proposal"]["area"])
|
||
if show_score:
|
||
_label = f"{_e['owl_label']} | score={_e['mask_score']:.2f} | area={_mask_area}"
|
||
else:
|
||
_label = f"{_e['owl_label']} | area={_mask_area}"
|
||
_draw.rectangle([_x1, _y1, _x2, _y2], outline=border_color, width=3)
|
||
|
||
try:
|
||
_tb = _draw.textbbox((_x1, _y1), _label)
|
||
except Exception:
|
||
_tb = (_x1, _y1, _x1 + 200, _y1 + 20)
|
||
_draw.rectangle(
|
||
[_tb[0], max(0, _tb[1] - 2), _tb[2] + 4, _tb[3] + 2],
|
||
fill=border_color,
|
||
)
|
||
_draw.text((_x1 + 2, max(0, _y1)), _label, fill="white")
|
||
return _rgba
|
||
|
||
_all_entries = _kept + _rejected
|
||
_all_entries.sort(key=lambda e: e["idx"])
|
||
|
||
_before_img = _overlay_masks(image, _all_entries, "red", show_score=False)
|
||
_after_img = _overlay_masks(image, _kept, (0, 180, 0), show_score=True)
|
||
|
||
_draw_after = ImageDraw.Draw(_after_img)
|
||
for _e in _rejected:
|
||
_x1, _y1, _x2, _y2 = [float(v) for v in _e["owl_bbox"]]
|
||
_feat = _e["features"]
|
||
_reason_parts = []
|
||
if _feat.area_ratio < config.min_area_ratio:
|
||
_reason_parts.append(f"area_ratio={_feat.area_ratio:.4f}<min")
|
||
if _feat.area_ratio > config.max_area_ratio:
|
||
_reason_parts.append(f"area_ratio={_feat.area_ratio:.4f}>max")
|
||
_aspect = max(_feat.aspect_ratio, 1.0 / max(_feat.aspect_ratio, 1e-6))
|
||
if _aspect > config.max_aspect_ratio:
|
||
_reason_parts.append(f"aspect={_aspect:.1f}>max")
|
||
if _feat.fill_ratio < config.min_fill_ratio_hard:
|
||
_reason_parts.append(f"fill={_feat.fill_ratio:.3f}<min")
|
||
if (
|
||
_feat.num_components > config.max_components
|
||
and _feat.largest_component_ratio < config.min_largest_component_ratio
|
||
):
|
||
_reason_parts.append(f"fragments={_feat.num_components}")
|
||
if _feat.touch_edge_count >= config.reject_edge_touch_count:
|
||
_reason_parts.append(f"edge_touch={_feat.touch_edge_count}")
|
||
if (
|
||
_feat.touch_edge_count >= config.reject_large_edge_touch_count
|
||
and _feat.area_ratio > config.reject_large_edge_area_ratio
|
||
):
|
||
_reason_parts.append("edge+large")
|
||
if _e["idx"] in meta.get("dropped_indices", []):
|
||
_reason_parts.append("pipeline_drop")
|
||
|
||
_reason = ", ".join(_reason_parts) if _reason_parts else "unknown"
|
||
_label = f"X {_e['owl_label']}: {_reason}"
|
||
|
||
_dash_len = 8
|
||
_gap_len = 4
|
||
for _seg_start in range(int(_x1), int(_x2), _dash_len + _gap_len):
|
||
_seg_end = min(_seg_start + _dash_len, int(_x2))
|
||
_draw_after.line([(_seg_start, _y1), (_seg_end, _y1)], fill="red", width=2)
|
||
_draw_after.line([(_seg_start, _y2), (_seg_end, _y2)], fill="red", width=2)
|
||
for _seg_start in range(int(_y1), int(_y2), _dash_len + _gap_len):
|
||
_seg_end = min(_seg_start + _dash_len, int(_y2))
|
||
_draw_after.line([(_x1, _seg_start), (_x1, _seg_end)], fill="red", width=2)
|
||
_draw_after.line([(_x2, _seg_start), (_x2, _seg_end)], fill="red", width=2)
|
||
|
||
_draw_after.text((_x1 + 2, max(0, _y1 - 18)), _label, fill="red")
|
||
|
||
_total = len(proposal_items)
|
||
_kept_count = len(_kept)
|
||
_rej_count = len(_rejected)
|
||
|
||
_rej_detail_lines = []
|
||
for _e in _rejected:
|
||
_f = _e["features"]
|
||
_rej_detail_lines.append(
|
||
f" - **{_e['owl_label']}** (idx={_e['idx']}): "
|
||
f"area_ratio={_f.area_ratio:.4f}, fill_ratio={_f.fill_ratio:.3f}, "
|
||
f"aspect_ratio={_f.aspect_ratio:.1f}, touch_edge={_f.touch_edge_count}, "
|
||
f"components={_f.num_components}, mask_score={_e['mask_score']:.3f}"
|
||
)
|
||
|
||
_kept_detail_lines = []
|
||
for _e in _kept:
|
||
_f = _e["features"]
|
||
_kept_detail_lines.append(
|
||
f" - **{_e['owl_label']}** (idx={_e['idx']}): "
|
||
f"area_ratio={_f.area_ratio:.4f}, mask_score={_e['mask_score']:.3f}"
|
||
)
|
||
|
||
_detail_parts = []
|
||
if _kept_detail_lines:
|
||
_detail_parts.append("**保留的 mask：**\n" + "\n".join(_kept_detail_lines))
|
||
if _rej_detail_lines:
|
||
_detail_parts.append("**过滤掉的 mask：**\n" + "\n".join(_rej_detail_lines))
|
||
|
||
_detail_text = "\n\n".join(_detail_parts) if _detail_parts else "无 mask 数据"
|
||
if fallback_reason is not None:
|
||
_detail_text = f"**Pipeline fallback:** `{fallback_reason}`\n\n" + _detail_text
|
||
|
||
mo.vstack(
|
||
[
|
||
mo.md(
|
||
"## Mask 过滤对比"
|
||
f"\n\n共 {_total} 个 proposal → 保留 **{_kept_count}** 个，过滤掉 **{_rej_count}** 个"
|
||
),
|
||
mo.hstack(
|
||
[
|
||
mo.vstack(
|
||
[
|
||
mo.md(f"### 过滤前（全部 {_total} 个 proposal）"),
|
||
mo.image(_before_img, width=480),
|
||
]
|
||
),
|
||
mo.vstack(
|
||
[
|
||
mo.md(
|
||
f"### 过滤后（{_kept_count} 个保留，{_rej_count} 个未保留）"
|
||
),
|
||
mo.image(_after_img, width=480),
|
||
]
|
||
),
|
||
]
|
||
),
|
||
mo.md(_detail_text),
|
||
]
|
||
)
|
||
return
|
||
|
||
|
||
if __name__ == "__main__":
|
||
app.run()
|