# -*- coding: utf-8 -*-
"""在测试集上推理并按真实类别(文件夹名)统计识别结果，同时保存标注图。"""
import glob
import json
import os
from collections import Counter

from ultralytics import YOLO

ROOT = "/home/martin/顾工交接/wzj/test"          # 4 个子文件夹 BGS/CGC/PSA/SGC
OUT = "/home/martin/顾工交接/wzj/test_predict"   # 标注图输出(保留子目录)
BEST = "/home/martin/顾工交接/wzj/ultralytics/runs/detect/card_v1/weights/best.pt"
CONF = 0.25
FOLDERS = ["BGS", "CGC", "PSA", "SGC"]

model = YOLO(BEST)
matrix = {f: Counter() for f in FOLDERS}
total = {f: 0 for f in FOLDERS}
records = []   # 每张图: {true, pred, conf, file}

for f in FOLDERS:
    imgs = []
    for ext in ("*.jpg", "*.jpeg", "*.png", "*.JPG", "*.JPEG", "*.PNG"):
        imgs += glob.glob(os.path.join(ROOT, f, ext))
    imgs = sorted(set(imgs))
    total[f] = len(imgs)
    print(f"推理 {f}: {len(imgs)} 张 ...", flush=True)
    outdir = os.path.join(OUT, f)
    os.makedirs(outdir, exist_ok=True)
    for imgpath in imgs:
        r = model(imgpath, conf=CONF, verbose=False, imgsz=640)[0]
        name = os.path.basename(imgpath)
        if r.boxes is None or len(r.boxes) == 0:
            matrix[f]["(无检测)"] += 1
            records.append({"true": f, "pred": "(无检测)", "conf": 0.0, "file": name})
            continue
        i = int(r.boxes.conf.argmax())
        pred = model.names[int(r.boxes.cls[i])]
        conf = float(r.boxes.conf[i])
        matrix[f][pred] += 1
        records.append({"true": f, "pred": pred, "conf": round(conf, 4), "file": name})
        r.save(filename=os.path.join(outdir, name))

# 汇总
print("\n================ 识别结果 ================")
tot = tot_correct_family = tot_correct_exact = 0
for f in FOLDERS:
    n = total[f]
    tot += n
    print(f"\n真实类别 {f} (共 {n} 张):")
    for pred, c in matrix[f].most_common():
        print(f"   预测 {pred:16} {c:4}  ({c / n * 100:5.1f}%)")
    fam = sum(c for p, c in matrix[f].items() if p.startswith(f))
    ex = matrix[f].get(f, 0)
    tot_correct_family += fam
    tot_correct_exact += ex
    print(f"   -> {f} 系列正确: {fam}/{n} = {fam / n * 100:.1f}%   完全相等: {ex}/{n} = {ex / n * 100:.1f}%")

print("\n================ 总体 ================")
print(f"主类(系列)正确: {tot_correct_family}/{tot} = {tot_correct_family / tot * 100:.2f}%")
print(f"完全相等:       {tot_correct_exact}/{tot} = {tot_correct_exact / tot * 100:.2f}%")

os.makedirs(OUT, exist_ok=True)
with open(os.path.join(OUT, "test_predictions.json"), "w", encoding="utf-8") as fp:
    json.dump({"totals": {f: total[f] for f in FOLDERS},
               "matrix": {f: dict(matrix[f]) for f in FOLDERS},
               "overall_family_acc": tot_correct_family / tot,
               "overall_exact_acc": tot_correct_exact / tot,
               "records": records}, fp, ensure_ascii=False, indent=2)
print(f"\n逐图预测明细已存: {OUT}/test_predictions.json")
print(f"标注图已存: {OUT}/<类别>/")