PokemonCardSearch/code/pokemon/pokemon_grading.py

# -*- coding: utf-8 -*-
"""
pokemon_grading.py - best.pt 在 images_pokemon 上跑: 裁评级标签条 + 标注图 + 预测记录

输出:
  {OUT}/crops/<folder>_<name>.jpg       评级标签条裁剪图
  {OUT}/predict/<folder>/<name>.jpg     best.pt 标注图(带框+标签+置信度)
  {OUT}/crop_records.json               裁剪记录(box, conf)
  {OUT}/test_predictions.json           best.pt 预测(true, pred, conf, file)
  {OUT}/grading_summary.csv             汇总(filename, true_label, detected_company, is_graded, conf, box)
运行环境: pytorch (~/miniconda3/envs/pytorch/bin/python)
"""
import os
import sys
import json
import csv
import re
import time
from pathlib import Path
import numpy as np
import cv2

sys.stdout.reconfigure(encoding="utf-8")

ROOT = "/home/martin/顾工交接/wzj/images_pokemon"
OUT = "/home/martin/顾工交接/wzj/images_pokemon_out"
BEST = "/home/martin/顾工交接/wzj/ultralytics/runs/detect/card_v1/weights/best.pt"
FOLDERS = ["BGS", "CGC", "PSA", "SGC"]
CONF = 0.25

os.makedirs(os.path.join(OUT, "crops"), exist_ok=True)
os.makedirs(os.path.join(OUT, "predict"), exist_ok=True)
for f in FOLDERS:
    os.makedirs(os.path.join(OUT, "predict", f), exist_ok=True)

print("[STEP] 加载 best.pt ...")
sys.path.insert(0, "/home/martin/顾工交接/wzj")
from modules.ultralytics_compat import import_yolo
YOLO = import_yolo()
_t_load0 = time.perf_counter()
model = YOLO(BEST, task="detect")
t_load = time.perf_counter() - _t_load0
print(f"[STEP] 加载完成 (耗时 {t_load:.2f}s)\n")

crop_records = []
pred_records = []
csv_rows = []

# 计时累加器
T = {"infer": 0.0, "imread": 0.0, "crop_write": 0.0, "annot_write": 0.0}
n_img = 0
_t_all0 = time.perf_counter()

for fld in FOLDERS:
    imgs = sorted(Path(os.path.join(ROOT, fld)).glob("*.jpg"))
    print(f"[STEP] {fld}: {len(imgs)} 张", flush=True)
    for imgp in imgs:
        name = imgp.name
        n_img += 1
        _t = time.perf_counter()
        results = model.predict(str(imgp), conf=CONF, imgsz=640, verbose=False)
        T["infer"] += time.perf_counter() - _t
        box = None
        conf_val = 0.0
        pred_label = "(无检测)"
        if results and results[0].boxes is not None and len(results[0].boxes) > 0:
            confs = results[0].boxes.conf.cpu().numpy()
            clses = results[0].boxes.cls.cpu().numpy().astype(int)
            xyxy = results[0].boxes.xyxy.cpu().numpy()
            best_i = int(np.argmax(confs))
            conf_val = float(confs[best_i])
            box = tuple(float(v) for v in xyxy[best_i])
            pred_label = model.names[clses[best_i]]

        # 公司前缀(从 label 取,如 BGS-AUTHENTIC -> BGS)
        detected_company = pred_label.split("-", 1)[0] if pred_label != "(无检测)" else ""
        is_graded = 1 if detected_company else 0

        # 裁剪评级标签条
        _t = time.perf_counter()
        img = cv2.imread(str(imgp))
        T["imread"] += time.perf_counter() - _t
        crop = None
        if box:
            x1, y1, x2, y2 = box
            h, w = img.shape[:2]
            x1i, y1i = max(0, int(x1)), max(0, int(y1))
            x2i, y2i = min(w, int(x2)), min(h, int(y2))
            if x2i > x1i and y2i > y1i:
                crop = img[y1i:y2i, x1i:x2i]
        if crop is None:
            crop = img  # 无检测则用原图(后续 OCR 兜底)
        crop_path = os.path.join(OUT, "crops", f"{fld}_{name}")
        _t = time.perf_counter()
        cv2.imwrite(crop_path, crop)
        T["crop_write"] += time.perf_counter() - _t

        # 标注图(带框+标签+置信度)
        _t = time.perf_counter()
        if results and results[0].boxes is not None and len(results[0].boxes) > 0:
            annotated = results[0].plot()  # ultralytics 自带标注
        else:
            annotated = img
        cv2.imwrite(os.path.join(OUT, "predict", fld, name), annotated)
        T["annot_write"] += time.perf_counter() - _t

        crop_records.append({
            "file": name, "true_label": fld, "box": box,
            "conf": round(conf_val, 4) if conf_val else None,
            "crop_path": crop_path,
        })
        pred_records.append({
            "true": fld, "pred": pred_label,
            "conf": round(conf_val, 4) if conf_val else 0.0, "file": name,
        })
        csv_rows.append({
            "filename": name, "true_label": fld,
            "detected_company": detected_company,
            "is_graded": is_graded,
            "conf": round(conf_val, 4) if conf_val else "",
            "box": str(box) if box else "",
        })

