PokemonCardSearch/code/pokemon/grading_ocr.py

# -*- coding: utf-8 -*-
"""
best.pt 裁评级标签条 + PaddleOCR 识别分行文本

流程:
  1. best.pt(yolo26) 在原图检测评级标签条 → 取置信度最高的框裁剪
  2. 对裁出的标签条跑 PaddleOCR(use_doc_orientation_classify=True)
  3. 输出:
     - crops/<原图名>.jpg: 裁出的标签条图
     - ocr_result.json: 每张图的 OCR 分行文本
     - ocr_summary.csv: 汇总表(文件名,评级框坐标,置信度,OCR分行文本)

运行环境: 先用 pytorch 环境跑裁剪, 再用 paddleocr 环境跑 OCR
"""
import os
import sys
import json
import csv
import hashlib
from pathlib import Path

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

ROOT = "/home/martin/顾工交接/wzj/test"           # 4个子目录 BGS/CGC/PSA/SGC
OUT = "/home/martin/顾工交接/wzj/grading_ocr"     # 输出目录
BEST = "/home/martin/顾工交接/wzj/ultralytics/runs/detect/card_v1/weights/best.pt"
FOLDERS = ["BGS", "CGC", "PSA", "SGC"]
CONF = 0.25

os.makedirs(OUT, exist_ok=True)
os.makedirs(os.path.join(OUT, "crops"), exist_ok=True)

# ============ 第一步: best.pt 裁剪 (pytorch 环境) ============
print("[STEP1] best.pt 裁评级标签条 ...")
sys.path.insert(0, "/home/martin/顾工交接/wzj")
from modules.ultralytics_compat import import_yolo
YOLO = import_yolo()
model = YOLO(BEST, task="detect")

crop_records = []  # [{file, true_label, box, conf, crop_path}, ...]
total_imgs = 0
for fld in FOLDERS:
    imgs = sorted(Path(os.path.join(ROOT, fld)).glob("*.jpg"))
    total_imgs += len(imgs)
    print(f"  {fld}: {len(imgs)} 张")
    for imgp in imgs:
        name = imgp.name
        results = model.predict(str(imgp), conf=CONF, imgsz=640, verbose=False)
        box, conf_val = None, 0.0
        if results and results[0].boxes is not None and len(results[0].boxes) > 0:
            import numpy as np
            confs = results[0].boxes.conf.cpu().numpy()
            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])
        # 裁剪
        import cv2
        img = cv2.imread(str(imgp))
        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  # 未检测到则用原图
        crop_path = os.path.join(OUT, "crops", f"{fld}_{name}")
        cv2.imwrite(crop_path, crop)
        crop_records.append({
            "file": name,
            "true_label": fld,
            "box": box,
            "conf": round(conf_val, 4) if conf_val else None,
            "crop_path": crop_path,
        })

print(f"[STEP1] 完成, 共处理 {total_imgs} 张, 裁剪图存于 {OUT}/crops/")

# ============ 第二步: PaddleOCR 识别 (需要 paddleocr 环境) ============
print("\n[STEP2] PaddleOCR 识别裁剪图 ...")
print("提示: 此步骤需要在 paddleocr conda 环境运行:")
print("  conda activate paddleocr")
print("  python grading_ocr.py --ocr-only")
print("\n当前脚本将只生成裁剪图, OCR 步骤请单独运行。")

# 保存裁剪记录供 OCR 步骤使用
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)
print(f"裁剪记录已保存: {OUT}/crop_records.json")

# ============ OCR 步骤(单独脚本) ============
# 见下方 grading_ocr_step2.py