xzhou/ModelTrainEval
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(detect): 约束预测框互不重叠

Browse Source 
人工标注的真值里一处缺陷只有一个框：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 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
2026-09-18 11:19:47 +08:00
co-authored by Claude Opus 5
parent 43c8c4a179
commit eba533a5b2
3 changed files with 140 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
engine/detection.py
+38
View File
@@ -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 图像并返回）。
scripts/visualize_detections.py
+41 -3
View File
@@ -24,6 +24,10 @@ runs/<name>/
```
`--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(
tests/test_detection.py
+61
View File
@@ -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]