diff --git a/python/rapidocr_openvino/config.yaml b/python/rapidocr_openvino/config.yaml index 3bd206126..81aee2083 100644 --- a/python/rapidocr_openvino/config.yaml +++ b/python/rapidocr_openvino/config.yaml @@ -7,8 +7,14 @@ Global: min_height: 30 width_height_ratio: 8 + use_cuda: &use_cuda false + + inference_num_threads: &infer_num_threads -1 + Det: - use_cuda: false + inference_num_threads: *infer_num_threads + + use_cuda: *use_cuda model_path: models/ch_PP-OCRv4_det_infer.onnx @@ -23,7 +29,9 @@ Det: score_mode: fast Cls: - use_cuda: false + inference_num_threads: *infer_num_threads + + use_cuda: *use_cuda model_path: models/ch_ppocr_mobile_v2.0_cls_infer.onnx @@ -33,7 +41,9 @@ Cls: label_list: ['0', '180'] Rec: - use_cuda: false + inference_num_threads: *infer_num_threads + + use_cuda: *use_cuda model_path: models/ch_PP-OCRv4_rec_infer.onnx diff --git a/python/rapidocr_openvino/utils.py b/python/rapidocr_openvino/utils.py index 21d6c3b76..ad2090da2 100644 --- a/python/rapidocr_openvino/utils.py +++ b/python/rapidocr_openvino/utils.py @@ -3,6 +3,7 @@ # @Contact: liekkaskono@163.com import argparse import math +import os import random from io import BytesIO from pathlib import Path @@ -20,11 +21,17 @@ class OpenVINOInferSession: def __init__(self, config): - ie = Core() + core = Core() self._verify_model(config["model_path"]) - model_onnx = ie.read_model(config["model_path"]) - compile_model = ie.compile_model(model=model_onnx, device_name="CPU") + 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: diff --git a/python/rapidocr_paddle/config.yaml b/python/rapidocr_paddle/config.yaml index 3ed6a4620..c30b4fda4 100644 --- a/python/rapidocr_paddle/config.yaml +++ b/python/rapidocr_paddle/config.yaml @@ -7,10 +7,18 @@ Global: min_height: 30 width_height_ratio: 8 + use_cuda: &use_cuda false + gpu_id: &gpu_id 0 + gpu_mem: &gpu_mem 500 + + cpu_math_library_num_threads: &infer_num_threads -1 + Det: - use_cuda: false - gpu_id: 0 - gpu_mem: 500 + use_cuda: *use_cuda + gpu_id: *gpu_id + gpu_mem: *gpu_mem + + cpu_math_library_num_threads: *infer_num_threads model_path: models/ch_PP-OCRv4_det_infer @@ -25,9 +33,11 @@ Det: score_mode: fast Cls: - use_cuda: false - gpu_id: 0 - gpu_mem: 500 + use_cuda: *use_cuda + gpu_id: *gpu_id + gpu_mem: *gpu_mem + + cpu_math_library_num_threads: *infer_num_threads model_path: models/ch_ppocr_mobile_v2_cls_infer @@ -37,9 +47,11 @@ Cls: label_list: ['0', '180'] Rec: - use_cuda: false - gpu_id: 0 - gpu_mem: 500 + use_cuda: *use_cuda + gpu_id: *gpu_id + gpu_mem: *gpu_mem + + cpu_math_library_num_threads: *infer_num_threads model_path: models/ch_PP-OCRv4_rec_infer diff --git a/python/rapidocr_paddle/utils.py b/python/rapidocr_paddle/utils.py index 6169da12e..cc4a33935 100644 --- a/python/rapidocr_paddle/utils.py +++ b/python/rapidocr_paddle/utils.py @@ -45,6 +45,11 @@ def __init__(self, config, mode: Optional[str] = None) -> None: else: infer_opts.disable_gpu() + cpu_nums = os.cpu_count() + infer_num_threads = config.get("cpu_math_library_num_threads", -1) + if infer_num_threads != -1 and 1 <= infer_num_threads <= cpu_nums: + infer_opts.set_cpu_math_library_num_threads(infer_num_threads) + # enable memory optim infer_opts.enable_memory_optim() infer_opts.disable_glog_info()