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
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(