feat(detect): 训练期的验证也约束预测框互不重叠
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>
This commit is contained in:
co-authored by
Claude Opus 5
parent
eba533a5b2
commit
9439cbb56a
@@ -48,6 +48,13 @@ train:
|
||||
hsv_s: 0.3
|
||||
hsv_v: 0.3
|
||||
|
||||
# 真值里一处缺陷只有一个框、零重叠(实测 11211 张图、151 张多框图,重叠为 0),
|
||||
# 所以验证时也把预测约束成互不重叠:早停与 best.pt 的选择才和最终推理一致。
|
||||
# 负数=不约束。注意这不是 ultralytics 的参数,由 detect.py 自己实现。
|
||||
max_overlap: 0.0
|
||||
# 端到端头的 top-k 每个 anchor 只取一个类别,压掉同一处缺陷被两个类别各框一次
|
||||
agnostic_nms: true
|
||||
|
||||
plots: true # results.png / PR 曲线 / 混淆矩阵 / val_batch 预测图
|
||||
save_period: -1
|
||||
|
||||
|
||||
@@ -6,8 +6,11 @@
|
||||
--checkpoint runs/detect/film/weights/best.pt
|
||||
|
||||
数据是在 Label Studio 里修正过的标签,导出并划分 train/val 之后放在 data/ 下的数据集。
|
||||
框来自 Grad-CAM / SHAP 的弱监督标注,不是人工标注——指标该按"与热力区域的一致性"
|
||||
来读,而不是当成真值上的检测精度。
|
||||
框来自 Grad-CAM / SHAP 的弱监督标注再经人工修正。
|
||||
|
||||
**预测框约束为互不重叠**(`train.max_overlap`,默认 0)。真值里一处缺陷只有一个框、
|
||||
零重叠,模型却会输出重叠框,所以验证阶段补一层抑制——早停与 best.pt 的选择都据此进行,
|
||||
指标才和最终推理一致。设为负数可关掉。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -19,6 +22,7 @@ from pathlib import Path
|
||||
|
||||
import yaml
|
||||
|
||||
from engine.detection import filter_overlapping_outputs
|
||||
from utils import load_config, seed_everything, setup_logging
|
||||
from utils.device import DeviceNotAvailableError, describe_device, resolve_device
|
||||
|
||||
@@ -35,6 +39,23 @@ def parse_args() -> argparse.Namespace:
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def install_overlap_suppression(max_overlap: float) -> None:
|
||||
"""给验证器的后处理包一层"框不重叠"约束。
|
||||
|
||||
只能用打补丁的方式:ultralytics 没有对应的开关——端到端模型(yolo26)的 NMS 早退,
|
||||
传入的 iou / agnostic_nms 都不参与 IoU 抑制(见 engine/detection.filter_overlapping_outputs
|
||||
的说明)。补在 postprocess 之后,训练期验证、早停、best.pt 的选择就都遵守这条约束。
|
||||
"""
|
||||
from ultralytics.models.yolo.detect import DetectionValidator
|
||||
|
||||
original = DetectionValidator.postprocess
|
||||
|
||||
def postprocess(self, preds):
|
||||
return filter_overlapping_outputs(original(self, preds), max_iou=max_overlap)
|
||||
|
||||
DetectionValidator.postprocess = postprocess
|
||||
|
||||
|
||||
def to_ultralytics_device(device) -> str | int:
|
||||
"""torch.device -> ultralytics 的 device 记法(cuda 用序号,其余用字符串)。"""
|
||||
if device.type == "cuda":
|
||||
@@ -79,10 +100,18 @@ def main() -> int:
|
||||
|
||||
train_cfg = dict(cfg.get("train", {}))
|
||||
swanlab_cfg = dict(train_cfg.pop("swanlab", {}))
|
||||
# 不是 ultralytics 的参数,得先摘出来,否则它会拒绝未知键
|
||||
max_overlap = float(train_cfg.pop("max_overlap", 0.0))
|
||||
|
||||
from ultralytics import YOLO
|
||||
from ultralytics import settings as ul_settings
|
||||
|
||||
if max_overlap >= 0:
|
||||
install_overlap_suppression(max_overlap)
|
||||
logger.info("验证时约束预测框互不重叠(最大 IoU %.2f)", max_overlap)
|
||||
else:
|
||||
logger.warning("未约束预测框重叠:指标里会计入真值中不可能出现的重叠框")
|
||||
|
||||
# ultralytics 默认把数据集根指到 ~/datasets 之类的全局目录;我们的 data.yaml 里写的是
|
||||
# 绝对路径,但它仍会拿 datasets_dir 去拼相对路径,统一指到项目根省得踩坑。
|
||||
ul_settings.update({"datasets_dir": str(Path.cwd())})
|
||||
|
||||
@@ -140,6 +140,46 @@ def _inside(inner: Box, outer: Box) -> bool:
|
||||
)
|
||||
|
||||
|
||||
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 图像并返回)。
|
||||
|
||||
@@ -10,6 +10,7 @@ from engine.detection import ( # noqa: E402
|
||||
COLOR_UNMATCHED,
|
||||
Box,
|
||||
draw_boxes,
|
||||
filter_overlapping_outputs,
|
||||
from_yolo,
|
||||
iou,
|
||||
match_boxes,
|
||||
@@ -190,3 +191,57 @@ def test_max_iou_allows_slight_overlap():
|
||||
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])
|
||||
|
||||
Reference in New Issue
Block a user