Files
ModelTrainEval/tests/test_detection.py
T
Shaun ZhouandClaude Opus 5 154677032f Revert "约束预测框互不重叠"(eba533a、9439cbb)
去掉强制非重叠:detect.py 不再给 DetectionValidator.postprocess 打补丁,
配置里移除 train.max_overlap 与 agnostic_nms,可视化脚本恢复为直接使用模型原始输出,
engine/detection 里的 suppress_overlaps / filter_overlapping_outputs 一并移除。

约束会把 mAP50 从 0.555 压到 0.483(precision 0.536->0.612),代价是按置信度抑制时
高分错框会压掉本可命中的低分框。需要时用 git revert 这次提交即可取回。

按类别分目录的可视化产出结构(36fff98)保留。

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-18 15:34:13 +08:00

132 lines
4.4 KiB
Python

"""检测结果的匹配与绘制。"""
import pytest
pytest.importorskip("PIL")
from engine.detection import ( # noqa: E402
COLOR_GT,
COLOR_MATCHED,
COLOR_UNMATCHED,
Box,
draw_boxes,
from_yolo,
iou,
match_boxes,
parse_label,
)
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)