Shaun ZhouandClaude Opus 5
43c8c4a179
model.predict 不理会传入的 batch 参数，把整个图片列表当成一个 batch，8969 张图在 16 GiB 显卡上直接 OOM（单次要分配 10.27 GiB）。改成自己按 --batch 切片调用。 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
265 lines
11 KiB
Python
265 lines
11 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 阈值、类别需一致），
|
||
所以图上看到的和指标算的是同一回事。
|
||
|
||
产出按"结果类型"分目录，直接对着看失败模式：
|
||
|
||
```
|
||
runs/<name>/
|
||
├── train/ val/ # 分开放：训练集的图模型见过，指标不可信
|
||
│ ├── ok/<图名>.png # 全部命中，没有多余预测
|
||
│ ├── missed/<图名>.png # 有漏检
|
||
│ ├── extra/<图名>.png # 有误检（含背景图上凭空框出的）
|
||
│ └── none/<图名>.png # 既无真值也无预测——负样本判对了
|
||
├── index.csv # 每张图一行：真值数/预测数/命中数/最高置信度
|
||
└── summary.json # 逐类 TP/FP/FN 与各结果类型的张数
|
||
```
|
||
|
||
`--only errors` 只输出有漏检或误检的图，扫起来快得多。
|
||
"""
|
||
|
||
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 Box, draw_boxes, match_boxes, parse_label # noqa: E402
|
||
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 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("--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", args.conf, args.iou_threshold, args.imgsz)
|
||
|
||
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
|
||
|
||
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,
|
||
)
|
||
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,
|
||
)
|
||
]
|
||
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 / split / kind / 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,
|
||
"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"
|
||
)
|
||
|
||
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())
|