project/tbsg/steps/ocr.py

"""
OCR 识别类 — 纯识别，不操作手机
基于 RapidOCR (rapidocr_onnxruntime)，无需 GPU
"""
import re
from pathlib import Path
from typing import Optional, Union

import cv2
import numpy as np
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)

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

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

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

        Returns:
            detail='all'时:
                [{'text': str, 'confidence': float,
                  'bbox': [[x1,y1],[x2,y1],[x2,y2],[x1,y2]]}, ...]
            detail='text'时:
                ['文字1', '文字2', ...]
            detail='first'时:
                str 或 None
        """
        # 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
        raw, _ = self._engine(img)

        # 5. 结构化 & 过滤，把坐标映射回原始图
        results = []
        for bbox, text, conf in (raw or []):
            if conf < confidence:
                continue
            if pattern and not re.search(pattern, text):
                continue
            results.append({
                "text": text,
                "confidence": round(float(conf), 4),
                "bbox": [[int(p[0] / scale) + crop_offset_l, int(p[1] / scale) + crop_offset_t] for p in bbox],
            })

        # 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)