ocr_platform/ocr_tools/cell_preprocess_lab/cell_sweep.py

#!/usr/bin/env python3
"""
单元格裁剪图预处理参数扫描：去水印 / threshold / contrast / upscale / det 阈值 / OCR 模式。

默认从 **原图**（`*_raw.png`）出发，与 pipeline 二次 OCR 一致，避免对已预处理 debug 图二次去水印。

用法:
    python cell_sweep.py cell219_empty_empty_raw.png -o ./out -t "ATM存折取款"
    python cell_sweep.py /path/to/tablecell_ocr/ -o ./out
    python cell_sweep.py cell.png --quick --no-save-images
    OCR_DET_MODEL_PATH=... OCR_REC_MODEL_PATH=... python cell_sweep.py cell.png
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from itertools import product
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple

import cv2
import numpy as np

_repo_root = Path(__file__).resolve().parents[2]
if str(_repo_root) not in sys.path:
    sys.path.insert(0, str(_repo_root))

from ocr_utils.watermark import WatermarkProcessor, merge_watermark_config
from ocr_utils.watermark.contrast import apply_contrast_enhancement_config

_IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff", ".webp"}
_DEFAULT_MODEL_DIR = Path(
    "/Users/zhch158/models/modelscope_cache/models/OpenDataLab/"
    "PDF-Extract-Kit-1___0/models/OCR/paddleocr_torch"
)


def _parse_csv_ints(s: str) -> List[Optional[int]]:
    out: List[Optional[int]] = []
    for part in s.split(","):
        part = part.strip()
        if not part or part.lower() in ("none", "d", "default"):
            out.append(None)
        else:
            out.append(int(part))
    return out


def _parse_csv_floats(s: str) -> List[float]:
    return [float(x.strip()) for x in s.split(",") if x.strip()]


def _parse_csv_bools(s: str) -> List[bool]:
    out: List[bool] = []
    for part in s.split(","):
        p = part.strip().lower()
        if p in ("1", "true", "yes", "on"):
            out.append(True)
        elif p in ("0", "false", "no", "off"):
            out.append(False)
        else:
            raise ValueError(f"无效的 bool 值: {part!r}")
    return out


def _default_model_dir() -> Path:
    det = os.environ.get("OCR_DET_MODEL_PATH")
    if det:
        return Path(det).parent
    return _DEFAULT_MODEL_DIR


def _upscale(img: np.ndarray, min_side: int) -> np.ndarray:
    h, w = img.shape[:2]
    if h >= min_side and w >= min_side:
        return img
    s = max(min_side / max(h, 1), min_side / max(w, 1), 1.0)
    return cv2.resize(img, None, fx=s, fy=s, interpolation=cv2.INTER_CUBIC)


def _preprocess(
    raw: np.ndarray,
    *,
    method: str,
    thresh: Optional[int],
    contrast: bool,
    upscale: int,
    text_black_target: int,
) -> np.ndarray:
    user: Dict[str, Any] = {"enabled": True, "method": method}
    if method == "threshold" and thresh is not None:
        user["threshold"] = thresh
    cfg = merge_watermark_config("cell", user)
    img, _ = WatermarkProcessor(cfg, scope="cell").process(raw, force=True)
    if contrast:
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        ce = dict(cfg.get("contrast_enhancement") or {})
        ce["enabled"] = True
        ce["text_black_target"] = text_black_target
        gray = apply_contrast_enhancement_config(gray, ce)
        img = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
    return _upscale(img, upscale)


def _parse_rec_pair(rec_part: Any) -> Tuple[str, float]:
    """从 OCR 返回的 (text, score) 或嵌套结构中解析识别结果。"""
    if rec_part is None:
        return "", 0.0
    if isinstance(rec_part, (list, tuple)) and len(rec_part) >= 2:
        if isinstance(rec_part[0], (list, tuple, dict)):
            return "", 0.0
        txt = str(rec_part[0] or "").strip()
        try:
            sc = float(rec_part[1] or 0.0)
        except (TypeError, ValueError):
            sc = 0.0
        return txt, sc if txt else 0.0
    if isinstance(rec_part, (list, tuple)) and len(rec_part) == 1:
        txt = str(rec_part[0] or "").strip()
        return txt, 0.0
    return "", 0.0


def _aggregate_rec_score(boxes: List[Dict[str, Any]]) -> float:
    """按字符数加权平均识别分（与 pipeline aggregate_line_ocr 一致）。"""
    total_len = sum(len(b.get("text") or "") for b in boxes)
    if total_len <= 0:
        return 0.0
    weighted = sum(
        len(b.get("text") or "") * float(b.get("score") or 0.0) for b in boxes
    )
    return weighted / total_len


def _ocr(engine: Any, img: np.ndarray, *, det: bool, rec: bool) -> Dict[str, Any]:
    empty: Dict[str, Any] = {
        "text": "",
        "score": 0.0,
        "boxes": [],
        "det": det,
        "rec": rec,
        "n_boxes": 0,
    }
    try:
        res = engine.ocr(img, det=det, rec=rec)
        items = res[0] if res and res[0] is not None else []
        boxes_out: List[Dict[str, Any]] = []

        if det:
            for item in items:
                if not item or len(item) < 2:
                    continue
                text, score = _parse_rec_pair(item[1])
                bbox = item[0]
                if hasattr(bbox, "tolist"):
                    bbox = bbox.tolist()
                entry: Dict[str, Any] = {
                    "text": text,
                    "score": round(score, 6),
                }
                if bbox is not None:
                    entry["det_bbox"] = bbox
                boxes_out.append(entry)
        else:
            for item in items:
                text, score = _parse_rec_pair(item)
                if not text and isinstance(item, (list, tuple)) and len(item) >= 1:
                    text, score = _parse_rec_pair(item[0])
                boxes_out.append({"text": text, "score": round(score, 6)})

        text = "".join(b["text"] for b in boxes_out if b.get("text")).strip()
        agg_score = _aggregate_rec_score(boxes_out)
        return {
            "text": text,
            "score": round(agg_score, 6),
            "boxes": boxes_out,
            "det": det,
            "rec": rec,
            "n_boxes": len(boxes_out),
        }
    except Exception as e:
        out = dict(empty)
        out["error"] = str(e)
        return out


def _make_engine(det_thresh: float, model_dir: Path) -> Any:
    from ocr_tools.pytorch_models.pytorch_paddle import PytorchPaddleOCR

    det_path = os.environ.get("OCR_DET_MODEL_PATH") or str(
        model_dir / "ch_PP-OCRv5_det_infer.pth"
    )
    rec_path = os.environ.get("OCR_REC_MODEL_PATH") or str(
        model_dir / "ch_PP-OCRv4_rec_server_doc_infer.pth"
    )
    return PytorchPaddleOCR(
        lang="ch",
        det_model_path=det_path,
        rec_model_path=rec_path,
        det_db_box_thresh=det_thresh,
    )


def resolve_input_image(path: Path, *, prefer_raw: bool) -> Path:
    """优先使用与 pipeline debug 配套的 *_raw.png。"""
    if not prefer_raw or path.stem.endswith("_raw"):
        return path
    raw_path = path.parent / f"{path.stem}_raw{path.suffix}"
    if raw_path.is_file():
        print(f"  使用原图: {raw_path.name}（跳过 {path.name}）")
        return raw_path
    return path


def collect_inputs(path: Path, *, prefer_raw: bool) -> List[Path]:
    if path.is_file():
        if path.suffix.lower() not in _IMAGE_SUFFIXES:
            raise ValueError(f"不支持的图像格式: {path}")
        return [resolve_input_image(path, prefer_raw=prefer_raw)]

    if not path.is_dir():
        raise FileNotFoundError(path)

    all_images = sorted(
        p
        for p in path.iterdir()
        if p.is_file() and p.suffix.lower() in _IMAGE_SUFFIXES
    )
    if not all_images:
        raise FileNotFoundError(f"目录内无图像: {path}")

    if prefer_raw:
        raws = [p for p in all_images if p.stem.endswith("_raw")]
        if raws:
            return raws

    chosen: List[Path] = []
    for p in all_images:
        if p.stem.endswith("_raw"):
            continue
        raw_sibling = p.parent / f"{p.stem}_raw{p.suffix}"
        if prefer_raw and raw_sibling.is_file():
            continue
        chosen.append(p)
    return chosen or all_images


def _match_hit(text: str, target: Optional[str]) -> Optional[str]:
    if not text:
        return None
    if not target:
        return "nonempty"
    if target in text:
        return "full"
    if len(target) >= 6 and target.isdigit() and len(text) >= 6 and text.isdigit():
        return "partial"
    return None


def run_sweep(
    input_path: Path,
    out_dir: Path,
    *,
    prefer_raw: bool,
    target: Optional[str],
    model_dir: Path,
    methods: Sequence[str],
    thresholds: Sequence[Optional[int]],
    contrasts: Sequence[bool],
    upscales: Sequence[int],
    det_threshs: Sequence[float],
    text_black_target: int,
    save_images: bool,
    run_baseline: bool,
    baseline_upscale: int,
) -> Dict[str, Any]:
    resolved = resolve_input_image(input_path, prefer_raw=prefer_raw)
    raw = cv2.imread(str(resolved))
    if raw is None:
        raise RuntimeError(f"无法读取图像: {resolved}")

    stem = resolved.stem.removesuffix("_raw") if resolved.stem.endswith("_raw") else resolved.stem
    cell_out = out_dir / stem
    cell_out.mkdir(parents=True, exist_ok=True)

    ocr_modes: List[Tuple[str, bool, bool]] = [
        ("det_rec", True, True),
        ("whole_rec", False, True),
    ]

    results: List[Dict[str, Any]] = []
    hits: List[Dict[str, Any]] = []
    engines: Dict[float, Any] = {}
    total = 0

    for method, thresh, contrast, upscale, det_th in product(
        methods, thresholds, contrasts, upscales, det_threshs
    ):
        if method != "threshold" and thresh is not None:
            continue
        if det_th not in engines:
            print(f"  [{stem}] 加载 OCR det_db_box_thresh={det_th} ...")
            engines[det_th] = _make_engine(det_th, model_dir)

        img = _preprocess(
            raw,
            method=method,
            thresh=thresh,
            contrast=contrast,
            upscale=upscale,
            text_black_target=text_black_target,
        )
        tag = f"{method}_t{thresh or 'd'}_c{int(contrast)}_u{upscale}_det{det_th}"
        if save_images:
            cv2.imwrite(str(cell_out / f"{tag}.png"), img)

        for mode_name, det, rec in ocr_modes:
            total += 1
            ocr = _ocr(engines[det_th], img, det=det, rec=rec)
            row: Dict[str, Any] = {
                "tag": tag,
                "method": method,
                "threshold": thresh,
                "contrast": contrast,
                "upscale": upscale,
                "det_db_box_thresh": det_th,
                "ocr_mode": mode_name,
                **ocr,
            }
            results.append(row)
            m = _match_hit(row.get("text", ""), target)
            if m:
                row["match"] = m
                hits.append(row)
                print(
                    f"  HIT [{m}] {mode_name} {tag} "
                    f"score={row.get('score')} -> {row.get('text')!r}"
                )

    if run_baseline:
        for det_th in det_threshs:
            if det_th not in engines:
                engines[det_th] = _make_engine(det_th, model_dir)
            base_img = _upscale(raw, baseline_upscale)
            if save_images:
                cv2.imwrite(str(cell_out / f"baseline_upscale{baseline_upscale}.png"), base_img)
            for mode_name, det, rec in ocr_modes:
                ocr = _ocr(engines[det_th], base_img, det=det, rec=rec)
                row = {
                    "tag": f"baseline_upscale{baseline_upscale}",
                    "det_db_box_thresh": det_th,
                    "ocr_mode": mode_name,
                    **ocr,
                }
                results.append(row)
                m = _match_hit(row.get("text", ""), target)
                if m:
                    row["match"] = m
                    hits.append(row)

    report = {
        "input": str(resolved),
        "input_requested": str(input_path),
        "output_dir": str(cell_out),
        "target": target,
        "total_trials": total,
        "hits": hits,
        "all_results": results,
    }
    report_path = cell_out / "sweep_report.json"
    report_path.write_text(
        json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    return report


def _build_arg_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        description="单元格图预处理 + OCR 参数网格扫描（对齐 pipeline 格级二次 OCR）",
    )
    p.add_argument(
        "input",
        type=Path,
        help="单元格裁剪图路径，或 tablecell_ocr 目录（批量扫描）",
    )
    p.add_argument(
        "-o",
        "--output",
        type=Path,
        default=None,
        help="输出目录，默认 <input_dir|input_parent>/sweep_out/<stem>",
    )
    p.add_argument(
        "-t",
        "--target",
        default=None,
        help="期望 OCR 文本；用于标记 HIT（子串匹配）。省略则任意非空为 HIT",
    )
    p.add_argument(
        "--model-dir",
        type=Path,
        default=None,
        help="PaddleOCR torch 模型目录（含 det/rec .pth），也可用 OCR_*_MODEL_PATH",
    )
    p.add_argument(
        "--no-prefer-raw",
        action="store_true",
        help="不自动选用同名的 *_raw.png",
    )
    p.add_argument(
        "--quick",
        action="store_true",
        help="缩小网格（threshold 170,175 × upscale 128,192 × det 0.3,0.5）",
    )
    p.add_argument(
        "--methods",
        default="threshold,masked_adaptive",
        help="去水印方式，逗号分隔",
    )
    p.add_argument(
        "--thresholds",
        default="155,165,170,175,180,none",
        help="threshold 法的阈值；none=预设默认",
    )
    p.add_argument(
        "--contrasts",
        default="false,true",
        help="是否 contrast，逗号分隔 false,true",
    )
    p.add_argument(
        "--upscales",
        default="64,96,128,192",
        help="最短边放大目标，逗号分隔整数",
    )
    p.add_argument(
        "--det-threshs",
        default="0.2,0.3,0.4,0.5",
        help="det_db_box_thresh，逗号分隔",
    )
    p.add_argument(
        "--text-black-target",
        type=int,
        default=88,
        help="contrast text_restore 目标黑度",
    )
    p.add_argument(
        "--no-save-images",
        action="store_true",
        help="不写出中间预处理 png（仅报告）",
    )
    p.add_argument(
        "--no-baseline",
        action="store_true",
        help="跳过「仅放大、不去水印」对照组",
    )
    p.add_argument(
        "--baseline-upscale",
        type=int,
        default=128,
        help="baseline 对照组的最短边放大",
    )
    return p


def main(argv: Optional[Sequence[str]] = None) -> None:
    args = _build_arg_parser().parse_args(argv)
    inputs = collect_inputs(args.input, prefer_raw=not args.no_prefer_raw)
    if not inputs:
        raise SystemExit("未找到可扫描的图像")

    if args.output is not None:
        out_root = args.output
    elif args.input.is_file():
        out_root = args.input.parent / "sweep_out"
    else:
        out_root = args.input / "sweep_out"
    out_root.mkdir(parents=True, exist_ok=True)

    model_dir = args.model_dir or _default_model_dir()
    methods = [m.strip() for m in args.methods.split(",") if m.strip()]

    if args.quick:
        thresholds = [170, 175]
        upscales = [128, 192]
        det_threshs = [0.3, 0.5]
        contrasts = [False, True]
    else:
        thresholds = _parse_csv_ints(args.thresholds)
        upscales = [int(x) for x in args.upscales.split(",") if x.strip()]
        det_threshs = _parse_csv_floats(args.det_threshs)
        contrasts = _parse_csv_bools(args.contrasts)

    print(f"扫描 {len(inputs)} 张图 -> {out_root}")
    print(f"  methods={methods} thresholds={thresholds} upscales={upscales}")
    if args.target:
        print(f"  target={args.target!r}")

    summary: List[Dict[str, Any]] = []
    for img_path in inputs:
        print(f"\n=== {img_path.name} ===")
        report = run_sweep(
            img_path,
            out_root,
            prefer_raw=not args.no_prefer_raw,
            target=args.target,
            model_dir=model_dir,
            methods=methods,
            thresholds=thresholds,
            contrasts=contrasts,
            upscales=upscales,
            det_threshs=det_threshs,
            text_black_target=args.text_black_target,
            save_images=not args.no_save_images,
            run_baseline=not args.no_baseline,
            baseline_upscale=args.baseline_upscale,
        )
        summary.append(
            {
                "input": report["input"],
                "hits": len(report["hits"]),
                "report": str(Path(report["output_dir"]) / "sweep_report.json"),
            }
        )

    index_path = out_root / "sweep_index.json"
    index_path.write_text(
        json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    print(f"\n全部完成，索引: {index_path}")
    for s in summary:
        print(f"  {s['input']}: {s['hits']} hits -> {s['report']}")


if __name__ == "__main__":
    if len(sys.argv) == 1:
        print("ℹ️  未提供命令行参数，使用默认配置运行...")
        default_config = {
            "input": "/Users/zhch158/workspace/data/流水分析/彭_广东兴宁农村商业银行/bank_statement_yusys_local/debug/table_recognition_wired/tablecell_ocr/彭_广东兴宁农村商业银行_page_002_0/cell219_empty_empty_raw.png",
            "output": "./output/彭_广东兴宁农村商业银行/cell219_sweep",
            "target": "ATM存折取款",
        }
        sys.argv = [sys.argv[0], default_config["input"]]
        for key, value in default_config.items():
            if key == "input":
                continue
            flag = f"--{key.replace('_', '-')}"
            if isinstance(value, bool) and value:
                sys.argv.append(flag)
            elif not isinstance(value, bool):
                sys.argv.extend([flag, str(value)])

    sys.exit(main())