Shaun ZhouandClaude Opus 5
9439cbb56a
ultralytics 没有可用的开关：yolo26 是端到端模型，non_max_suppression 对它早退 （utils/nms.py 里 end2end 分支），只按置信度与 max_det 截断，传入的 iou 和 agnostic 都不参与 IoU 抑制；agnostic_nms 只作用于检测头 top-k（每个 anchor 取一个类别）， 压不掉不同 anchor 产生的重叠框。 所以给 DetectionValidator.postprocess 包一层 filter_overlapping_outputs。训练期验证、 早停、best.pt 的选择都随之遵守这条约束，指标与最终推理一致。配置项 train.max_overlap （默认 0，负数关闭），由 detect.py 摘出后自行实现，不传给 ultralytics。 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
210 lines
8.1 KiB
Python
210 lines
8.1 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 suppress_overlaps(
|
||
boxes: Sequence[Box], max_iou: float = 0.0, contained: bool = True
|
||
) -> list[int]:
|
||
"""贪心地挑出互不重叠的框，返回保留下来的下标（按输入顺序）。
|
||
|
||
按置信度从高到低考察，与任一已保留框重叠超过 `max_iou` 就丢弃。
|
||
|
||
为什么需要这一步：NMS 只在**同类**内、且按固定 IoU 阈值抑制，所以同一处缺陷被两个
|
||
类别各框一次（跨类重复），或两个框 IoU 低于 NMS 阈值时，都会留下重叠框。而人工标注的
|
||
真值里一处缺陷只有一个框，零重叠——推理侧不加这层约束，指标和图上都会出现真值里
|
||
不可能有的形态。
|
||
|
||
Args:
|
||
max_iou: 允许的最大 IoU。0 表示碰到就算重叠，完全不允许。
|
||
contained: 是否额外丢弃被已保留框完全包住的框。整图大框套小框时 IoU 可能很小
|
||
（面积悬殊），但那显然是同一处目标的重复表达。
|
||
"""
|
||
order = sorted(range(len(boxes)), key=lambda i: -(boxes[i].score or 0.0))
|
||
kept: list[int] = []
|
||
for index in order:
|
||
box = boxes[index]
|
||
if any(
|
||
iou(box, boxes[other]) > max_iou or (contained and _inside(box, boxes[other]))
|
||
for other in kept
|
||
):
|
||
continue
|
||
kept.append(index)
|
||
return sorted(kept)
|
||
|
||
|
||
def _inside(inner: Box, outer: Box) -> bool:
|
||
"""inner 是否完全落在 outer 里（含边界）。"""
|
||
return (
|
||
outer.x1 <= inner.x1 and outer.y1 <= inner.y1
|
||
and outer.x2 >= inner.x2 and outer.y2 >= inner.y2
|
||
)
|
||
|
||
|
||
def filter_overlapping_outputs(
|
||
outputs: list[dict], max_iou: float = 0.0, contained: bool = True
|
||
) -> list[dict]:
|
||
"""给 ultralytics 验证器 postprocess 的输出补上"框不重叠"约束。
|
||
|
||
每个元素是一张图的结果：`bboxes`（xyxy）、`conf`、`cls` 等张量。坐标是不是归一化的
|
||
无所谓，IoU 与尺度无关。
|
||
|
||
为什么要在这里补：yolo26 这类**端到端模型不走 IoU 抑制**。ultralytics 的
|
||
`non_max_suppression` 对它们早退（`utils/nms.py` 里 `prediction.shape[-1] == 6 or end2end`
|
||
的分支），只按置信度与 `max_det` 截断，传入的 `iou` 和 `agnostic` 都不生效；
|
||
`agnostic_nms` 只作用于检测头的 top-k（每个 anchor 取一个类别），压不掉不同 anchor
|
||
产生的重叠框。不补这一步，训练期的验证指标就和最终推理对不上。
|
||
"""
|
||
filtered = []
|
||
for output in outputs:
|
||
boxes = [
|
||
Box(int(class_id), *(float(v) for v in xyxy), float(score))
|
||
for xyxy, score, class_id in zip(
|
||
output["bboxes"].tolist(), output["conf"].tolist(), output["cls"].tolist(),
|
||
strict=True,
|
||
)
|
||
]
|
||
kept = suppress_overlaps(boxes, max_iou=max_iou, contained=contained)
|
||
if len(kept) == len(boxes):
|
||
filtered.append(output)
|
||
continue
|
||
# 张量用列表下标筛选；空结果不能用 [] 索引，改用零长切片
|
||
filtered.append({
|
||
key: (value[kept] if kept else value[:0]) if _is_per_box(value, len(boxes)) else value
|
||
for key, value in output.items()
|
||
})
|
||
return filtered
|
||
|
||
|
||
def _is_per_box(value, count: int) -> bool:
|
||
"""这个张量是不是"每个框一行"——只有这种才跟着筛选。"""
|
||
return hasattr(value, "shape") and len(getattr(value, "shape", ())) >= 1 and len(value) == count
|
||
|
||
|
||
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
|