From 9439cbb56a27101f9139fe2cfad08e095f7341d3 Mon Sep 17 00:00:00 2001 From: Shaun Zhou Date: Fri, 18 Sep 2026 11:30:48 +0800 Subject: [PATCH] =?UTF-8?q?feat(detect):=20=E8=AE=AD=E7=BB=83=E6=9C=9F?= =?UTF-8?q?=E7=9A=84=E9=AA=8C=E8=AF=81=E4=B9=9F=E7=BA=A6=E6=9D=9F=E9=A2=84?= =?UTF-8?q?=E6=B5=8B=E6=A1=86=E4=BA=92=E4=B8=8D=E9=87=8D=E5=8F=A0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- configs/detect/film.yaml | 7 +++++ detect.py | 33 ++++++++++++++++++++++-- engine/detection.py | 40 +++++++++++++++++++++++++++++ tests/test_detection.py | 55 ++++++++++++++++++++++++++++++++++++++++ 4 files changed, 133 insertions(+), 2 deletions(-) diff --git a/configs/detect/film.yaml b/configs/detect/film.yaml index 1d87720..cd90ee3 100644 --- a/configs/detect/film.yaml +++ b/configs/detect/film.yaml @@ -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 diff --git a/detect.py b/detect.py index 2b33e14..ef33020 100644 --- a/detect.py +++ b/detect.py @@ -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())}) diff --git a/engine/detection.py b/engine/detection.py index 3240ddc..bb5defa 100644 --- a/engine/detection.py +++ b/engine/detection.py @@ -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 图像并返回）。 diff --git a/tests/test_detection.py b/tests/test_detection.py index b06900c..1532354 100644 --- a/tests/test_detection.py +++ b/tests/test_detection.py @@ -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])