mirror of
https://github.com/RapidAI/RapidOCR.git
synced 2026-11-11 21:40:52 +08:00
86 lines
3.1 KiB
Python
86 lines
3.1 KiB
Python
# -*- encoding: utf-8 -*-
|
|
# @Author: SWHL
|
|
# @Contact: liekkaskono@163.com
|
|
import warnings
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import yaml
|
|
from onnxruntime import (GraphOptimizationLevel, InferenceSession,
|
|
SessionOptions, get_available_providers, get_device)
|
|
|
|
root_dir = Path(__file__).resolve().parent
|
|
|
|
|
|
class OrtInferSession():
|
|
def __init__(self, config):
|
|
sess_opt = SessionOptions()
|
|
sess_opt.log_severity_level = 4
|
|
sess_opt.enable_cpu_mem_arena = False
|
|
sess_opt.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL
|
|
|
|
cuda_ep = 'CUDAExecutionProvider'
|
|
cpu_ep = 'CPUExecutionProvider'
|
|
cpu_provider_options = {
|
|
"arena_extend_strategy": "kSameAsRequested",
|
|
}
|
|
|
|
EP_list = []
|
|
if config['use_cuda'] and get_device() == 'GPU' \
|
|
and cuda_ep in get_available_providers():
|
|
EP_list = [(cuda_ep, config[cuda_ep])]
|
|
EP_list.append((cpu_ep, cpu_provider_options))
|
|
|
|
config['model_path'] = str(root_dir / config['model_path'])
|
|
self._verify_model(config['model_path'])
|
|
self.session = InferenceSession(config['model_path'],
|
|
sess_options=sess_opt,
|
|
providers=EP_list)
|
|
|
|
if config['use_cuda'] and cuda_ep not in self.session.get_providers():
|
|
warnings.warn(f'{cuda_ep} is not avaiable for current env, the inference part is automatically shifted to be executed under {cpu_ep}.\n'
|
|
'Please ensure the installed onnxruntime-gpu version matches your cuda and cudnn version, '
|
|
'you can check their relations from the offical web site: '
|
|
'https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html',
|
|
RuntimeWarning)
|
|
|
|
def __call__(self, input_content: np.ndarray) -> np.ndarray:
|
|
input_dict = dict(zip(self.get_input_names(), [input_content]))
|
|
try:
|
|
return self.session.run(self.get_output_names(), input_dict)
|
|
except Exception as e:
|
|
raise ONNXRuntimeError('ONNXRuntime inferece failed.') from e
|
|
|
|
def get_input_names(self, ):
|
|
return [v.name for v in self.session.get_inputs()]
|
|
|
|
def get_output_names(self,):
|
|
return [v.name for v in self.session.get_outputs()]
|
|
|
|
def get_character_list(self, key: str = 'character'):
|
|
return self.meta_dict[key].splitlines()
|
|
|
|
def have_key(self, key: str = 'character') -> bool:
|
|
self.meta_dict = self.session.get_modelmeta().custom_metadata_map
|
|
if key in self.meta_dict.keys():
|
|
return True
|
|
return False
|
|
|
|
@staticmethod
|
|
def _verify_model(model_path):
|
|
model_path = Path(model_path)
|
|
if not model_path.exists():
|
|
raise FileNotFoundError(f'{model_path} does not exists.')
|
|
if not model_path.is_file():
|
|
raise FileExistsError(f'{model_path} is not a file.')
|
|
|
|
|
|
class ONNXRuntimeError(Exception):
|
|
pass
|
|
|
|
|
|
def read_yaml(yaml_path):
|
|
with open(yaml_path, 'rb') as f:
|
|
data = yaml.load(f, Loader=yaml.Loader)
|
|
return data
|