Class OpenDetector

tools/infer_det.py:153–329

class OpenDetector(object):

    def __init__(self,
                 config=None,
                 backend='torch',
                 onnx_model_path=None,
                 use_gpu='auto',
                 numId=0):
        """
        Args:
            config (dict, optional): 配置信息。默认为None。
            backend (str): 'torch' 或 'onnx'
            onnx_model_path (str): ONNX模型路径（仅当backend='onnx'时需要）
            use_gpu (str, optional): GPU使用策略，可选值为'auto'/'true'/'false'。默认为'auto'。
            numId (int, optional): 设备编号。默认为0。
        """

        if config is None:
            config = Config(DEFAULT_CFG_PATH_DET).cfg

        # 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}'")

        config['Global']['device'] = device

        self._init_common(config)
        backend = backend if config['Global'].get(
            'backend', None) is None else config['Global']['backend']
        self.backend = backend
        if backend == 'torch':
            import torch
            self.torch = torch
            if config['Architecture']['algorithm'] == 'DB_mobile':
                if not os.path.exists(config['Global']['pretrained_model']):
                    config['Global'][
                        'pretrained_model'] = check_and_download_model(
                            MODEL_NAME_DET, DOWNLOAD_URL_DET)
            self._init_torch_model(config, numId)
        elif backend == 'onnx':
            from tools.infer.onnx_engine import ONNXEngine
            onnx_model_path = onnx_model_path if config['Global'].get(
                'onnx_model_path',
                None) is None else config['Global']['onnx_model_path']
            if onnx_model_path is None:
                if config['Architecture']['algorithm'] == 'DB_mobile':
                    onnx_model_path = check_and_download_model(
                        MODEL_NAME_DET_ONNX, DOWNLOAD_URL_DET_ONNX)