PokemonCardSearch/code/ocr_judge.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""第二步：对 YOLO 裁剪图做 OCR + 语种判定（PaddleOCR 3.2.0）

流程: 裁剪图 → PaddleOCR 识别 → 文本截断500字符 → 语种判定
运行环境: paddleocr conda 环境

输出 language_labels_320.json，每条含:
  - language: 语种判断结果 (tcg us / tcg jp / 简中 / 繁中 / None)
  - detail: 判定依据 (kana/cjk/latin/s/t 计数)
  - text: OCR 识别文本 (最多500字符)
  - angle: 图像方向角度 0/90/180/270（供主流程做卡牌正立翻转，不影响 OCR/语种）

注: use_doc_orientation_classify=True 是必要的——实测关闭时旋转卡的 OCR 会乱码，
开启后旋转卡文本识别正确(语种判定跟着对)，并顺带给出方向角度。
"""
import os, sys, csv, hashlib, json
sys.stdout.reconfigure(encoding="utf-8")
sys.path.insert(0, "/home/martin/顾工交接/wzj")
import config

# ============ 配置 ============
MAX_TEXT_LEN = 500  # OCR 文本最大字符数

DATA = config.DATA_DIR
def md5n(u): return hashlib.md5(u.encode()).hexdigest() + ".jpg"
CROP_DIR = os.path.join(DATA, "query_crops")

# ============ OCR 初始化 ============
print("[OCR] 初始化 PaddleOCR 3.2.0...", flush=True)
from paddleocr import PaddleOCR
ocr = PaddleOCR(
    use_doc_orientation_classify=True,
    use_doc_unwarping=False,
    use_textline_orientation=True,
)
print("[OCR] 加载完成", flush=True)


def ocr_full_text(img_path):
    """OCR 识别，返回 (完整文本, 方向角度)。
    角度来自 doc_preprocessor_res.angle（0/90/180/270，输入相对正立顺时针的旋转）。"""
    result = ocr.predict(img_path)
    texts, angle = [], 0
    for res in result:
        # res 是 dict-like，用 dict 访问
        if isinstance(res, dict):
            t = res.get("rec_texts")
            if t:
                texts.extend(list(t))
            dpr = res.get("doc_preprocessor_res")
            if dpr is not None:
                a = dpr.get("angle") if isinstance(dpr, dict) else getattr(dpr, "angle", None)
                if a is not None:
                    try:
                        angle = int(a)
                    except Exception:
                        angle = a
        elif hasattr(res, "rec_texts") and res.rec_texts:
            texts.extend(list(res.rec_texts))
            dpr = getattr(res, "doc_preprocessor_res", None)
            if dpr is not None:
                a = getattr(dpr, "angle", None)
                if a is not None:
                    try:
                        angle = int(a)
                    except Exception:
                        angle = a
    return "".join(texts), angle


# ============ 语种判定 ============
import hanzidentifier as hz


def judge(text):
    """语种判定逻辑:
       ① 假名>=10 → tcg jp (日文铁证)
       ② 中文占比>25% (cjk*3 > latin) → 是中文卡 → hanzidentifier 判简繁
       ③ 否则 (cjk*3 <= latin) → tcg us (英文)
    卡牌上英文元素(HP/Weakness/版权/编号)天然拉高latin，故中文阈值放宽到25%。
    """
    kana = sum(1 for ch in text if 0x3040 <= ord(ch) < 0x3100)
    cjk = sum(1 for ch in text if 0x4E00 <= ord(ch) < 0x9FFF)
    latin = sum(1 for ch in text if ch.isascii() and ch.isalpha())
    if kana >= 10:
        return "tcg jp", {"kana": kana, "cjk": cjk, "latin": latin}
    # 中文判定阈值: 中文字符*3 > 英文字符 (中文相对英文占比>25%) → 判为中文卡
    if cjk * 3 > latin:
        s = sum(1 for ch in text if hz.identify(ch) == 2)  # 简体字
        t = sum(1 for ch in text if hz.identify(ch) == 1)  # 繁体字
        script = "简中" if s >= t else "繁中"
        return script, {"kana": kana, "s": s, "t": t, "cjk": cjk}
    # cjk*3 <= latin: 英文占主导 → 英文卡
    return "tcg us", {"kana": kana, "latin": latin, "cjk": cjk}


# ============ 主循环 ============
txs = list(csv.DictReader(open(os.path.join(DATA, "transactions_ebay.csv"), encoding="utf-8-sig")))
print("[OCR] 处理 %d 张裁剪图（文本上限%d字符）" % (len(txs), MAX_TEXT_LEN), flush=True)
print("=" * 60, flush=True)

labels = {}
empty_count = 0
for i, tx in enumerate(txs):
    tid = tx["tx_id"]
    crop_path = os.path.join(CROP_DIR, md5n(tx["img_url"]))
    if not os.path.exists(crop_path):
        labels[tid] = {"language": None, "detail": {"err": "no_crop"}, "text": ""}
        continue
    try:
        full_text, angle = ocr_full_text(crop_path)
        # 截断到最大字符数
        text = full_text[:MAX_TEXT_LEN]
        if not text:
            empty_count += 1
        # 基于截断后的文本做语种判定
        lang, det = judge(text)
    except Exception as e:
        lang, det = None, {"err": str(e)[:60]}
        text, angle = "", 0
    labels[tid] = {"language": lang, "detail": det, "text": text, "angle": angle}
    print("[%3d] %-30s -> %-7s ang=%s %s" % (i + 1, tx["title"][:30], lang or "?", angle, det), flush=True)

# ============ 保存 ============
out = os.path.join(DATA, "language_labels_320.json")
with open(out, "w", encoding="utf-8") as f:
    json.dump(labels, f, ensure_ascii=False, indent=1)

from collections import Counter
cnt = Counter(v["language"] for v in labels.values())
print("=" * 60)
print("===== 语种判断模型结果 (PaddleOCR 3.2.0 + YOLO裁剪, 文本≤%d) =====" % MAX_TEXT_LEN)
for k, v in cnt.most_common():
    print("  %-8s %d" % (k, v))
print("OCR空文本: %d" % empty_count)
print("已保存: %s" % out)