From eba533a5b2009f46ff3937d9526b162c7fad4251 Mon Sep 17 00:00:00 2001 From: Shaun Zhou Date: Fri, 18 Sep 2026 11:19:47 +0800 Subject: [PATCH] =?UTF-8?q?feat(detect):=20=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 人工标注的真值里一处缺陷只有一个框：11211 张图、151 张多框图，重叠框对为 0。 而模型预测在验证集上有 6.9% 的图存在重叠框——NMS 只在同类内、按固定 IoU 阈值抑制， 压不掉两类情况：同一处缺陷被两个类别各框一次（实测 IoU 高达 0.99），以及 IoU 低于 NMS 阈值的同类重叠（grease 66 对、wrinkle 45 对）。 - engine/detection.suppress_overlaps：按置信度贪心保留互不重叠的框，可设允许的 最大 IoU；另可丢弃被大框完全包住的小框（面积悬殊时 IoU 很小，但显然是重复表达）。 - visualize_detections 默认 --max-overlap 0（与真值一致），并放出 --nms-iou 与 类别无关 NMS 开关；--max-overlap -1 可看模型原始输出。 Co-Authored-By: Claude Opus 5 --- engine/detection.py | 38 ++++++++++++++++++++ scripts/visualize_detections.py | 44 ++++++++++++++++++++++-- tests/test_detection.py | 61 +++++++++++++++++++++++++++++++++ 3 files changed, 140 insertions(+), 3 deletions(-) diff --git a/engine/detection.py b/engine/detection.py index 3fd01f3..3240ddc 100644 --- a/engine/detection.py +++ b/engine/detection.py @@ -102,6 +102,44 @@ def match_boxes( ) +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 draw_boxes(image, ground_truth: Sequence[Box], predictions: Sequence[Box], classes: Sequence[str], match: Match, width: int = 2): """把真值与预测画在图上（就地修改传入的 PIL 图像并返回）。 diff --git a/scripts/visualize_detections.py b/scripts/visualize_detections.py index a4d8a3a..b4a7d4b 100644 --- a/scripts/visualize_detections.py +++ b/scripts/visualize_detections.py @@ -24,6 +24,10 @@ runs// ``` `--only errors` 只输出有漏检或误检的图，扫起来快得多。 + +**预测框默认约束为互不重叠**（`--max-overlap 0`）。人工标注的真值里一处缺陷只有一个框、 +零重叠，而 NMS 只在同类内按固定阈值抑制，留得下跨类重复和 IoU 低于阈值的重叠框。 +`--max-overlap -1` 可关掉这层约束，看模型的原始输出。 """ from __future__ import annotations @@ -41,7 +45,13 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[1])) import yaml # noqa: E402 from datasets.yolo import IMAGE_SUFFIXES # noqa: E402 -from engine.detection import Box, draw_boxes, match_boxes, parse_label # noqa: E402 +from engine.detection import ( # noqa: E402 + Box, + draw_boxes, + match_boxes, + parse_label, + suppress_overlaps, +) from utils import setup_logging # noqa: E402 from utils.device import DeviceNotAvailableError, describe_device, resolve_device # noqa: E402 @@ -110,6 +120,16 @@ def main() -> int: parser.add_argument("--imgsz", type=int, default=640) parser.add_argument("--conf", type=float, default=0.25, help="置信度阈值") parser.add_argument("--iou-threshold", type=float, default=0.5, help="判定命中的 IoU") + parser.add_argument( + "--max-overlap", type=float, default=0.0, + help="预测框之间允许的最大 IoU，超出的丢低分框。0=完全不允许重叠（默认，与真值一致），" + "-1=不约束，看模型原始输出", + ) + parser.add_argument("--nms-iou", type=float, default=0.7, help="NMS 的 IoU 阈值") + parser.add_argument( + "--no-agnostic-nms", action="store_true", + help="关掉类别无关 NMS。默认开启，用来压掉同一处缺陷被两个类别各框一次的重复", + ) parser.add_argument("--batch", type=int, default=16) parser.add_argument( "--only", choices=("all", "errors"), default="all", @@ -137,7 +157,12 @@ def main() -> int: "数据集: %s（%d 类，%s）", args.data, len(classes), "，".join(f"{name} {len(images)} 张" for name, images in splits.items()), ) - logger.info("conf=%.2f IoU=%.2f imgsz=%d", args.conf, args.iou_threshold, args.imgsz) + logger.info( + "conf=%.2f 命中 IoU=%.2f imgsz=%d NMS IoU=%.2f 类别无关 NMS=%s 框间最大重叠=%s", + args.conf, args.iou_threshold, args.imgsz, args.nms_iou, + not args.no_agnostic_nms, + "不约束" if args.max_overlap < 0 else f"{args.max_overlap:.2f}", + ) if out_dir.exists() and any(out_dir.iterdir()): if not args.overwrite: @@ -158,6 +183,7 @@ def main() -> int: scores: dict[str, list[float]] = defaultdict(list) rows = [] processed = 0 + suppressed = 0 def predict_in_batches(images): """自己按批切分：ultralytics 的 predict 不理会 batch 参数， @@ -166,7 +192,8 @@ def main() -> int: chunk = images[start:start + args.batch] results = model.predict( source=[str(p) for p in chunk], imgsz=args.imgsz, conf=args.conf, - device=device_arg, verbose=False, + device=device_arg, verbose=False, iou=args.nms_iou, + agnostic_nms=not args.no_agnostic_nms, ) yield from zip(chunk, results, strict=True) @@ -179,6 +206,11 @@ def main() -> int: result.boxes.conf.tolist(), strict=True, ) ] + if args.max_overlap >= 0: + kept = suppress_overlaps(predictions, max_iou=args.max_overlap) + suppressed += len(predictions) - len(kept) + predictions = [predictions[i] for i in kept] + label_file = label_path_for(image_path) truth = parse_label(label_file.read_text(encoding="utf-8")) if label_file.exists() else [] match = match_boxes(predictions, truth, args.iou_threshold) @@ -225,6 +257,10 @@ def main() -> int: "conf": args.conf, "iou_threshold": args.iou_threshold, "imgsz": args.imgsz, + "nms_iou": args.nms_iou, + "agnostic_nms": not args.no_agnostic_nms, + "max_overlap": args.max_overlap, + "suppressed_overlapping_boxes": suppressed, "outcomes": {split: dict(counter) for split, counter in outcomes.items()}, "per_class": { split: { @@ -241,6 +277,8 @@ def main() -> int: json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8" ) + if suppressed: + logger.info("因与更高分的框重叠而丢弃: %d 个", suppressed) for split, counter in outcomes.items(): shown = sum(counter[k] for k in ("missed", "extra")) logger.info( diff --git a/tests/test_detection.py b/tests/test_detection.py index 4e1e4f1..b06900c 100644 --- a/tests/test_detection.py +++ b/tests/test_detection.py @@ -14,6 +14,7 @@ from engine.detection import ( # noqa: E402 iou, match_boxes, parse_label, + suppress_overlaps, ) @@ -129,3 +130,63 @@ def test_draw_handles_boxes_at_the_border(): 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]