onnx_ocr_matcher.py

import os
import threading
import time
from collections.abc import Callable
from logging import DEBUG
from typing import Any

from cv2.typing import MatLike

from one_dragon.base.matcher.match_result import MatchResult, MatchResultList
from one_dragon.base.matcher.ocr import ocr_utils
from one_dragon.base.matcher.ocr.ocr_match_result import OcrMatchResult
from one_dragon.base.matcher.ocr.ocr_matcher import OcrMatcher
from one_dragon.base.web.common_downloader import CommonDownloaderParam
from one_dragon.base.web.zip_downloader import ZipDownloader
from one_dragon.utils import os_utils, str_utils
from one_dragon.utils.i18_utils import gt
from one_dragon.utils.log_utils import log

DEFAULT_OCR_MODEL_NAME: str = 'ppocrv5'
PPOCRV6_MODEL_NAME: str = 'ppocrv6'
GITHUB_DOWNLOAD_URL: str = 'https://github.com/OneDragon-Anything/OneDragon-Env/releases/download'
CNB_DOWNLOAD_URL: str = 'https://cnb.cool/OneDragon-Anything/OneDragon-Env/-/releases/download'


def get_ocr_model_dir(ocr_model_name: str) -> str:
    return os_utils.get_path_under_work_dir('assets', 'models', 'onnx_ocr', ocr_model_name)


def get_ocr_download_url_github(ocr_model_name: str) -> str:
    return get_ocr_download_url(GITHUB_DOWNLOAD_URL, ocr_model_name)


def get_ocr_download_url_cnb(ocr_model_name: str) -> str:
    return get_ocr_download_url(CNB_DOWNLOAD_URL, ocr_model_name)


def get_ocr_download_url(website: str, ocr_model_name: str) -> str:
    return f'{website}/{ocr_model_name}/{ocr_model_name}.zip'


def get_ocr_model_dict_name(ocr_model_name: str) -> str | None:
    """
    获取模型对应的字典文件名。直接扫描本地文件夹中的字典文件。
    """
    base_dir = get_ocr_model_dir(ocr_model_name)
    if os.path.exists(base_dir):
        for f in os.listdir(base_dir):
            if f.endswith('_dict.txt'):
                return f
    return None


def get_final_file_list(ocr_model_name: str) -> list[str]:
    """
    下载成功后 整个模型的所有文件
    :param ocr_model_name: 模型名称
    :return:
    """
    base_dir = get_ocr_model_dir(ocr_model_name)
    files = [
        os.path.join(base_dir, 'det.onnx'),
        os.path.join(base_dir, 'rec.onnx'),
        os.path.join(base_dir, 'cls.onnx'),
        os.path.join(base_dir, 'simfang.ttf'),
    ]
    dict_name = get_ocr_model_dict_name(ocr_model_name)
    if dict_name is not None:
        files.append(os.path.join(base_dir, dict_name))
    return files


