mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-13 04:25:32 +08:00
SikongJueluo
3638ffdb8d
- Add FramePacket dataclass to encapsulate per-image pipeline state - Rename internal methods with underscore prefix convention - Replace separate filter_batch/crop_batch with unified process_batch method - Update notebook to use new HashPipeline API
371 lines
12 KiB
Python
371 lines
12 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 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", [])
|
||
filtered_items = list(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 _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, meta
|
||
|
||
|
||
@app.cell
|
||
def _(meta):
|
||
proposals = meta.get("masks", [])
|
||
return (proposals,)
|
||
|
||
|
||
@app.cell
|
||
def _(Image, ImageDraw, filtered_items, image, mo, np, proposals):
|
||
from compressors.filter import (
|
||
MaskScoringConfig,
|
||
compute_mask_features,
|
||
score_mask,
|
||
should_reject_mask,
|
||
)
|
||
|
||
image_shape = (image.height, image.width)
|
||
config = MaskScoringConfig()
|
||
|
||
_kept = []
|
||
_rejected = []
|
||
for _idx, proposal in enumerate(proposals):
|
||
_feat = compute_mask_features(proposal, image_shape)
|
||
_is_rejected = should_reject_mask(_feat, config)
|
||
_score = score_mask(proposal, image_shape, config)
|
||
_owl_label = (
|
||
filtered_items[_idx][2] if _idx < len(filtered_items) else f"obj_{_idx}"
|
||
)
|
||
_owl_score = filtered_items[_idx][1] if _idx < len(filtered_items) else 0.0
|
||
_owl_bbox = (
|
||
filtered_items[_idx][0] if _idx < len(filtered_items) else [0, 0, 0, 0]
|
||
)
|
||
|
||
_entry = {
|
||
"idx": _idx,
|
||
"proposal": proposal,
|
||
"features": _feat,
|
||
"mask_score": _score,
|
||
"owl_label": _owl_label,
|
||
"owl_score": _owl_score,
|
||
"owl_bbox": _owl_bbox,
|
||
}
|
||
if _is_rejected:
|
||
_rejected.append(_entry)
|
||
else:
|
||
_kept.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")
|
||
|
||
_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(proposals)
|
||
_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 数据"
|
||
|
||
mo.vstack(
|
||
[
|
||
mo.md(
|
||
"## Mask 过滤对比"
|
||
f"\n\n共 {_total} 个 mask → 保留 **{_kept_count}** 个，过滤掉 **{_rej_count}** 个"
|
||
),
|
||
mo.hstack(
|
||
[
|
||
mo.vstack(
|
||
[
|
||
mo.md(f"### 过滤前（{_total} 个）"),
|
||
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()
|