# -*- coding: utf-8 -*-
"""
grading_ocr_step2.py - 对裁剪出的评级标签条跑 PaddleOCR，输出分行文本

运行:
  conda activate paddleocr
  python grading_ocr_step2.py

输出:
  - ocr_result.json: 每张图的 OCR 分行文本列表
  - ocr_summary.csv: 汇总表(文件名,真实类别,评级框,置信度,OCR分行文本)
"""
import os
import sys
import json
import csv

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

OUT = "/home/martin/顾工交接/wzj/grading_ocr"
CROP_DIR = os.path.join(OUT, "crops")
RECORDS = os.path.join(OUT, "crop_records.json")

print("[OCR] 初始化 PaddleOCR ...")
from paddleocr import PaddleOCR
ocr = PaddleOCR(
    use_doc_orientation_classify=True,   # 开启方向分类,旋转卡文本识别更准
    use_doc_unwarping=False,
    use_textline_orientation=True,
)
print("[OCR] 加载完成\n")

# 读裁剪记录
with open(RECORDS, "r", encoding="utf-8") as f:
    records = json.load(f)

ocr_results = {}
rows_csv = []

print(f"[OCR] 处理 {len(records)} 张裁剪图 ...\n")
for i, rec in enumerate(records):
    crop_path = rec["crop_path"]
    name = rec["file"]
    true_label = rec["true_label"]
    box = rec["box"]
    conf = rec["conf"]

    try:
        result = ocr.predict(crop_path)
    except Exception as e:
        print(f"[{i+1}/{len(records)}] {name} OCR失败: {e}")
        ocr_results[name] = {"lines": [], "error": str(e)}
        rows_csv.append([name, true_label, str(box) if box else "", conf or "", ""])
        continue

    # 提取分行文本
    lines = []
    for res in result:
        if isinstance(res, dict):
            texts = res.get("rec_texts")
            if texts:
                lines.extend(list(texts))
        elif hasattr(res, "rec_texts") and res.rec_texts:
            lines.extend(list(res.rec_texts))

    # 合并分行文本(用 | 分隔,便于CSV查看)
    text_joined = " | ".join(lines) if lines else ""
    ocr_results[name] = {"lines": lines, "text": text_joined}

    rows_csv.append([name, true_label, str(box) if box else "", conf or "", text_joined])

    # 打印前20条示例
    if i < 20:
        print(f"[{i+1}] {true_label}/{name}")
        for ln in lines:
            print(f"    {ln}")
        print()

# 保存 JSON
with open(os.path.join(OUT, "ocr_result.json"), "w", encoding="utf-8") as f:
    json.dump(ocr_results, f, ensure_ascii=False, indent=2)
print(f"\n[OCR] JSON 已保存: {OUT}/ocr_result.json")

# 保存 CSV
csv_path = os.path.join(OUT, "ocr_summary.csv")
with open(csv_path, "w", encoding="utf-8-sig", newline="") as f:
    w = csv.writer(f)
    w.writerow(["filename", "true_label", "box_xyxy", "conf", "ocr_text_lines"])
    w.writerows(rows_csv)
print(f"[OCR] CSV 已保存: {csv_path}")

# 统计
n_empty = sum(1 for v in ocr_results.values() if not v.get("lines"))
print(f"\n统计: 共 {len(records)} 张, OCR空文本 {n_empty} 张")