refactor(visualization): enhance proposal filtering visualization

This commit is contained in:
2026-04-18 23:24:05 +08:00
parent e3cdc1c6e7
commit d9745f45dc

View File

@@ -22,7 +22,6 @@ def _(np):
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
@@ -52,13 +51,12 @@ def _(np):
)
)
render_topdown_scene_map(
topdown_image = 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
return agent, room_nodes, sim, topdown_image, views_per_room
@app.cell
@@ -75,7 +73,7 @@ def _(agent, room_nodes, sim, views_per_room):
@app.cell
def _(ImageDraw, ImageFont, all_room_views, mo):
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
@@ -113,7 +111,40 @@ def _(ImageDraw, ImageFont, all_room_views, mo):
boxes = meta.get("boxes_xyxy", [])
scores = meta.get("scores", [])
labels = meta.get("labels", [])
filtered_items = list(zip(boxes, scores, labels, strict=False))
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)
@@ -123,11 +154,12 @@ def _(ImageDraw, ImageFont, all_room_views, mo):
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}"
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="red", width=3)
_draw.rectangle([_x1, _y1, _x2, _y2], outline=_outline_color, width=3)
try:
_tx1, _ty1, _tx2, _ty2 = _draw.textbbox((_x1, _y1), _label, font=font)
@@ -137,15 +169,15 @@ def _(ImageDraw, ImageFont, all_room_views, mo):
_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.rectangle(text_bg, fill=_outline_color)
_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]
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"- {_text_label}: score={_score:.3f}, box={box_rounded}"
f"- {_item['label']}: score={_item['score']:.3f}, box={box_rounded}, status={status}"
)
if not detection_lines:
@@ -159,59 +191,56 @@ def _(ImageDraw, ImageFont, all_room_views, mo):
"## 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 device, filtered_items, image, meta
return detected_items, device, image, meta, proposal_items
@app.cell
def _(meta):
proposals = meta.get("masks", [])
return (proposals,)
@app.cell
def _(Image, ImageDraw, filtered_items, image, mo, np, proposals):
def _(Image, ImageDraw, image, meta, mo, np, proposal_items):
from compressors.filter import (
MaskScoringConfig,
compute_mask_features,
score_mask,
should_reject_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 _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]
)
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": proposal,
"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_rejected:
_rejected.append(_entry)
else:
if _is_selected:
_kept.append(_entry)
else:
_rejected.append(_entry)
_colors = [
(255, 0, 0, 90),
@@ -289,6 +318,8 @@ def _(Image, ImageDraw, filtered_items, image, mo, np, proposals):
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}"
@@ -306,7 +337,7 @@ def _(Image, ImageDraw, filtered_items, image, mo, np, proposals):
_draw_after.text((_x1 + 2, max(0, _y1 - 18)), _label, fill="red")
_total = len(proposals)
_total = len(proposal_items)
_kept_count = len(_kept)
_rej_count = len(_rejected)
@@ -335,25 +366,27 @@ def _(Image, ImageDraw, filtered_items, image, mo, np, proposals):
_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}mask → 保留 **{_kept_count}** 个,过滤掉 **{_rej_count}** 个"
f"\n\n{_total}proposal → 保留 **{_kept_count}** 个,过滤掉 **{_rej_count}** 个"
),
mo.hstack(
[
mo.vstack(
[
mo.md(f"### 过滤前({_total} 个)"),
mo.md(f"### 过滤前(全部 {_total} proposal"),
mo.image(_before_img, width=480),
]
),
mo.vstack(
[
mo.md(
f"### 过滤后({_kept_count} 个保留,{_rej_count}过滤"
f"### 过滤后({_kept_count} 个保留,{_rej_count}未保留"
),
mo.image(_after_img, width=480),
]