Class OpenOCRE2E

tools/infer_e2e.py:78–402

class OpenOCRE2E(object):

    def __init__(self,
                 mode='mobile',
                 backend='torch',
                 onnx_det_model_path=None,
                 onnx_rec_model_path=None,
                 drop_score=0.5,
                 det_box_type='quad',
                 use_gpu='auto'):
        """
        初始化函数，用于初始化OCR引擎的相关配置和组件。

        Args:
            mode (str, optional): 运行模式，可选值为'mobile'或'server'。默认为'mobile'。
            drop_score (float, optional): 检测框的置信度阈值，低于该阈值的检测框将被丢弃。默认为0.5。
            det_box_type (str, optional): 检测框的类型，可选值为'quad' and 'poly'。默认为'quad'。
            use_gpu (str, optional): GPU使用策略，可选值为'auto'/'true'/'false'。默认为'auto'。

        Returns:
            无返回值。

        """
        # Auto-switch backend for server mode
        if mode == 'server' and backend != 'torch':
            logger.warning(
                f"Server mode only supports 'torch' backend, got '{backend}'. "
                f"Automatically switching to 'torch' backend. "
                f"Please make sure 'torch' and 'torchvision' are installed: "
                f"pip install torch torchvision")
            backend = 'torch'

        # Parse use_gpu parameter
        if use_gpu == 'auto':
            try:
                import torch
                device = 'gpu' if torch.cuda.is_available() else 'cpu'
            except:
                device = 'cpu'
        elif use_gpu == 'true':
            device = 'gpu'
        elif use_gpu == 'false':
            device = 'cpu'
        else:
            raise ValueError(f"use_gpu must be 'auto', 'true', or 'false', got '{use_gpu}'")

        cfg_det = Config(DEFAULT_CFG_PATH_DET).cfg  # mobile model
        cfg_det['Global']['device'] = device
        if mode == 'server':
            cfg_rec = Config(DEFAULT_CFG_PATH_REC_SERVER).cfg  # server model
        else:
            cfg_rec = Config(DEFAULT_CFG_PATH_REC).cfg  # mobile model

        cfg_rec['Global']['device'] = device

        self.text_detector = OpenDetector(cfg_det,
                                          backend=backend,
                                          onnx_model_path=onnx_det_model_path,