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>
248 lines
8.6 KiB
Python
248 lines
8.6 KiB
Python
"""检测结果的匹配与绘制。"""
|
|
|
|
import pytest
|
|
|
|
pytest.importorskip("PIL")
|
|
|
|
from engine.detection import ( # noqa: E402
|
|
COLOR_GT,
|
|
COLOR_MATCHED,
|
|
COLOR_UNMATCHED,
|
|
Box,
|
|
draw_boxes,
|
|
filter_overlapping_outputs,
|
|
from_yolo,
|
|
iou,
|
|
match_boxes,
|
|
parse_label,
|
|
suppress_overlaps,
|
|
)
|
|
|
|
|
|
def test_from_yolo_converts_center_to_corners():
|
|
assert from_yolo(2, 0.5, 0.5, 0.2, 0.4) == Box(2, 0.4, 0.3, 0.6, 0.7)
|
|
|
|
|
|
def test_parse_label_reads_boxes_and_empty_files():
|
|
boxes = parse_label("0 0.5 0.5 0.2 0.2\n3 0.1 0.1 0.1 0.1\n")
|
|
assert [b.class_id for b in boxes] == [0, 3]
|
|
assert parse_label("") == [] # 负样本
|
|
assert parse_label("\n\n") == []
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"a,b,expected",
|
|
[
|
|
((0, 0, 1, 1), (0, 0, 1, 1), 1.0), # 完全重合
|
|
((0, 0, 1, 1), (2, 2, 3, 3), 0.0), # 不相交
|
|
((0, 0, 1, 1), (1, 1, 2, 2), 0.0), # 只碰到角
|
|
((0, 0, 2, 2), (1, 1, 3, 3), 1 / 7), # 交 1，并 7
|
|
],
|
|
)
|
|
def test_iou(a, b, expected):
|
|
assert iou(Box(0, *a), Box(0, *b)) == pytest.approx(expected)
|
|
|
|
|
|
def test_matches_by_confidence_first():
|
|
"""两个预测抢同一个真值时，置信度高的赢——与 mAP 的规则一致。"""
|
|
truth = [from_yolo(0, 0.5, 0.5, 0.2, 0.2)]
|
|
weak = Box(0, 0.4, 0.4, 0.6, 0.6, 0.3)
|
|
strong = Box(0, 0.41, 0.41, 0.59, 0.59, 0.9)
|
|
match = match_boxes([weak, strong], truth, iou_threshold=0.3)
|
|
|
|
assert match.pairs == [(1, 0)] # 下标 1 = strong
|
|
assert match.unmatched_predictions == [0]
|
|
assert match.unmatched_ground_truth == []
|
|
|
|
|
|
def test_low_iou_counts_as_both_miss_and_false_positive():
|
|
truth = [from_yolo(0, 0.2, 0.2, 0.1, 0.1)]
|
|
prediction = Box(0, 0.7, 0.7, 0.9, 0.9, 0.9)
|
|
match = match_boxes([prediction], truth)
|
|
assert match.pairs == []
|
|
assert match.unmatched_predictions == [0] # 误检
|
|
assert match.unmatched_ground_truth == [0] # 同时也是漏检
|
|
|
|
|
|
def test_class_mismatch_is_not_a_hit_unless_allowed():
|
|
"""定位准但分类错：默认不算命中，same_class=False 时可单看位置。"""
|
|
truth = [from_yolo(1, 0.5, 0.5, 0.2, 0.2)]
|
|
prediction = Box(0, 0.4, 0.4, 0.6, 0.6, 0.9)
|
|
|
|
assert match_boxes([prediction], truth).pairs == []
|
|
assert match_boxes([prediction], truth, same_class=False).pairs == [(0, 0)]
|
|
|
|
|
|
def test_each_ground_truth_is_claimed_once():
|
|
"""一个真值只能被一个预测认领，多出来的预测算误检。"""
|
|
truth = [from_yolo(0, 0.5, 0.5, 0.4, 0.4)]
|
|
predictions = [
|
|
Box(0, 0.31, 0.31, 0.69, 0.69, 0.9),
|
|
Box(0, 0.32, 0.32, 0.68, 0.68, 0.8),
|
|
]
|
|
match = match_boxes(predictions, truth)
|
|
assert len(match.pairs) == 1
|
|
assert len(match.unmatched_predictions) == 1
|
|
|
|
|
|
def test_no_prediction_on_background_is_clean():
|
|
match = match_boxes([], [])
|
|
assert match == ([], [], [])
|
|
|
|
|
|
def test_multiple_boxes_are_matched_pairwise():
|
|
truth = [from_yolo(0, 0.2, 0.2, 0.1, 0.1), from_yolo(1, 0.8, 0.8, 0.1, 0.1)]
|
|
predictions = [
|
|
Box(1, 0.75, 0.75, 0.85, 0.85, 0.7),
|
|
Box(0, 0.15, 0.15, 0.25, 0.25, 0.9),
|
|
]
|
|
match = match_boxes(predictions, truth)
|
|
assert sorted(match.pairs) == [(0, 1), (1, 0)]
|
|
assert not match.unmatched_predictions and not match.unmatched_ground_truth
|
|
|
|
|
|
def _colors(image):
|
|
return {pixel for pixel in image.getdata()}
|
|
|
|
|
|
def test_draw_uses_a_distinct_colour_per_role():
|
|
"""三种角色三种颜色：绿=真值，蓝=命中的预测，红=没对上的预测。"""
|
|
from PIL import Image
|
|
|
|
truth = [from_yolo(0, 0.25, 0.25, 0.2, 0.2), from_yolo(0, 0.75, 0.25, 0.2, 0.2)]
|
|
predictions = [
|
|
Box(0, 0.16, 0.16, 0.34, 0.34, 0.9), # 命中第一个
|
|
Box(0, 0.5, 0.8, 0.6, 0.9, 0.4), # 没有对应真值
|
|
]
|
|
match = match_boxes(predictions, truth, iou_threshold=0.5)
|
|
assert match.pairs == [(0, 0)]
|
|
|
|
image = Image.new("RGB", (200, 200), "black")
|
|
draw_boxes(image, truth, predictions, ["blackdot"], match)
|
|
present = _colors(image)
|
|
assert COLOR_GT in present and COLOR_MATCHED in present and COLOR_UNMATCHED in present
|
|
|
|
|
|
def test_draw_handles_boxes_at_the_border():
|
|
"""框贴着上边缘时，标签文字要画到框内，不能画到图外去。"""
|
|
from PIL import Image
|
|
|
|
image = Image.new("RGB", (64, 64), "black")
|
|
prediction = Box(0, 0.0, 0.0, 0.5, 0.05, 0.9)
|
|
draw_boxes(image, [], [prediction], ["blackdot"], match_boxes([prediction], []))
|
|
assert COLOR_UNMATCHED in _colors(image)
|
|
|
|
|
|
# ---------------- 不重叠约束 ----------------
|
|
|
|
|
|
def test_suppress_keeps_the_highest_score_of_an_overlapping_pair():
|
|
"""同一处缺陷被框两次时，留置信度高的那个。"""
|
|
strong = Box(0, 0.3, 0.3, 0.7, 0.7, 0.9)
|
|
weak = Box(1, 0.35, 0.35, 0.75, 0.75, 0.4)
|
|
assert suppress_overlaps([weak, strong]) == [1] # 下标 1 = strong
|
|
|
|
|
|
def test_suppress_is_cross_class():
|
|
"""NMS 只在同类内抑制；跨类重复要靠这一步。"""
|
|
a = Box(0, 0.3, 0.3, 0.7, 0.7, 0.9)
|
|
b = Box(5, 0.31, 0.31, 0.69, 0.69, 0.8)
|
|
assert suppress_overlaps([a, b]) == [0]
|
|
|
|
|
|
def test_suppress_keeps_disjoint_boxes():
|
|
"""真正分开的多个缺陷不能被误伤。"""
|
|
boxes = [
|
|
Box(0, 0.0, 0.0, 0.2, 0.2, 0.9),
|
|
Box(0, 0.5, 0.5, 0.7, 0.7, 0.8),
|
|
Box(1, 0.8, 0.1, 0.95, 0.3, 0.7),
|
|
]
|
|
assert suppress_overlaps(boxes) == [0, 1, 2]
|
|
|
|
|
|
def test_touching_edges_are_not_overlap():
|
|
"""边缘相接、面积不相交，IoU=0，应当都保留。"""
|
|
left = Box(0, 0.0, 0.0, 0.5, 1.0, 0.9)
|
|
right = Box(0, 0.5, 0.0, 1.0, 1.0, 0.8)
|
|
assert suppress_overlaps([left, right]) == [0, 1]
|
|
|
|
|
|
def test_contained_box_is_dropped_even_when_iou_is_below_the_threshold():
|
|
"""整图大框套一个小框：面积悬殊导致 IoU 很小，但显然是同一处目标的重复表达。
|
|
|
|
这正是 `contained` 的用武之地——只在允许轻微重叠（max_iou>0）时才有区别，
|
|
因为 max_iou=0 时任何包含都已经因 IoU>0 被抑制了。
|
|
"""
|
|
big = Box(0, 0.0, 0.0, 1.0, 1.0, 0.9)
|
|
small = Box(1, 0.45, 0.45, 0.55, 0.55, 0.8)
|
|
assert iou(big, small) < 0.02
|
|
|
|
assert suppress_overlaps([big, small], max_iou=0.1) == [0]
|
|
assert suppress_overlaps([big, small], max_iou=0.1, contained=False) == [0, 1]
|
|
|
|
|
|
def test_max_iou_allows_slight_overlap():
|
|
a = Box(0, 0.0, 0.0, 0.5, 0.5, 0.9)
|
|
b = Box(0, 0.4, 0.4, 0.9, 0.9, 0.8) # IoU 约 0.02
|
|
assert suppress_overlaps([a, b]) == [0]
|
|
assert suppress_overlaps([a, b], max_iou=0.1) == [0, 1]
|
|
|
|
|
|
def test_suppress_handles_empty_and_single():
|
|
assert suppress_overlaps([]) == []
|
|
assert suppress_overlaps([Box(0, 0.1, 0.1, 0.2, 0.2, 0.5)]) == [0]
|
|
|
|
|
|
# ---------------- 验证器输出的过滤 ----------------
|
|
|
|
|
|
def test_filter_overlapping_outputs_drops_the_lower_score():
|
|
"""ultralytics 验证器的输出是张量字典，筛选要对每个"每框一行"的张量同步生效。"""
|
|
torch = pytest.importorskip("torch")
|
|
|
|
output = {
|
|
"bboxes": torch.tensor([[0.0, 0.0, 10.0, 10.0], [1.0, 1.0, 11.0, 11.0],
|
|
[50.0, 50.0, 60.0, 60.0]]),
|
|
"conf": torch.tensor([0.9, 0.5, 0.8]),
|
|
"cls": torch.tensor([0.0, 1.0, 2.0]),
|
|
"extra": torch.zeros((3, 0)),
|
|
}
|
|
filtered = filter_overlapping_outputs([output])[0]
|
|
|
|
assert filtered["conf"].tolist() == pytest.approx([0.9, 0.8]) # 中间那个与第一个重叠
|
|
assert filtered["cls"].tolist() == [0.0, 2.0]
|
|
assert filtered["bboxes"].shape == (2, 4)
|
|
assert filtered["extra"].shape == (2, 0)
|
|
|
|
|
|
def test_filter_keeps_the_object_untouched_when_nothing_overlaps():
|
|
torch = pytest.importorskip("torch")
|
|
|
|
output = {
|
|
"bboxes": torch.tensor([[0.0, 0.0, 1.0, 1.0], [5.0, 5.0, 6.0, 6.0]]),
|
|
"conf": torch.tensor([0.9, 0.8]),
|
|
"cls": torch.tensor([0.0, 0.0]),
|
|
}
|
|
filtered = filter_overlapping_outputs([output])
|
|
assert filtered[0] is output # 没有变化时原样返回，不做无谓的拷贝
|
|
|
|
|
|
def test_filter_handles_images_with_no_detection():
|
|
torch = pytest.importorskip("torch")
|
|
|
|
empty = {"bboxes": torch.zeros((0, 4)), "conf": torch.zeros(0), "cls": torch.zeros(0)}
|
|
assert filter_overlapping_outputs([empty])[0]["conf"].numel() == 0
|
|
|
|
|
|
def test_filter_can_empty_an_image_entirely():
|
|
"""全部互相重叠时只留一个；这里构造两个几乎重合的框。"""
|
|
torch = pytest.importorskip("torch")
|
|
|
|
output = {
|
|
"bboxes": torch.tensor([[0.0, 0.0, 10.0, 10.0], [0.1, 0.1, 10.1, 10.1]]),
|
|
"conf": torch.tensor([0.3, 0.9]),
|
|
"cls": torch.tensor([1.0, 1.0]),
|
|
}
|
|
filtered = filter_overlapping_outputs([output])[0]
|
|
assert filtered["conf"].tolist() == pytest.approx([0.9])
|