PokemonCardSearch/code/pokemon/flat_card_no_ocr.py

"""
flat_card_no_ocr.py - 步骤②：编号区域 OCR + 单次宽容 + 低分不合规置空

链路：
  card_no 裁剪图 → PP-OCRv6_medium_rec → apply_leniency(单次)
  → rec_score≤0.9 且不合规格式 → card_no 置空（blur_unreadable）

空结果区分：
  no_detection     — 上层无检测框（卡本身无编号），本脚本不产生
  blur_unreadable  — 有框但模糊不可读
  rec_empty        — 有框但 rec 空文本

用法：
  export OUT_DIR=/path/to/output
  ~/miniconda3/envs/paddleocr37/bin/python flat_card_no_ocr.py
"""
import os, sys, json, time, datetime
sys.stdout.reconfigure(encoding="utf-8")

_HERE = os.path.dirname(os.path.abspath(__file__))
for root in (
    "/home/martin/顾工交接/wzj",
    os.path.expanduser("~/顾工交接/wzj"),
    os.path.abspath(os.path.join(_HERE, "..", "..", "..")),
):
    if os.path.isdir(os.path.join(root, "modules")):
        sys.path.insert(0, root)
        break

from modules.card_no_leniency import finalize_card_no

OUT_DIR = os.environ.get("OUT_DIR", "/home/martin/顾工交接/wzj/_infer_out")
REC_MODEL = os.environ.get("REC_MODEL", "PP-OCRv6_medium_rec")
RECORDS = os.path.join(OUT_DIR, "card_no_crop_records.json")
SCORE_THR = float(os.environ.get("REC_SCORE_THR", "0.9"))

REC_KWARGS = {"model_name": REC_MODEL}
if REC_MODEL.startswith("PP-OCRv6"):
    REC_KWARGS.update(engine="transformers", device="cpu")

print(f"[OCR] 初始化 TextRecognition({REC_MODEL}) ...")
from paddleocr import TextRecognition
_t = time.perf_counter()
rec = TextRecognition(**REC_KWARGS)
t_load = time.perf_counter() - _t
print(f"[OCR] 加载完成 ({t_load:.2f}s)\n")

records = json.load(open(RECORDS, encoding="utf-8"))
paths = [r["crop_path"] for r in records]
print(f"[RUN] {len(paths)} 张裁剪图 (rec-only + 单次宽容 + 低分置空 thr={SCORE_THR})\n")

_t = time.perf_counter()
outs = []
for r in rec.predict(paths):
    outs.append((r.get("rec_text", "") or "", float(r.get("rec_score", 0.0) or 0.0)))
t_ocr = time.perf_counter() - _t

ocr_results = {}
n_empty_rec = n_blur = n_lenient = n_keep = 0
for meta, (text, score) in zip(records, outs):
    fin = finalize_card_no(text, score, score_threshold=SCORE_THR)
    if fin["leniency_applied"]:
        n_lenient += 1
    er = fin["empty_reason"]
    if er == "rec_empty":
        n_empty_rec += 1
    elif er == "blur_unreadable":
        n_blur += 1
    else:
        n_keep += 1

    ocr_results[meta["file"]] = {
        "rec_text": (text or "").strip(),
        "card_no": fin["card_no"],
        "leniency_applied": fin["leniency_applied"],
        "empty_reason": er,  # None / rec_empty / blur_unreadable
        "rec_score": score,
        "det_conf": meta["conf"],
        "crop_wh": meta.get("crop_wh"),
        "status": "detected",  # 本文件只处理有检测框的样本
    }

n = max(len(paths), 1)
avg = sum(s for _, s in outs) / n if n else 0.0
print(f"\n[RESULT] {n} 张(有检测框)")
print(f"  保留编号 {n_keep} | 宽容命中 {n_lenient}")
print(f"  rec空 {n_empty_rec} | 模糊置空(blur_unreadable) {n_blur}")
print(f"  平均 rec_score {avg:.3f}")
print("=" * 50)
print(f"⏱ flat_card_no_ocr | {n} 张 | 加载 {t_load:.2f}s | rec {t_ocr:.2f}s ({n/t_ocr:.1f}/s)")
print("=" * 50)

with open(os.path.join(OUT_DIR, "timing.log"), "a", encoding="utf-8") as lf:
    lf.write(f"\n[{datetime.datetime.now():%Y-%m-%d %H:%M:%S}] === flat_card_no_ocr.py ===\n")
    lf.write(f"  有检测框 {n}, 保留 {n_keep}, 宽容 {n_lenient}, rec空 {n_empty_rec}, 模糊置空 {n_blur}\n")
    lf.write(f"  加载 {t_load:.2f}s, rec {t_ocr:.2f}s ({n/t_ocr:.1f}/s)\n")

json.dump(ocr_results, open(os.path.join(OUT_DIR, "card_no_ocr.json"), "w", encoding="utf-8"),
          ensure_ascii=False, indent=2)
print(f"✓ {OUT_DIR}/card_no_ocr.json")