中间件的生成与查看
1.标注文本转换
import os
import json
from PIL import Image
LABEL_DIR = r"D:\01-AI\20260914\dataset\images\images\YOLODataset\labels\val"
IMAGE_DIR = r"D:\01-AI\20260914\dataset\images\images\YOLODataset\images\val"
OUT_JSON = "ground_truth.json"
gt_list = []
for img_file in sorted(os.listdir(IMAGE_DIR)):
if not img_file.lower().endswith((".jpg", ".jpeg", ".png", ".bmp")):
continue
stem = os.path.splitext(img_file)[0]
label_path = os.path.join(LABEL_DIR, stem + ".txt")
if not os.path.exists(label_path):
continue
读图片尺寸
W, H = Image.open(os.path.join(IMAGE_DIR, img_file)).size
逐行读 YOLO 标注
for line in open(label_path):
cls_id, xc, yc, w, h = map(float, line.split())
反归一化:中心点 → 左上角
x = (xc - w / 2) * W
y = (yc - h / 2) * H
gt_list.append({
"image_id": stem,
"file_name": img_file,
"category_id": int(cls_id) + 1, # COCO 类别从 1 开始
"bbox": [round(x, 2), round(y, 2), round(w * W, 2), round(h * H, 2)]
})
json.dump(gt_list, open(OUT_JSON, "w", encoding="utf-8"),
ensure_ascii=False, indent=2)
print(f"完成:{len(gt_list)} 个框")
2.生成matrix
from ultralytics import YOLO
import pandas as pd
if name == 'main':
model = YOLO('runs/detect/train/weights/best.pt')
metrics = model.val(data=r'D:\01-AI\20260914\dataset\images\images\YOLODataset\dataset.yaml' )
获取混淆矩阵的原始数值
cm = metrics.confusion_matrix.matrix # numpy 数组,形状为 (nc+1, nc+1)
names = metrics.names # 类名列表
加上背景类,构建行列标签
labels = list(names) + ['background']
转为 DataFrame 并保存
df = pd.DataFrame(cm, index=labels, columns=labels)
df.to_csv('confusion_matrix.csv', index=True)
print("混淆矩阵已保存为 confusion_matrix.csv")
生成roc曲线
import json, csv
import numpy as np
import matplotlib.pyplot as plt
from collections import defaultdict
from sklearn.metrics import roc_curve, auc, precision_recall_curve, average_precision_score
def iou(b1, b2):
x1, y1 = max(b1[0], b2[0]), max(b1[1], b2[1])
x2, y2 = min(b1[2], b2[2]), min(b1[3], b2[3])
inter = max(0, x2 - x1) * max(0, y2 - y1)
a1 = (b1[2] - b1[0]) * (b1[3] - b1[1])
a2 = (b2[2] - b2[0]) * (b2[3] - b2[1])
return inter / (a1 + a2 - inter + 1e-6)
def xywh2xyxy(b):
x, y, w, h = b
return [x, y, x + w, y + h]
---------- 读数据 ----------
preds = json.load(open("runs/detect/val-2/predictions.json", encoding="utf-8"))
gts = json.load(open("ground_truth.json", encoding="utf-8"))
gt_group, pred_group = defaultdict(list), defaultdict(list)
for g in gts:
gt_group[(g["image_id"], g["category_id"])].append(xywh2xyxy(g["bbox"]))
for p in preds:
pred_group[(p["image_id"], p["category_id"])].append(p)
category_ids = sorted({g["category_id"] for g in gts})
---------- 匹配 & 统计 ----------
IOU_THRESH = 0.5
scores = {c: {"y": [], "s": []} for c in category_ids}
conf = defaultdict(lambda: defaultdict(int))
tp = defaultdict(int); fp = defaultdict(int); fn = defaultdict(int)
for (img_id, cat_id), plist in pred_group.items():
gboxes = gt_group.get((img_id, cat_id), [])
used = set()
for p in sorted(plist, key=lambda x: x["score"], reverse=True):
pbox = xywh2xyxy(p["bbox"])
best_iou, best_idx = 0, -1
for i, gbox in enumerate(gboxes):
if i not in used and (v := iou(pbox, gbox)) > best_iou:
best_iou, best_idx = v, i
if best_iou >= IOU_THRESH:
used.add(best_idx)
scores[cat_id]["y"].append(1); tp[cat_id] += 1
conf[cat_id][cat_id] += 1
else:
scores[cat_id]["y"].append(0); fp[cat_id] += 1
conf["background"][cat_id] += 1
scores[cat_id]["s"].append(p["score"])
for i in range(len(gboxes)):
if i not in used:
scores[cat_id]["y"].append(1); scores[cat_id]["s"].append(0.0)
fn[cat_id] += 1
conf[cat_id]["background"] += 1
---------- 指标 ----------
rows, roc_data, pr_data = [], {}, {}
for c in category_ids:
y, s = np.array(scores[c]["y"]), np.array(scores[c]["s"])
ok = len(np.unique(y)) > 1
ap = average_precision_score(y, s) if ok else (1.0 if y.all() else 0.0)
f1 = 2 * tp[c] / (2 * tp[c] + fp[c] + fn[c] + 1e-9) # 直接由 TP/FP/FN 得 F1
if ok:
fpr, tpr, _ = roc_curve(y, s)
roc_auc = auc(fpr, tpr)
roc_data[c] = (fpr, tpr, roc_auc)
prec, rec, _ = precision_recall_curve(y, s)
pr_data[c] = (rec, prec, ap)
else:
roc_auc = float("nan")
rows.append([c, f"{ap:.4f}", f"{f1:.4f}",
"nan" if np.isnan(roc_auc) else f"{roc_auc:.4f}"])
val.csv
with open("val.csv", "a+", newline="", encoding="utf-8") as f:
w = csv.writer(f); w.writerow(["class", "mAP@0.5", "F1", "AUC"]); w.writerows(rows)
valid = [r for r in rows if r[3] != "nan"]
w.writerow(["macro_avg"] + [f"{np.mean([float(r[i]) for r in (rows if i == 2 else valid)]):.4f}"
for i in (1, 2, 3)])
print("已保存 val.csv")
---------- 画图 ----------
def save_curve(data, xlabel, ylabel, title, fmt, fname):
plt.figure(figsize=(7, 6))
for c, d in data.items():
x, y = d[0], d[1]
plt.plot(x, y, label=f"class_{c} ({fmt}={d[2]:.3f})")
if fmt == "AUC":
plt.plot([0, 1], [0, 1], "k--")
plt.xlabel(xlabel); plt.ylabel(ylabel); plt.title(title)
plt.legend(fontsize=8); plt.grid(alpha=0.3); plt.tight_layout()
plt.savefig(fname, dpi=200); plt.close()
print(f"已生成 {fname}")
save_curve(roc_data, "FPR", "TPR", "ROC per Class (IoU=0.5)", "AUC", "roc.jpg")
save_curve(pr_data, "Recall", "Precision", "PR per Class (IoU=0.5)", "AP", "pr.jpg")

浙公网安备 33010602011771号