class OnnxOcrParam:
    """
    OCR配置实体类，包含OCR引擎的各项参数设置
    默认值见 onnxocr.utils
    """

    def __init__(
            self,
            ocr_model_name: str = DEFAULT_OCR_MODEL_NAME,
            det_model_name: str = 'det.onnx',
            rec_model_name: str = 'rec.onnx',
            cls_model_name: str = 'cls.onnx',
            dict_name: str | None = None,
            font_name: str = 'simfang.ttf',
            use_gpu: bool = False,
            use_angle_cls: bool = False,
            det_limit_side_len: float = 960.0,
            ocr_model_size: str | None = None,
    ):
        self.ocr_model_name: str = ocr_model_name
        self.models_dir: str = get_ocr_model_dir(ocr_model_name)
        if dict_name is None:
            dict_name = get_ocr_model_dict_name(ocr_model_name)
            if dict_name is None:
                # 首次运行未下载时，根据模型名推导一个默认的字典文件名，避免崩溃
                dict_name = f"{self.ocr_model_name}_dict.txt"
        # ===================================================================
        # I. 设备与性能 (Device & Performance)
        # ===================================================================
        self.use_gpu = use_gpu  # 是否使用GPU进行计算

        # ===================================================================
        # II. 模型路径 (Model Paths)
        # ===================================================================
        self.det_model_dir = os.path.join(self.models_dir, det_model_name)  # 文字检测模型文件路径
        self.rec_model_dir = os.path.join(self.models_dir, rec_model_name)  # 文字识别模型文件路径
        self.cls_model_dir = os.path.join(self.models_dir, cls_model_name)  # 方向分类模型文件路径
        self.rec_char_dict_path = os.path.join(self.models_dir, dict_name)  # 字符字典文件路径
        self.vis_font_path = os.path.join(self.models_dir, font_name)  # 可视化字体文件路径

        # ===================================================================
        # III. 核心功能开关 (Core Feature Switches)
        # ===================================================================
        self.use_angle_cls = use_angle_cls  # 是否加载并使用方向分类模型
        if self.ocr_model_name == PPOCRV6_MODEL_NAME or ocr_model_size is None:
            ocr_model_size = 'small'
        self.ocr_model_size: str | None = ocr_model_size

        # ===================================================================
        # IV. 文字检测超参数 (Detection Hyperparameters)
        # ===================================================================
        self.det_limit_side_len = det_limit_side_len  # 输入图像的长边限制

    def to_dict(self) -> dict[str, Any]:
        """将OCR配置转换为字典格式"""
        return {
            'use_gpu': self.use_gpu,
            'det_model_dir': self.det_model_dir,
            'rec_model_dir': self.rec_model_dir,
            'cls_model_dir': self.cls_model_dir,
            'rec_char_dict_path': self.rec_char_dict_path,
            'vis_font_path': self.vis_font_path,
            'use_angle_cls': self.use_angle_cls,
            'det_limit_side_len': self.det_limit_side_len,
            'ocr_model_size': self.ocr_model_size,
        }