print("\n[STEP] 保存记录 ...")
with open(os.path.join(OUT, "crop_records.json"), "w", encoding="utf-8") as f:
    json.dump(crop_records, f, ensure_ascii=False, indent=2)
with open(os.path.join(OUT, "test_predictions.json"), "w", encoding="utf-8") as f:
    json.dump({"records": pred_records}, f, ensure_ascii=False, indent=2)
with open(os.path.join(OUT, "grading_summary.csv"), "w", encoding="utf-8-sig", newline="") as f:
    w = csv.DictWriter(f, fieldnames=["filename", "true_label", "detected_company", "is_graded", "conf", "box"])
    w.writeheader()
    w.writerows(csv_rows)

# 准确率统计(仅评级卡)
graded = [r for r in csv_rows if r["is_graded"] == 1]
correct = sum(1 for r in graded if r["detected_company"] == r["true_label"])
print(f"\n[RESULT] 评级卡 {len(graded)} 张, 非评级卡 {len(csv_rows)-len(graded)} 张")
print(f"[RESULT] 公司识别准确率(仅评级卡): {correct}/{len(graded)} = {correct/len(graded)*100:.2f}%")

# ===== 计时报告 =====
t_all = time.perf_counter() - _t_all0
rep = []
rep.append("=" * 50)
rep.append("⏱  计时报告 (pokemon_grading.py)")
rep.append("=" * 50)
rep.append(f"  模型加载(best.pt)     : {t_load:7.2f}s")
rep.append(f"  best.pt 推理(累计)    : {T['infer']:7.2f}s  | 均 {T['infer']/n_img*1000:6.1f} ms/张")
rep.append(f"  读图 cv2.imread(累计) : {T['imread']:7.2f}s  | 均 {T['imread']/n_img*1000:6.1f} ms/张")
rep.append(f"  写裁剪图(累计)        : {T['crop_write']:7.2f}s  | 均 {T['crop_write']/n_img*1000:6.1f} ms/张")
rep.append(f"  写标注图(累计)        : {T['annot_write']:7.2f}s  | 均 {T['annot_write']/n_img*1000:6.1f} ms/张")
rep.append(f"  {'-'*46}")
rep.append(f"  处理总耗时({n_img}张)  : {t_all:7.2f}s  | 均 {t_all/n_img*1000:6.1f} ms/张  | 吞吐 {n_img/t_all:.1f} 张/s")
rep.append("=" * 50)
print("\n" + "\n".join(rep))

# 写 timing.log (追加, 带时间戳)
import datetime
with open(os.path.join(OUT, "timing.log"), "a", encoding="utf-8") as lf:
    lf.write(f"\n[{datetime.datetime.now():%Y-%m-%d %H:%M:%S}] === 步骤① pokemon_grading.py (评级检测 best.pt) ===\n")
    lf.write(f"  样本数: {n_img} 张\n")
    lf.write(f"  模型加载 best.pt          : {t_load:.2f}s\n")
    lf.write(f"  best.pt 推理累计          : {T['infer']:.2f}s  (均 {T['infer']/n_img*1000:.1f} ms/张)\n")
    lf.write(f"  读图 cv2.imread 累计      : {T['imread']:.2f}s  (均 {T['imread']/n_img*1000:.1f} ms/张)\n")
    lf.write(f"  写裁剪图累计              : {T['crop_write']:.2f}s  (均 {T['crop_write']/n_img*1000:.1f} ms/张)\n")
    lf.write(f"  写标注图累计              : {T['annot_write']:.2f}s  (均 {T['annot_write']/n_img*1000:.1f} ms/张)\n")
    lf.write(f"  处理总耗时                : {t_all:.2f}s  (均 {t_all/n_img*1000:.1f} ms/张, 吞吐 {n_img/t_all:.1f} 张/s)\n")
print(f"\n✓ 完成. 输出目录: {OUT}")
print(f"✓ 计时已写入: {os.path.join(OUT, 'timing.log')}")