# -*- coding: utf-8 -*-
"""
flat_ocr.py - 平铺版 PaddleOCR (对 crops/ 里的评级标签条识别分行文本)
环境: paddleocr
输出: OUT/ocr_result.json
"""
import os, sys, json, time, datetime
sys.stdout.reconfigure(encoding="utf-8")

OUT = os.environ.get("OUT_DIR", "/home/martin/顾工交接/wzj/_infer_out")
RECORDS = os.path.join(OUT, "crop_records.json")

print("[OCR] 初始化 PaddleOCR ...")
from paddleocr import PaddleOCR
_t = time.perf_counter()
ocr = PaddleOCR(use_doc_orientation_classify=True, use_doc_unwarping=False,
                use_textline_orientation=True)
t_load = time.perf_counter() - _t
print(f"[OCR] 加载完成 ({t_load:.2f}s)\n")

records = json.load(open(RECORDS, encoding="utf-8"))
ocr_results = {}
t_ocr = 0.0
_t_all = time.perf_counter()
for i, rec in enumerate(records):
    name = rec["file"]; cp = rec["crop_path"]
    try:
        _t = time.perf_counter()
        result = ocr.predict(cp)
        t_ocr += time.perf_counter() - _t
    except Exception as e:
        ocr_results[name] = {"lines": [], "error": str(e)}; continue
    lines = []
    for res in result:
        if isinstance(res, dict):
            t = res.get("rec_texts")
            if t: lines.extend(list(t))
        elif hasattr(res, "rec_texts") and res.rec_texts:
            lines.extend(list(res.rec_texts))
    ocr_results[name] = {"lines": lines, "text": " | ".join(lines)}
    if (i + 1) % 200 == 0:
        print(f"  [{i+1}/{len(records)}] 完成", flush=True)

json.dump(ocr_results, open(os.path.join(OUT, "ocr_result.json"), "w", encoding="utf-8"),
          ensure_ascii=False, indent=2)
n = max(len(records), 1); t_all = time.perf_counter() - _t_all
print(f"\n[OCR] 完成 {n} 张")
print("=" * 50)
print(f"⏱ flat_ocr (PaddleOCR) | {n} 张")
print(f"  加载 {t_load:.2f}s | OCR推理 {t_ocr:.2f}s (均{t_ocr/n*1000:.1f}ms) | 总 {t_all:.2f}s ({n/t_all:.2f}张/s)")
print("=" * 50)
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}] === 步骤② flat_ocr.py (PaddleOCR 评级标签) ===\n")
    lf.write(f"  样本 {n} 张\n")
    lf.write(f"  PaddleOCR 加载       : {t_load:.2f}s\n")
    lf.write(f"  OCR 推理累计         : {t_ocr:.2f}s (均 {t_ocr/n*1000:.1f} ms/张)\n")
    lf.write(f"  总耗时               : {t_all:.2f}s (均 {t_all/n*1000:.1f} ms/张, 吞吐 {n/t_all:.2f} 张/s)\n")
print(f"✓ 输出: {OUT}/ocr_result.json | 计时→ timing.log")