Shaun ZhouandClaude Opus 5
36fff984b7
原先按 split 和结果类型（ok/missed/extra/none）分目录，层级太深。改成 <类别>/<图名>.png，与 runs/boxes_shap 的结构一致。类别取真值里框最多的那个； 负样本没有真值框，按文件名前缀归类（<类别>_<序号> 是 prepare_label_studio 的命名）， 认不出来放 background/。 split 与结果类型仍记在 index.csv 的列里，排序筛选即可。 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
319 lines
14 KiB
Python
319 lines
14 KiB
Python
#!/usr/bin/env python
|
||
"""用训练好的检测模型在整个数据集上推理，把预测框与真值框画在一起供人工核对。
|
||
|
||
uv run python scripts/visualize_detections.py \\
|
||
--weights runs/detect/film/weights/best.pt \\
|
||
--data data/film_yolo/data.yaml \\
|
||
--out runs/detect/film_vis
|
||
|
||
图上三种颜色：**绿=真值**，**蓝=命中的预测**，**红=没对上真值的预测**（误检）。
|
||
只有绿框没有蓝框，就是漏检。匹配规则与 mAP 一致（按置信度贪心、IoU 阈值、类别需一致），
|
||
所以图上看到的和指标算的是同一回事。
|
||
|
||
产出按**类别**分目录，与 `generate_boxes.py` 的产出结构一致：
|
||
|
||
```
|
||
runs/<name>/
|
||
├── <类别>/<图名>.png # 画好框的图。类别取真值里框最多的那个；
|
||
│ # 没有真值框时按文件名前缀归类（<类别>_<序号> 是
|
||
│ # prepare_label_studio.py 的命名），认不出来就放 background/
|
||
├── index.csv # 每张图一行：split、结果类型、真值数/预测数/命中数/最高置信度
|
||
└── summary.json # 逐类 TP/FP/FN 与各结果类型的张数
|
||
```
|
||
|
||
train / val 不分目录，但 `index.csv` 里有 `split` 列——训练集的图模型见过，看结果时要留意。
|
||
同样地，"命中 / 漏检 / 误检"记在 `index.csv` 的 `outcome` 列里，按它排序筛选即可。
|
||
|
||
`--only errors` 只输出有漏检或误检的图，扫起来快得多。
|
||
|
||
**预测框默认约束为互不重叠**（`--max-overlap 0`）。人工标注的真值里一处缺陷只有一个框、
|
||
零重叠，而 NMS 只在同类内按固定阈值抑制，留得下跨类重复和 IoU 低于阈值的重叠框。
|
||
`--max-overlap -1` 可关掉这层约束，看模型的原始输出。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import csv
|
||
import json
|
||
import shutil
|
||
import sys
|
||
from collections import Counter, defaultdict
|
||
from pathlib import Path
|
||
|
||
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 ( # 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
|
||
|
||
logger = None
|
||
|
||
OUTCOMES = ("ok", "missed", "extra", "none")
|
||
|
||
|
||
def outcome(n_gt: int, n_pred: int, match) -> str:
|
||
"""一张图的结果类型。missed 优先于 extra：漏检比误检更该先看。"""
|
||
if n_gt == 0 and n_pred == 0:
|
||
return "none"
|
||
if match.unmatched_ground_truth:
|
||
return "missed"
|
||
if match.unmatched_predictions:
|
||
return "extra"
|
||
return "ok"
|
||
|
||
|
||
def class_dir_for(truth: list[Box], classes: list[str], stem: str) -> str:
|
||
"""一张图归到哪个类别目录。
|
||
|
||
有真值框时取框最多的那个类别；没有框（负样本）时退回文件名前缀——
|
||
数据集里的图名是 `<类别>_<序号>`，前缀就是人工归类的类别；认不出来就放 background。
|
||
"""
|
||
if truth:
|
||
return classes[Counter(box.class_id for box in truth).most_common(1)[0][0]]
|
||
prefix = stem.rsplit("_", 1)[0]
|
||
return prefix if prefix in classes else "background"
|
||
|
||
|
||
def load_dataset(data_yaml: Path) -> tuple[list[str], dict[str, list[Path]]]:
|
||
"""从 data.yaml 读类别与各 split 的图片列表。"""
|
||
meta = yaml.safe_load(data_yaml.read_text(encoding="utf-8"))
|
||
names = meta["names"]
|
||
classes = [names[i] for i in sorted(names)] if isinstance(names, dict) else list(names)
|
||
root = Path(meta.get("path", data_yaml.parent))
|
||
if not root.is_absolute():
|
||
root = (data_yaml.parent / root).resolve()
|
||
|
||
splits = {}
|
||
for split in ("train", "val"):
|
||
entry = meta.get(split)
|
||
if not entry:
|
||
continue
|
||
directory = root / entry
|
||
if not directory.is_dir():
|
||
raise FileNotFoundError(f"{data_yaml} 里的 {split} 指向 {directory}，但它不存在")
|
||
images = sorted(p for p in directory.iterdir() if p.suffix.lower() in IMAGE_SUFFIXES)
|
||
# train 与 val 指向同一目录时（未划分的数据集）只跑一次
|
||
if any(directory == root / meta.get(other, "") for other in splits):
|
||
continue
|
||
splits[split] = images
|
||
if not splits:
|
||
raise FileNotFoundError(f"{data_yaml} 里没有可用的 train / val")
|
||
return classes, splits
|
||
|
||
|
||
def label_path_for(image: Path) -> Path:
|
||
"""YOLO 约定：把路径里的 images/ 换成 labels/，后缀换成 .txt。"""
|
||
parts = list(image.parts)
|
||
for index in range(len(parts) - 1, -1, -1):
|
||
if parts[index] == "images":
|
||
parts[index] = "labels"
|
||
break
|
||
return Path(*parts).with_suffix(".txt")
|
||
|
||
|
||
def main() -> int:
|
||
global logger
|
||
parser = argparse.ArgumentParser(
|
||
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
|
||
)
|
||
parser.add_argument("--weights", required=True, help="训练好的 .pt")
|
||
parser.add_argument("--data", required=True, help="数据集的 data.yaml")
|
||
parser.add_argument("--out", required=True, help="产出目录，建议放在 runs/ 下")
|
||
parser.add_argument("--device", default="auto", help="auto / cuda / mps / cpu")
|
||
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",
|
||
help="errors=只输出有漏检或误检的图，统计仍按全量计",
|
||
)
|
||
parser.add_argument("--limit", type=int, default=0, help="每个 split 只跑前 N 张（试参数用）")
|
||
parser.add_argument("--overwrite", action="store_true", help="产出目录已存在时先清空")
|
||
args = parser.parse_args()
|
||
|
||
out_dir = Path(args.out)
|
||
logger = setup_logging()
|
||
|
||
try:
|
||
device = resolve_device(args.device)
|
||
classes, splits = load_dataset(Path(args.data))
|
||
except (DeviceNotAvailableError, FileNotFoundError) as exc:
|
||
logger.error("%s", exc)
|
||
return 1
|
||
|
||
if args.limit:
|
||
splits = {name: images[: args.limit] for name, images in splits.items()}
|
||
total = sum(len(images) for images in splits.values())
|
||
logger.info("设备: %s | 权重: %s", describe_device(device), args.weights)
|
||
logger.info(
|
||
"数据集: %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 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:
|
||
logger.error("%s 已存在且非空；确认要覆盖请加 --overwrite", out_dir)
|
||
return 1
|
||
shutil.rmtree(out_dir)
|
||
|
||
from PIL import Image
|
||
from ultralytics import YOLO
|
||
|
||
model = YOLO(args.weights)
|
||
device_arg = device.index or 0 if device.type == "cuda" else device.type
|
||
|
||
outcomes: dict[str, Counter[str]] = {name: Counter() for name in splits}
|
||
stats: dict[str, dict[str, Counter[str]]] = {
|
||
name: {"tp": Counter(), "fp": Counter(), "fn": Counter()} for name in splits
|
||
}
|
||
scores: dict[str, list[float]] = defaultdict(list)
|
||
rows = []
|
||
processed = 0
|
||
suppressed = 0
|
||
|
||
def predict_in_batches(images):
|
||
"""自己按批切分：ultralytics 的 predict 不理会 batch 参数，
|
||
一次性把整个列表当成一个 batch，上万张图必然 OOM。"""
|
||
for start in range(0, len(images), args.batch):
|
||
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, iou=args.nms_iou,
|
||
agnostic_nms=not args.no_agnostic_nms,
|
||
)
|
||
yield from zip(chunk, results, strict=True)
|
||
|
||
for split, images in splits.items():
|
||
for image_path, result in predict_in_batches(images):
|
||
predictions = [
|
||
Box(int(cls), *(float(v) for v in xyxy), float(score))
|
||
for xyxy, cls, score in zip(
|
||
result.boxes.xyxyn.tolist(), result.boxes.cls.tolist(),
|
||
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)
|
||
|
||
kind = outcome(len(truth), len(predictions), match)
|
||
outcomes[split][kind] += 1
|
||
for _, truth_index in match.pairs:
|
||
stats[split]["tp"][classes[truth[truth_index].class_id]] += 1
|
||
for index in match.unmatched_predictions:
|
||
stats[split]["fp"][classes[predictions[index].class_id]] += 1
|
||
for index in match.unmatched_ground_truth:
|
||
stats[split]["fn"][classes[truth[index].class_id]] += 1
|
||
for box in predictions:
|
||
scores[classes[box.class_id]].append(box.score)
|
||
|
||
rows.append({
|
||
"split": split, "name": image_path.name, "outcome": kind,
|
||
"n_gt": len(truth), "n_pred": len(predictions), "n_matched": len(match.pairs),
|
||
"max_conf": round(max((b.score for b in predictions), default=0.0), 4),
|
||
"gt_classes": "|".join(sorted({classes[b.class_id] for b in truth})),
|
||
"pred_classes": "|".join(sorted({classes[b.class_id] for b in predictions})),
|
||
})
|
||
|
||
if args.only == "all" or kind in ("missed", "extra"):
|
||
destination = (
|
||
out_dir / class_dir_for(truth, classes, image_path.stem)
|
||
/ f"{image_path.stem}.png"
|
||
)
|
||
destination.parent.mkdir(parents=True, exist_ok=True)
|
||
with Image.open(image_path) as raw:
|
||
canvas = raw.convert("RGB")
|
||
draw_boxes(canvas, truth, predictions, classes, match)
|
||
canvas.save(destination)
|
||
|
||
processed += 1
|
||
if processed % 1000 == 0:
|
||
logger.info("处理 [%d/%d]", processed, total)
|
||
|
||
with (out_dir / "index.csv").open("w", encoding="utf-8-sig", newline="") as handle:
|
||
writer = csv.DictWriter(handle, fieldnames=list(rows[0]))
|
||
writer.writeheader()
|
||
writer.writerows(rows)
|
||
|
||
summary = {
|
||
"weights": str(Path(args.weights).resolve()),
|
||
"data": str(Path(args.data).resolve()),
|
||
"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: {
|
||
name: {
|
||
"tp": counters["tp"][name], "fp": counters["fp"][name],
|
||
"fn": counters["fn"][name],
|
||
}
|
||
for name in classes
|
||
}
|
||
for split, counters in stats.items()
|
||
},
|
||
}
|
||
(out_dir / "summary.json").write_text(
|
||
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(
|
||
"%s：全部命中 %d，漏检 %d，误检 %d，负样本判对 %d（有问题的 %d 张，占 %.1f%%）",
|
||
split, counter["ok"], counter["missed"], counter["extra"], counter["none"],
|
||
shown, 100 * shown / max(sum(counter.values()), 1),
|
||
)
|
||
logger.info("%-12s %8s %8s %8s %8s", "类别", "TP", "FP", "FN", "召回")
|
||
for name in classes:
|
||
tp, fp, fn = (stats[split][k][name] for k in ("tp", "fp", "fn"))
|
||
if tp or fp or fn:
|
||
logger.info(
|
||
"%-12s %8d %8d %8d %7.1f%%",
|
||
name, tp, fp, fn, 100 * tp / max(tp + fn, 1),
|
||
)
|
||
logger.info("产出目录: %s（index.csv 可按 outcome / max_conf 排序筛选）", out_dir)
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|