feat(detect): 全量推理可视化，预测框与真值框画在一起
scripts/visualize_detections.py：在整个数据集上推理，绿=真值、蓝=命中的预测、 红=没对上真值的预测，按结果类型分目录（ok / missed / extra / none），再配一份 index.csv（逐图的真值数、预测数、命中数、最高置信度）与 summary.json（逐类 TP/FP/FN）。 --only errors 只输出有问题的图。 匹配逻辑放在 engine/detection.py：按置信度贪心、IoU 阈值、类别需一致，与 mAP 的规则 一致，所以图上看到的漏检误检和指标算的是同一回事。same_class=False 可以单看定位准不准 ——定位对但分类错，和完全没找到是两种问题。 train 与 val 分开放：训练集的图模型见过，那部分结果不能当泛化能力看。 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
co-authored by
Claude Opus 5
Claude Opus 5
parent
ad50c38f41
commit
788bc36ab8
@@ -0,0 +1,257 @@
|
||||
#!/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
|
||||
|
||||
for split, images in splits.items():
|
||||
results = model.predict(
|
||||
source=[str(p) for p in images], stream=True, imgsz=args.imgsz, conf=args.conf,
|
||||
device=device_arg, batch=args.batch, verbose=False,
|
||||
)
|
||||
for image_path, result in zip(images, results, strict=True):
|
||||
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())
|
||||
Reference in New Issue
Block a user