class OnnxOcrMatcher(OcrMatcher, ZipDownloader):
    """
    使用onnx的ocr模型 速度更快
    """

    def __init__(self, ocr_param: OnnxOcrParam | None =  None):
        if ocr_param is None:
            ocr_param = OnnxOcrParam()
        OcrMatcher.__init__(self)
        param = CommonDownloaderParam(
            save_file_path=ocr_param.models_dir,
            save_file_name=f'{ocr_param.ocr_model_name}.zip',
            cnb_release_download_url=get_ocr_download_url_cnb(ocr_param.ocr_model_name),
            github_release_download_url=get_ocr_download_url_github(ocr_param.ocr_model_name),
            mirror_chan_download_url='',
            check_existed_list=get_final_file_list(ocr_param.ocr_model_name)
        )
        ZipDownloader.__init__(
            self,
            param=param,
        )
        self._ocr_param: OnnxOcrParam = ocr_param
        self._model = None
        self._init_lock = threading.Lock()
        self._loading: bool = False
        self.overlay_debug_bus = None

    @staticmethod
    def _rect_from_anchor(anchor_position) -> tuple[int, int, int, int] | None:
        """
        Convert OCR detector quadrilateral points to an axis-aligned rectangle.

        Some models return rotated polygons; using only point[0]/point[1]/point[3]
        can produce shifted overlay boxes. We normalize to min/max bounds.
        """
        if anchor_position is None:
            return None

        try:
            xs: list[int] = []
            ys: list[int] = []
            for point in anchor_position:
                if point is None or len(point) < 2:
                    continue
                xs.append(int(round(float(point[0]))))
                ys.append(int(round(float(point[1]))))

            if len(xs) == 0 or len(ys) == 0:
                return None

            x1 = min(xs)
            y1 = min(ys)
            x2 = max(xs)
            y2 = max(ys)
            return x1, y1, max(1, x2 - x1), max(1, y2 - y1)
        except Exception:
            return None

    def init_model(
            self,
            download_by_cnb: bool = False,
            download_by_github: bool = True,
            download_by_gitee: bool = False,
            download_by_mirror_chan: bool = False,
            proxy_url: str | None = None,
            ghproxy_url: str | None = None,
            skip_if_existed: bool = True,
            progress_callback: Callable[[float, str], None] | None = None
            ) -> bool:
        with self._init_lock:
            log.info('正在加载OCR模型')

            if self._model is not None:
                log.info('加载OCR模型完毕')
                return True

            # 先检查模型文件和下载模型
            done: bool = self.download(
                download_by_cnb=download_by_cnb,
                download_by_github=download_by_github,
                download_by_gitee=download_by_gitee,
                download_by_mirror_chan=download_by_mirror_chan,
                proxy_url=proxy_url,
                ghproxy_url=ghproxy_url,
                skip_if_existed=skip_if_existed,
                progress_callback=progress_callback
            )
            if not done:
                log.error('下载OCR模型失败')
                return False

            # 加载模型
            from onnxocr.onnx_paddleocr import ONNXPaddleOcr

            try:
                args = self._ocr_param.to_dict()
                self._model = ONNXPaddleOcr(**args)
                log.info('加载OCR模型完毕')
                return True
            except Exception:
                log.error('OCR模型加载出错', exc_info=True)
                return False

    def cleanup(self) -> None:
        """
        释放底层模型实例资源，协助 GC 回收 ONNX 会话
        """
        with self._init_lock:
            if self._model is not None:
                del self._model
                self._model = None

    def update_use_gpu(self, use_gpu: bool) -> None:
        """
        更新是否使用GPU

        Args:
            use_gpu: 是否使用GPU
        """
        if self._ocr_param.use_gpu == use_gpu:
            return

        self._ocr_param.use_gpu = use_gpu
        del self._model
        self._model = None

    def is_use_gpu(self) -> bool:
        return self._ocr_param.use_gpu

    def run_ocr_single_line(self, image: MatLike, threshold: float = 0, strict_one_line: bool = True) -> str:
        """
        单行文本识别 手动合成一行 按匹配结果从左到右 从上到下
        理论中文情况不会出现过长分行的 这里只是为了兼容英语的情况
        :param image: 图片
        :param threshold: 阈值
        :param strict_one_line: True时认为当前只有单行文本 False时依赖程序合并成一行
        :return:
        """
        if strict_one_line:
            return self._run_ocr_without_det(image, threshold)
        else:
            ocr_map: dict = self.run_ocr(image, threshold)
            tmp = ocr_utils.merge_ocr_result_to_single_line(ocr_map, join_space=False)
            return tmp

    def run_ocr(self, image: MatLike, threshold: float | None = 0,
                merge_line_distance: float = -1) -> dict[str, MatchResultList]:
        """
        对图片进行OCR 返回所有匹配结果
        :param image: 图片
        :param threshold: 匹配阈值
        :param merge_line_distance: 多少行距内合并结果 -1为不合并 理论中文情况不会出现过长分行的 这里只是为了兼容英语的情况
        :return: {key_word: []}
        """
        if image is None:
            log.warning('OCR输入的图片为None')
            return {}
        if self._model is None and not self.init_model():
            return {}
        start_time = time.time()
        result_map: dict = {}
        scan_result_list: list = self._model.ocr(
            image,
            det=True,
            rec=True,
            cls=self._ocr_param.use_angle_cls
        )
        if len(scan_result_list) == 0:
            if log.isEnabledFor(DEBUG):
                log.debug('OCR结果 %s 耗时 %.2f', result_map.keys(), time.time() - start_time)
            return result_map

        scan_result = scan_result_list[0]
        for anchor in scan_result:
            anchor_position = anchor[0]
            anchor_text = anchor[1][0]
            anchor_score = anchor[1][1]
            if anchor_score < threshold:
                continue
            rect = self._rect_from_anchor(anchor_position)
            if rect is None:
                continue
            if anchor_text not in result_map:
                result_map[anchor_text] = MatchResultList(only_best=False)
            result_map[anchor_text].append(
                MatchResult(
                    anchor_score,
                    rect[0],
                    rect[1],
                    rect[2],
                    rect[3],
                    data=anchor_text,
                )
            )

        if merge_line_distance != -1:
            result_map = ocr_utils.merge_ocr_result_to_multiple_line(result_map, join_space=True,
                                                                     merge_line_distance=merge_line_distance)

        elapsed_ms = (time.time() - start_time) * 1000.0
        self._emit_overlay_vision(result_map)
        self._emit_overlay_perf_and_timeline(elapsed_ms, len(result_map))

        if log.isEnabledFor(DEBUG):
            log.debug('OCR结果 %s 耗时 %.2f', result_map.keys(), time.time() - start_time)
        return result_map

    def _run_ocr_without_det(self, image: MatLike, threshold: float = 0) -> str:
        """
        不使用检测模型分析图片内文字的分布
        默认传入的图片仅有文字信息
        :param image: 图片
        :param threshold: 匹配阈值
        :return: [[("text", "score"),]] 由于禁用了空格，可以直接取第一个元素
        """
        if self._model is None and not self.init_model():
            return ""
        start_time = time.time()
        scan_result: list = self._model.ocr(
            image,
            det=False,
            rec=True,
            cls=self._ocr_param.use_angle_cls
        )
        img_result = scan_result[0]  # 取第一张图片
        if len(img_result) > 1:
            log.debug("禁检测的OCR模型返回多个识别结果")  # 目前没有出现这种情况

        if img_result[0][1] < threshold:
            log.debug("OCR模型返回的识别结果置信度低于阈值")
            return ""
        if log.isEnabledFor(DEBUG):
            log.debug('OCR结果 %s 耗时 %.2f', scan_result, time.time() - start_time)
        return img_result[0][0]

    def match_words(
            self,
            image: MatLike, words: list[str],
            threshold: float = 0,
            same_word: bool = False,
            ignore_case: bool = True,
            lcs_percent: float = -1,
            merge_line_distance: float = -1,
    ) -> dict[str, MatchResultList]:
        """
        在图片中查找关键词 返回所有词对应的位置
        :param image: 图片
        :param words: 关键词
        :param threshold: 匹配阈值
        :param same_word: 要求整个词一样
        :param ignore_case: 忽略大小写
        :param lcs_percent: 最长公共子序列长度百分比 -1代表不使用 same_word=True时不生效
        :param merge_line_distance: 多少行距内合并结果 -1为不合并
        :return: {key_word: []}
        """
        all_match_result: dict = self.run_ocr(image, threshold, merge_line_distance=merge_line_distance)
        match_key = set()
        for k in all_match_result:
            for w in words:
                ocr_result: str = k
                ocr_target = gt(w, 'ocr')
                if ignore_case:
                    ocr_result = ocr_result.lower()
                    ocr_target = ocr_target.lower()

                if same_word:
                    if ocr_result == ocr_target:
                        match_key.add(k)
                else:
                    if lcs_percent == -1:
                        if ocr_result.find(ocr_target) != -1:
                            match_key.add(k)
                    else:
                        if str_utils.find_by_lcs(ocr_target, ocr_result, percent=lcs_percent):
                            match_key.add(k)

        return {key: all_match_result[key] for key in match_key if key in all_match_result}

    def ocr(
            self,
            image: MatLike,
            threshold: float = 0,
            merge_line_distance: float = -1,
    ) -> list[OcrMatchResult]:
        """
        对图片进行OCR 返回所有识别结果

        Args:
            image: 图片
            threshold: 匹配阈值
            merge_line_distance: 多少行距内合并结果 -1为不合并 理论中文情况不会出现过长分行的 这里只是为了兼容英语的情况

        Returns:
            ocr_result_list: 识别结果列表
        """
        start_time = time.time()
        ocr_result_list: list[OcrMatchResult] = []

        scan_result_list: list = self._model.ocr(image, cls=False)
        if len(scan_result_list) == 0:
            if log.isEnabledFor(DEBUG):
                log.debug('OCR结果 [] 耗时 %.2f', time.time() - start_time)
            return ocr_result_list

        scan_result = scan_result_list[0]  # 只取第一张图片
        for anchor in scan_result:
            anchor_position = anchor[0]
            anchor_text = anchor[1][0]
            anchor_score = anchor[1][1]
            if anchor_score < threshold:
                continue
            rect = self._rect_from_anchor(anchor_position)
            if rect is None:
                continue

            result = OcrMatchResult(
                anchor_score,
                rect[0],
                rect[1],
                rect[2],
                rect[3],
                data=anchor_text)
            ocr_result_list.append(result)

        if merge_line_distance != -1:
            pass  # TODO

        elapsed_ms = (time.time() - start_time) * 1000.0
        self._emit_overlay_vision_from_ocr_results(ocr_result_list)
        self._emit_overlay_perf_and_timeline(elapsed_ms, len(ocr_result_list))

        if log.isEnabledFor(DEBUG):
            log.debug('OCR结果 %s 耗时 %.2f', [i.data for i in ocr_result_list], time.time() - start_time)

        return ocr_result_list

    def _emit_overlay_vision(
        self,
        result_map: dict[str, MatchResultList],
    ) -> None:
        bus = getattr(self, "overlay_debug_bus", None)
        if bus is None or not result_map:
            return

        try:
            from one_dragon.base.operation.overlay_debug_bus import VisionDrawItem
        except Exception:
            return

        ox, oy = bus.crop_offset
        pushed = 0
        max_items = 60
        for text, match_list in result_map.items():
            if match_list is None:
                continue
            for match in match_list.arr:
                if pushed >= max_items:
                    return
                label = str(text or "").strip()
                if len(label) > 32:
                    label = label[:29] + "..."
                bus.add_vision(
                    VisionDrawItem(
                        source="ocr",
                        label=label,
                        x1=match.x + ox,
                        y1=match.y + oy,
                        x2=match.x + match.w + ox,
                        y2=match.y + match.h + oy,
                        score=match.confidence,
                        color="#ff6ac1",
                        ttl_seconds=1.4,
                    )
                )
                pushed += 1

    def _emit_overlay_vision_from_ocr_results(
        self,
        ocr_results: list[OcrMatchResult],
    ) -> None:
        bus = getattr(self, "overlay_debug_bus", None)
        if bus is None or not ocr_results:
            return

        try:
            from one_dragon.base.operation.overlay_debug_bus import VisionDrawItem
        except Exception:
            return

        offset_x, offset_y = bus.crop_offset
        for result in ocr_results[:60]:
            label = str(result.data or "").strip()
            if len(label) > 32:
                label = label[:29] + "..."
            bus.add_vision(
                VisionDrawItem(
                    source="ocr",
                    label=label,
                    x1=result.x + offset_x,
                    y1=result.y + offset_y,
                    x2=result.x + result.w + offset_x,
                    y2=result.y + result.h + offset_y,
                    score=result.confidence,
                    color="#ff6ac1",
                    ttl_seconds=1.4,
                )
            )

    def _emit_overlay_perf_and_timeline(self, elapsed_ms: float, item_count: int) -> None:
        bus = getattr(self, "overlay_debug_bus", None)
        if bus is None:
            return
        try:
            from one_dragon.base.operation.overlay_debug_bus import (
                PerfMetricSample,
                TimelineItem,
            )
        except Exception:
            return
        bus.add_performance(
            PerfMetricSample(
                metric="ocr_ms",
                value=float(elapsed_ms),
                unit="ms",
                ttl_seconds=20.0,
                meta={"text_items": item_count},
            )
        )
        bus.add_timeline(
            TimelineItem(
                category="vision",
                title="ocr",
                detail=f"{item_count} items / {elapsed_ms:.1f}ms",
                level="DEBUG",
                ttl_seconds=15.0,
            )
        )


def __debug():
    ocr = OnnxOcrMatcher()
    ocr.init_model(
        download_by_github=False,
        download_by_gitee=True)

    from one_dragon.utils import debug_utils
    img = debug_utils.get_debug_image('1')
    print(ocr.run_ocr(img))


if __name__ == '__main__':
    __debug()