project/tbsg/steps/ocr.py

"""
OCR 识别类 — 纯识别，不操作手机
默认 RapidOCR (rapidocr_onnxruntime)，本地免费，用于导航大字检测
engine="cloud" 切换百度云 OCR（标准版，用于说明书等小字页面），失败自动回退本地
"""
import json
import re
from pathlib import Path
from typing import Optional, Union

import cv2
import numpy as np
from paddleocr import PaddleOCR
from rapidocr_onnxruntime import RapidOCR


class OCR:
    """RapidOCR 封装，只负责图片文字识别，不参与设备操作"""

    def __init__(self):
        import os as _os
        _os.environ.setdefault("OMP_NUM_THREADS", "1")
        self._engine = RapidOCR(text_det_thresh=0.3, text_det_box_thresh=0.3)
        self._baidu_client = None  # 百度OCR客户端，惰性初始化
        self.oom = False  # 内存崩溃熔断标志：本地引擎OOM后置True，上层据此跳过依赖OCR的步骤

    # ── 百度云 OCR ──────────────────────────────────────────

    def _load_baidu_keys(self) -> dict:
        """从 config.json 或环境变量读取百度 OCR 密钥"""
        import os as _os
        cfg_path = Path(__file__).parent.parent / "config.json"
        cfg = {}
        if cfg_path.exists():
            try:
                with open(cfg_path, encoding="utf-8") as f:
                    cfg = json.load(f)
            except Exception:
                pass
        ocr_cfg = cfg.get("baidu_ocr", {})
        return {
            "app_id": _os.environ.get("BAIDU_OCR_APP_ID") or ocr_cfg.get("app_id", ""),
            "api_key": _os.environ.get("BAIDU_OCR_API_KEY") or ocr_cfg.get("api_key", ""),
            "secret_key": _os.environ.get("BAIDU_OCR_SECRET_KEY") or ocr_cfg.get("secret_key", ""),
        }

    def _get_baidu_client(self):
        if self._baidu_client is None:
            from aip import AipOcr
            k = self._load_baidu_keys()
            if not (k["app_id"] and k["api_key"] and k["secret_key"]):
                raise RuntimeError("config.json 缺少 baidu_ocr 密钥")
            self._baidu_client = AipOcr(k["app_id"], k["api_key"], k["secret_key"])
        return self._baidu_client

    @staticmethod
    def _baidu_to_standard(words_result: list) -> list:
        """百度 words_result → 项目统一格式 [{text, confidence, bbox, box}]"""
        results = []
        for item in words_result or []:
            loc = item.get("location", {})
            l, t = loc.get("left", 0), loc.get("top", 0)
            w, h = loc.get("width", 0), loc.get("height", 0)
            prob = item.get("probability") or {}
            conf = prob.get("average") if isinstance(prob, dict) and prob.get("average") is not None else 1.0
            results.append({
                "text": item.get("words", ""),
                "confidence": round(float(conf), 4),
                # 四角点：左上、右上、右下、左下
                "bbox": [[l, t], [l + w, t], [l + w, t + h], [l, t + h]],
                # 简化坐标：[左上x, 左上y, 右下x, 右下y]，方便直接查看
                "box": [l, t, l + w, t + h],
            })
        return results

    def recognize_cloud(
        self,
        image: Union[str, Path, np.ndarray],
        rect: Optional[list] = None,
        confidence: float = 0.1,
        pattern: Optional[str] = None,
        detail: str = "all",
    ) -> list:
        """百度云 OCR（通用文字识别标准版），返回格式与 recognize 一致"""
        if isinstance(image, (str, Path)):
            img = cv2.imread(str(image))
            if img is None:
                raise FileNotFoundError(f"无法读取图片: {image}")
        else:
            img = image
        if rect:
            l, t, r, b = rect
            img = img[t:b, l:r]

        client = self._get_baidu_client()
        # JPEG 压缩发送，减小体积（百度推荐）；标准版上限 4M/4096px
        ok, enc = cv2.imencode(".jpg", img, [cv2.IMWRITE_JPEG_QUALITY, 90])
        if not ok:
            raise RuntimeError("图片编码失败")
        # 标准版（高精度版 accurate 配额小、天天触发日限 error_code=17，已弃用；
        # 生僻字场景已由 PP-OCRv6 + 交叉验证覆盖）
        res = client.basicGeneral(enc.tobytes())
        words_result = res.get("words_result", []) if isinstance(res, dict) else []
        if isinstance(res, dict) and res.get("error_code"):
            print(f"[ocr] 标准版受限(error_code={res.get('error_code')}: {res.get('error_msg')})")
            words_result = []

        results = self._baidu_to_standard(words_result)
        if confidence:
            results = [r for r in results if r["confidence"] >= confidence]
        if pattern:
            results = [r for r in results if re.search(pattern, r["text"])]

        if detail == "text":
            return [r["text"] for r in results]
        elif detail == "first":
            return results[0]["text"] if results else None
        return results

    def recognize_license(self, image: Union[str, Path, np.ndarray], rect: Optional[list] = None) -> dict:
        """
        百度营业执照专用接口（美团同款 /rest/2.0/ocr/v1/business_license）。
        返回结构化字段 dict，如 {"注册号": "...", "公司名称": "...", "地址": "...", ...}
        失败返回 {}

        Args:
            image: 图片路径或 numpy 数组
            rect:  裁剪区域 [left, top, right, bottom]，不传则整图
        """
        if isinstance(image, (str, Path)):
            img = cv2.imread(str(image))
            if img is None:
                raise FileNotFoundError(f"无法读取图片: {image}")
        else:
            img = image
        if rect:
            l, t, r, b = rect
            img = img[t:b, l:r]
        ok, enc = cv2.imencode(".jpg", img, [cv2.IMWRITE_JPEG_QUALITY, 90])
        if not ok:
            return {}
        client = self._get_baidu_client()
        try:
            res = client.businessLicense(enc.tobytes())
        except Exception as e:
            print(f"[ocr] 百度营业执照OCR失败: {e}")
            return {}
        words_result = res.get("words_result", {}) if isinstance(res, dict) else {}
        return {k: v.get("words", "") for k, v in (words_result or {}).items()}

    # ── 核心识别 ────────────────────────────────────────────

    def recognize(
        self,
        image: Union[str, Path, np.ndarray],
        rect: Optional[list] = None,
        confidence: float = 0.1,
        pattern: Optional[str] = None,
        detail: str = "all",
        engine: str = "local",
    ) -> list:
        """
        对图片执行 OCR 识别

        Args:
            image: 图片路径(str/Path) 或 numpy 数组
            rect:  裁剪区域 [left, top, right, bottom]，不传则全图
            confidence: 置信度阈值，默认 0.1
            pattern:    正则过滤，只返回匹配的文字
            detail:     'all'=返回完整结果, 'text'=只返回文字列表, 'first'=只返回第一个匹配文字
            engine:     'local'=RapidOCR(默认), 'cloud'=百度云OCR(失败自动回退本地)

        Returns:
            detail='all'时:
                [{'text': str, 'confidence': float,
                  'bbox': [[x1,y1],[x2,y1],[x2,y2],[x1,y2]]}, ...]
            detail='text'时:
                ['文字1', '文字2', ...]
            detail='first'时:
                str 或 None
        """
        if engine == "cloud":
            try:
                return self.recognize_cloud(image, rect, confidence, pattern, detail)
            except Exception as e:
                print(f"[ocr] 百度OCR失败，回退本地: {e}")
        # 1. 读取图片
        if isinstance(image, (str, Path)):
            img = cv2.imread(str(image))
            if img is None:
                raise FileNotFoundError(f"无法读取图片: {image}")
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        elif isinstance(image, np.ndarray):
            img = image
        else:
            raise TypeError(f"不支持的图片类型: {type(image)}")

        # 2. 裁剪区域（在压缩前，用原始坐标）
        crop_offset_l, crop_offset_t = 0, 0
        if rect:
            l, t, r, b = rect
            img = img[t:b, l:r]
            crop_offset_l, crop_offset_t = l, t

        # 3. 压缩大图，避免ONNX OOM
        orig_h, orig_w = img.shape[:2]
        if orig_h < 1 or orig_w < 1:
            return [] if detail != "first" else ""
        scale = 1.0
        max_size = 960
        if max(orig_h, orig_w) > max_size:
            scale = max_size / max(orig_h, orig_w)
            img = cv2.resize(img, (max(1, int(orig_w * scale)), max(1, int(orig_h * scale))))

        # 4. 执行 OCR（内存崩溃时置熔断标志后抛出，由上层跳过依赖OCR的步骤）
        try:
            raw, _ = self._engine(img)
        except Exception as e:
            if any(k in str(e).lower() for k in ("bad allocation", "out of memory", "insufficient memory", "unable to allocate")):
                self.oom = True
                print("[ocr] 本地OCR内存崩溃，已熔断后续OCR步骤")
            raise

        # 5. 结构化 & 过滤，把坐标映射回原始图
        results = []
        for bbox, text, conf in (raw or []):
            if conf < confidence:
                continue
            if pattern and not re.search(pattern, text):
                continue
            pts = [[int(p[0] / scale) + crop_offset_l, int(p[1] / scale) + crop_offset_t] for p in bbox]
            results.append({
                "text": text,
                "confidence": round(float(conf), 4),
                # 四角点：左上、右上、右下、左下
                "bbox": pts,
                # 简化坐标：[左上x, 左上y, 右下x, 右下y]，方便直接查看
                "box": [pts[0][0], pts[0][1], pts[2][0], pts[2][1]],
            })

        # 5. 按需返回格式
        if detail == "text":
            return [r["text"] for r in results]
        elif detail == "first":
            return results[0]["text"] if results else None
        return results  # detail='all'

    # ── 快捷方法 ────────────────────────────────────────────

    def recognize_file(self, path: str, **kwargs) -> list:
        """从文件路径识别"""
        return self.recognize(path, **kwargs)

    def recognize_array(self, arr: np.ndarray, **kwargs) -> list:
        """从 numpy 数组识别"""
        return self.recognize(arr, **kwargs)


