Shaun ZhouandClaude Opus 5
788bc36ab8
scripts/visualize_detections.py：在整个数据集上推理，绿=真值、蓝=命中的预测、 红=没对上真值的预测，按结果类型分目录（ok / missed / extra / none），再配一份 index.csv（逐图的真值数、预测数、命中数、最高置信度）与 summary.json（逐类 TP/FP/FN）。 --only errors 只输出有问题的图。 匹配逻辑放在 engine/detection.py：按置信度贪心、IoU 阈值、类别需一致，与 mAP 的规则 一致，所以图上看到的漏检误检和指标算的是同一回事。same_class=False 可以单看定位准不准 ——定位对但分类错，和完全没找到是两种问题。 train 与 val 分开放：训练集的图模型见过，那部分结果不能当泛化能力看。 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
132 lines
4.7 KiB
Python
132 lines
4.7 KiB
Python
"""检测结果的匹配与绘制：把预测框和真值框放在一起看。
|
|
|
|
匹配用标准的贪心策略——预测按置信度从高到低，each 认领 IoU 最高且未被认领的真值框。
|
|
这与 mAP 的匹配规则一致，所以图上看到的"漏检/误检"和指标算的是同一回事。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Sequence
|
|
from typing import NamedTuple
|
|
|
|
#: 真值绿、命中的预测蓝、未命中的预测红。红绿对比最刺眼，留给"错"的那一类
|
|
COLOR_GT = (0, 220, 0)
|
|
COLOR_MATCHED = (0, 160, 255)
|
|
COLOR_UNMATCHED = (255, 40, 40)
|
|
|
|
|
|
class Box(NamedTuple):
|
|
"""归一化的 xyxy 框。`score` 仅预测框有。"""
|
|
|
|
class_id: int
|
|
x1: float
|
|
y1: float
|
|
x2: float
|
|
y2: float
|
|
score: float | None = None
|
|
|
|
|
|
def from_yolo(class_id: int, cx: float, cy: float, width: float, height: float) -> Box:
|
|
"""YOLO 的中心点记法 -> xyxy。"""
|
|
return Box(class_id, cx - width / 2, cy - height / 2, cx + width / 2, cy + height / 2)
|
|
|
|
|
|
def parse_label(text: str) -> list[Box]:
|
|
"""读一个 YOLO 标签文件的内容；空文件返回空列表（负样本）。"""
|
|
boxes = []
|
|
for line in text.splitlines():
|
|
parts = line.split()
|
|
if len(parts) == 5:
|
|
boxes.append(from_yolo(int(parts[0]), *(float(v) for v in parts[1:])))
|
|
return boxes
|
|
|
|
|
|
def iou(a: Box, b: Box) -> float:
|
|
"""两个框的交并比。不相交时为 0。"""
|
|
left, top = max(a.x1, b.x1), max(a.y1, b.y1)
|
|
right, bottom = min(a.x2, b.x2), min(a.y2, b.y2)
|
|
if right <= left or bottom <= top:
|
|
return 0.0
|
|
overlap = (right - left) * (bottom - top)
|
|
area_a = max(a.x2 - a.x1, 0) * max(a.y2 - a.y1, 0)
|
|
area_b = max(b.x2 - b.x1, 0) * max(b.y2 - b.y1, 0)
|
|
union = area_a + area_b - overlap
|
|
return overlap / union if union > 0 else 0.0
|
|
|
|
|
|
class Match(NamedTuple):
|
|
"""一张图的匹配结果。下标分别指向 predictions 与 ground_truth。"""
|
|
|
|
pairs: list[tuple[int, int]] # (预测下标, 真值下标)
|
|
unmatched_predictions: list[int] # 误检
|
|
unmatched_ground_truth: list[int] # 漏检
|
|
|
|
|
|
def match_boxes(
|
|
predictions: Sequence[Box],
|
|
ground_truth: Sequence[Box],
|
|
iou_threshold: float = 0.5,
|
|
same_class: bool = True,
|
|
) -> Match:
|
|
"""贪心匹配，与 mAP 的规则一致。
|
|
|
|
Args:
|
|
iou_threshold: 低于它不算命中。
|
|
same_class: 是否要求类别一致。设 False 可以单看"位置对不对"——
|
|
定位准但分类错，和完全没找到是两种问题，值得分开看。
|
|
"""
|
|
order = sorted(range(len(predictions)), key=lambda i: -(predictions[i].score or 0.0))
|
|
taken: set[int] = set()
|
|
pairs = []
|
|
unmatched_predictions = []
|
|
for index in order:
|
|
prediction = predictions[index]
|
|
best, best_iou = None, iou_threshold
|
|
for candidate, truth in enumerate(ground_truth):
|
|
if candidate in taken:
|
|
continue
|
|
if same_class and truth.class_id != prediction.class_id:
|
|
continue
|
|
overlap = iou(prediction, truth)
|
|
if overlap >= best_iou:
|
|
best, best_iou = candidate, overlap
|
|
if best is None:
|
|
unmatched_predictions.append(index)
|
|
else:
|
|
taken.add(best)
|
|
pairs.append((index, best))
|
|
return Match(
|
|
sorted(pairs),
|
|
sorted(unmatched_predictions),
|
|
[i for i in range(len(ground_truth)) if i not in taken],
|
|
)
|
|
|
|
|
|
def draw_boxes(image, ground_truth: Sequence[Box], predictions: Sequence[Box],
|
|
classes: Sequence[str], match: Match, width: int = 2):
|
|
"""把真值与预测画在图上（就地修改传入的 PIL 图像并返回）。
|
|
|
|
真值绿色、命中的预测蓝色、未命中的预测红色；预测标注类别与置信度。
|
|
"""
|
|
from PIL import ImageDraw
|
|
|
|
draw = ImageDraw.Draw(image)
|
|
w, h = image.size
|
|
matched = {index for index, _ in match.pairs}
|
|
|
|
def rectangle(box: Box, color, label: str | None) -> None:
|
|
xy = [box.x1 * w, box.y1 * h, box.x2 * w, box.y2 * h]
|
|
draw.rectangle(xy, outline=color, width=width)
|
|
if label:
|
|
# 文字画在框上方，贴边时改画到框内，免得被裁掉
|
|
y = xy[1] - 11 if xy[1] > 12 else xy[1] + 1
|
|
draw.text((xy[0] + 1, y), label, fill=color)
|
|
|
|
for box in ground_truth:
|
|
rectangle(box, COLOR_GT, None)
|
|
for index, box in enumerate(predictions):
|
|
color = COLOR_MATCHED if index in matched else COLOR_UNMATCHED
|
|
name = classes[box.class_id] if box.class_id < len(classes) else str(box.class_id)
|
|
rectangle(box, color, f"{name} {box.score:.2f}" if box.score is not None else name)
|
|
return image
|