mirror of
https://github.com/RapidAI/RapidOCR.git
synced 2026-11-11 13:30:51 +08:00
46 lines
1.4 KiB
Python
46 lines
1.4 KiB
Python
# -*- encoding: utf-8 -*-
|
|
# @Author: SWHL
|
|
# @Contact: liekkaskono@163.com
|
|
import os
|
|
import traceback
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
from openvino.runtime import Core
|
|
|
|
|
|
class OpenVINOInferSession:
|
|
def __init__(self, config):
|
|
core = Core()
|
|
|
|
self._verify_model(config["model_path"])
|
|
model_onnx = core.read_model(config["model_path"])
|
|
|
|
cpu_nums = os.cpu_count()
|
|
infer_num_threads = config.get("inference_num_threads", -1)
|
|
if infer_num_threads != -1 and 1 <= infer_num_threads <= cpu_nums:
|
|
core.set_property("CPU", {"INFERENCE_NUM_THREADS": str(infer_num_threads)})
|
|
|
|
compile_model = core.compile_model(model=model_onnx, device_name="CPU")
|
|
self.session = compile_model.create_infer_request()
|
|
|
|
def __call__(self, input_content: np.ndarray) -> np.ndarray:
|
|
try:
|
|
self.session.infer(inputs=[input_content])
|
|
return self.session.get_output_tensor().data
|
|
except Exception as e:
|
|
error_info = traceback.format_exc()
|
|
raise OpenVIONError(error_info) from e
|
|
|
|
@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 OpenVIONError(Exception):
|
|
pass
|