# -*- coding: utf-8 -*- """ pokemon_grading.py - best.pt 在 images_pokemon 上跑: 裁评级标签条 + 标注图 + 预测记录 输出: {OUT}/crops/_.jpg 评级标签条裁剪图 {OUT}/predict//.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')}")