if __name__ == "__main__":
    """
    验证 OCR 对指定图片的识别效果（本地RapidOCR / 百度云高精度 / PaddleOCR 三引擎对比）。
    用法: python steps/ocr.py <图片路径> [关键词]
    例:   python steps/ocr.py "C:/Users/18357/Downloads/image (5).png" 胃舒
    """
    import sys as _sys

    img_path = _sys.argv[1] if len(_sys.argv) > 1 else r"C:\Users\18357\Downloads\image.png"
    keyword = _sys.argv[2] if len(_sys.argv) > 2 else ""
    print(f"验证图片: {img_path}")
    print(f"过滤关键词: {keyword}\n")

    _ocr = OCR()

    # 1. 本地 RapidOCR
    print("=== 本地 RapidOCR ===")
    try:
        raw = _ocr.recognize(img_path, detail="all")
        hits = [r for r in raw if keyword in r["text"]]
        if not hits:
            print("  (未识别到含关键词的文字)")
        for r in hits:
            print(f"  conf={r['confidence']:.2f} box={r['box']} [{r['text']}]")
    except Exception as e:
        print(f"  识别失败: {e}")

    # 2. 百度云标准版
    print("=== 百度云标准版 ===")
    try:
        raw_c = _ocr.recognize(img_path, detail="all", engine="cloud")
        hits = [r for r in raw_c if keyword in r["text"]]
        if not hits:
            print("  (未识别到含关键词的文字)")
        for r in hits:
            print(f"  conf={r['confidence']:.2f} box={r['box']} [{r['text']}]")
    except Exception as e:
        print(f"  云端失败: {e}")

    # 3. PaddleOCR
    print("=== PaddleOCR ===")
    try:
        import cv2 as _cv2
        img = _cv2.imread(img_path)
        img = _cv2.cvtColor(img, _cv2.COLOR_BGR2RGB)
        _paddle = PaddleOCR(lang="ch", show_log=False)
        res = _paddle.ocr(img)
        hits = [line[1] for line in (res[0] or []) if keyword in line[1][0]]
        if not hits:
            print("  (未识别到含关键词的文字)")
        for text, conf in hits:
            print(f"  conf={conf:.2f} [{text}]")
    except Exception as e:
        print(f"  PaddleOCR失败: {e}")