From 8bde98f1bbaf7e0ee41961a22ee8a1dc6cdcc06a Mon Sep 17 00:00:00 2001 From: SWHL Date: Tue, 11 Feb 2025 09:39:20 +0800 Subject: [PATCH] chore: update files --- python/rapidocr/config.yaml | 39 ++- python/rapidocr/inference_engine/__init__.py | 2 + .../rapidocr/inference_engine/onnxruntime.py | 6 +- .../rapidocr/inference_engine/paddlepaddle.py | 4 +- python/rapidocr/main.py | 3 + python/rapidocr/utils/parse_parameters.py | 249 +----------------- 6 files changed, 31 insertions(+), 272 deletions(-) diff --git a/python/rapidocr/config.yaml b/python/rapidocr/config.yaml index 76fe799..6d314c7 100644 --- a/python/rapidocr/config.yaml +++ b/python/rapidocr/config.yaml @@ -4,6 +4,7 @@ Global: use_det: true use_cls: true use_rec: true + print_verbose: false min_height: 30 width_height_ratio: 8 @@ -11,22 +12,26 @@ Global: min_side_len: 30 return_word_box: false - intra_op_num_threads: &intra_nums -1 - inter_op_num_threads: &inter_nums -1 - with_onnx: false - with_openvino: true + with_openvino: false with_paddle: false with_torch: false +InferEngineConfig: + onnxruntime: + intra_op_num_threads: -1 + inter_op_num_threads: -1 + use_cuda: false + use_dml: false + + paddlepaddle: + cpu_math_library_num_threads: -1 + gpu_id: 0 + gpu_mem: 500 + Det: - intra_op_num_threads: *intra_nums - inter_op_num_threads: *inter_nums - - use_cuda: false - use_dml: false - model_path: models/ch_PP-OCRv4_det_infer.onnx + paddle_model_dir: models/paddle/ch_PP-OCRv4_det_infer limit_side_len: 736 limit_type: min @@ -41,13 +46,8 @@ Det: score_mode: fast Cls: - intra_op_num_threads: *intra_nums - inter_op_num_threads: *inter_nums - - use_cuda: false - use_dml: false - model_path: models/ch_ppocr_mobile_v2.0_cls_infer.onnx + paddle_model_dir: models/paddle/ch_ppocr_mobile_v2_cls_infer cls_image_shape: [3, 48, 192] cls_batch_num: 6 @@ -55,13 +55,8 @@ Cls: label_list: ['0', '180'] Rec: - intra_op_num_threads: *intra_nums - inter_op_num_threads: *inter_nums - - use_cuda: false - use_dml: false - model_path: models/ch_PP-OCRv4_rec_infer.onnx + paddle_model_dir: models/paddle/ch_PP-OCRv4_rec_infer rec_img_shape: [3, 48, 320] rec_batch_num: 6 diff --git a/python/rapidocr/inference_engine/__init__.py b/python/rapidocr/inference_engine/__init__.py index 436d073..eda5cf5 100644 --- a/python/rapidocr/inference_engine/__init__.py +++ b/python/rapidocr/inference_engine/__init__.py @@ -19,6 +19,8 @@ class Engine(Enum): def get_engine(engine_name: str): + print("engine_name:", engine_name) + if engine_name == Engine.ONNXRUNTIME.value: return OrtInferSession diff --git a/python/rapidocr/inference_engine/onnxruntime.py b/python/rapidocr/inference_engine/onnxruntime.py index 09da63f..446ac3c 100644 --- a/python/rapidocr/inference_engine/onnxruntime.py +++ b/python/rapidocr/inference_engine/onnxruntime.py @@ -33,13 +33,13 @@ class OrtInferSession(InferSession): model_path = config.get("model_path", None) self._verify_model(model_path) - self.cfg_use_cuda = config.get("use_cuda", None) - self.cfg_use_dml = config.get("use_dml", None) + self.cfg_use_cuda = config.engine_cfg.get("use_cuda", None) + self.cfg_use_dml = config.engine_cfg.get("use_dml", None) self.had_providers: List[str] = get_available_providers() EP_list = self._get_ep_list() - sess_opt = self._init_sess_opts(config) + sess_opt = self._init_sess_opts(config.engine_cfg) self.session = InferenceSession( model_path, sess_options=sess_opt, diff --git a/python/rapidocr/inference_engine/paddlepaddle.py b/python/rapidocr/inference_engine/paddlepaddle.py index a212efb..076c7e5 100644 --- a/python/rapidocr/inference_engine/paddlepaddle.py +++ b/python/rapidocr/inference_engine/paddlepaddle.py @@ -18,7 +18,7 @@ class PaddleInferSession: self.logger = get_logger("PaddleInferSession") self.mode = mode - model_dir = Path(config["model_path"]) + model_dir = Path(config.padde_model_dir) pdmodel_path = model_dir / "inference.pdmodel" pdiparams_path = model_dir / "inference.pdiparams" @@ -65,7 +65,7 @@ class PaddleInferSession: outputs.append(output) self.predictor.try_shrink_memory() - return outputs + return outputs[0] @staticmethod def _verify_model(model_path): diff --git a/python/rapidocr/main.py b/python/rapidocr/main.py index f7b321b..2df4199 100644 --- a/python/rapidocr/main.py +++ b/python/rapidocr/main.py @@ -58,15 +58,18 @@ class RapidOCR: self.use_det = config.Global.use_det config.Det.lang = det_lang config.Det.engine_name = engine_name + config.Det.engine_cfg = config.InferEngineConfig[engine_name] self.text_det = TextDetector(config.Det) self.use_cls = config.Global.use_cls config.Cls.engine_name = engine_name + config.Cls.engine_cfg = config.InferEngineConfig[engine_name] self.text_cls = TextClassifier(config.Cls) self.use_rec = config.Global.use_rec config.Rec.lang = rec_lang config.Rec.engine_name = engine_name + config.Rec.engine_cfg = config.InferEngineConfig[engine_name] self.text_rec = TextRecognizer(config.Rec) self.load_img = LoadImage() diff --git a/python/rapidocr/utils/parse_parameters.py b/python/rapidocr/utils/parse_parameters.py index 88cf47f..f1d79cd 100644 --- a/python/rapidocr/utils/parse_parameters.py +++ b/python/rapidocr/utils/parse_parameters.py @@ -6,253 +6,12 @@ from pathlib import Path from typing import Dict, List, Optional, Union import numpy as np -from omegaconf import DictConfig +from omegaconf import DictConfig, OmegaConf from PIL import Image root_dir = Path(__file__).resolve().parent.parent InputType = Union[str, np.ndarray, bytes, Path, Image.Image] -DEFAULT_OCR_MODEL_VERSION = "PP-OCRv4" -SUPPORT_OCR_MODEL_VERSION = ["PP-OCR", "PP-OCRv2", "PP-OCRv3", "PP-OCRv4"] -MODEL_URLS = { - "OCR": { - "PP-OCRv4": { - "det": { - "ch": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv4/chinese/ch_PP-OCRv4_det_infer.tar", - }, - "en": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/english/en_PP-OCRv3_det_infer.tar", - }, - "ml": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/Multilingual_PP-OCRv3_det_infer.tar" - }, - }, - "rec": { - "ch": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv4/chinese/ch_PP-OCRv4_rec_infer.tar", - "dict_path": "./ppocr/utils/ppocr_keys_v1.txt", - }, - "en": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv4/english/en_PP-OCRv4_rec_infer.tar", - "dict_path": "./ppocr/utils/en_dict.txt", - }, - "korean": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv4/multilingual/korean_PP-OCRv4_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/korean_dict.txt", - }, - "japan": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv4/multilingual/japan_PP-OCRv4_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/japan_dict.txt", - }, - "chinese_cht": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/chinese_cht_PP-OCRv3_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/chinese_cht_dict.txt", - }, - "ta": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv4/multilingual/ta_PP-OCRv4_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/ta_dict.txt", - }, - "te": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv4/multilingual/te_PP-OCRv4_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/te_dict.txt", - }, - "ka": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv4/multilingual/ka_PP-OCRv4_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/ka_dict.txt", - }, - "latin": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/latin_PP-OCRv3_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/latin_dict.txt", - }, - "arabic": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv4/multilingual/arabic_PP-OCRv4_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/arabic_dict.txt", - }, - "cyrillic": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/cyrillic_PP-OCRv3_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/cyrillic_dict.txt", - }, - "devanagari": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv4/multilingual/devanagari_PP-OCRv4_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/devanagari_dict.txt", - }, - }, - "cls": { - "ch": { - "url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_cls_infer.tar", - } - }, - }, - "PP-OCRv3": { - "det": { - "ch": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/chinese/ch_PP-OCRv3_det_infer.tar", - }, - "en": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/english/en_PP-OCRv3_det_infer.tar", - }, - "ml": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/Multilingual_PP-OCRv3_det_infer.tar" - }, - }, - "rec": { - "ch": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/chinese/ch_PP-OCRv3_rec_infer.tar", - "dict_path": "./ppocr/utils/ppocr_keys_v1.txt", - }, - "en": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/english/en_PP-OCRv3_rec_infer.tar", - "dict_path": "./ppocr/utils/en_dict.txt", - }, - "korean": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/korean_PP-OCRv3_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/korean_dict.txt", - }, - "japan": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/japan_PP-OCRv3_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/japan_dict.txt", - }, - "chinese_cht": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/chinese_cht_PP-OCRv3_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/chinese_cht_dict.txt", - }, - "ta": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/ta_PP-OCRv3_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/ta_dict.txt", - }, - "te": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/te_PP-OCRv3_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/te_dict.txt", - }, - "ka": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/ka_PP-OCRv3_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/ka_dict.txt", - }, - "latin": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/latin_PP-OCRv3_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/latin_dict.txt", - }, - "arabic": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/arabic_PP-OCRv3_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/arabic_dict.txt", - }, - "cyrillic": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/cyrillic_PP-OCRv3_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/cyrillic_dict.txt", - }, - "devanagari": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv3/multilingual/devanagari_PP-OCRv3_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/devanagari_dict.txt", - }, - }, - "cls": { - "ch": { - "url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_cls_infer.tar", - } - }, - }, - "PP-OCRv2": { - "det": { - "ch": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv2/chinese/ch_PP-OCRv2_det_infer.tar", - }, - }, - "rec": { - "ch": { - "url": "https://paddleocr.bj.bcebos.com/PP-OCRv2/chinese/ch_PP-OCRv2_rec_infer.tar", - "dict_path": "./ppocr/utils/ppocr_keys_v1.txt", - } - }, - "cls": { - "ch": { - "url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_cls_infer.tar", - } - }, - }, - "PP-OCR": { - "det": { - "ch": { - "url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_det_infer.tar", - }, - "en": { - "url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/en_ppocr_mobile_v2.0_det_infer.tar", - }, - "structure": { - "url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_table_det_infer.tar" - }, - }, - "rec": { - "ch": { - "url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_rec_infer.tar", - "dict_path": "./ppocr/utils/ppocr_keys_v1.txt", - }, - "en": { - "url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/en_number_mobile_v2.0_rec_infer.tar", - "dict_path": "./ppocr/utils/en_dict.txt", - }, - "french": { - "url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/french_mobile_v2.0_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/french_dict.txt", - }, - "german": { - "url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/german_mobile_v2.0_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/german_dict.txt", - }, - "korean": { - "url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/korean_mobile_v2.0_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/korean_dict.txt", - }, - "japan": { - "url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/japan_mobile_v2.0_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/japan_dict.txt", - }, - "chinese_cht": { - "url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/chinese_cht_mobile_v2.0_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/chinese_cht_dict.txt", - }, - "ta": { - "url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/ta_mobile_v2.0_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/ta_dict.txt", - }, - "te": { - "url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/te_mobile_v2.0_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/te_dict.txt", - }, - "ka": { - "url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/ka_mobile_v2.0_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/ka_dict.txt", - }, - "latin": { - "url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/latin_ppocr_mobile_v2.0_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/latin_dict.txt", - }, - "arabic": { - "url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/arabic_ppocr_mobile_v2.0_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/arabic_dict.txt", - }, - "cyrillic": { - "url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/cyrillic_ppocr_mobile_v2.0_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/cyrillic_dict.txt", - }, - "devanagari": { - "url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/devanagari_ppocr_mobile_v2.0_rec_infer.tar", - "dict_path": "./ppocr/utils/dict/devanagari_dict.txt", - }, - "structure": { - "url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_table_rec_infer.tar", - "dict_path": "ppocr/utils/dict/table_dict.txt", - }, - }, - "cls": { - "ch": { - "url": "https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_cls_infer.tar", - } - }, - }, - } -} - def parse_lang(lang): latin_lang = [ @@ -360,9 +119,9 @@ def parse_lang(lang): def update_model_path(config: DictConfig) -> DictConfig: key = "model_path" - config["Det"][key] = str(root_dir / config["Det"][key]) - config["Rec"][key] = str(root_dir / config["Rec"][key]) - config["Cls"][key] = str(root_dir / config["Cls"][key]) + OmegaConf.update(config, f"Det.{key}", str(root_dir / config.Det[key])) + OmegaConf.update(config, f"Cls.{key}", str(root_dir / config.Cls[key])) + OmegaConf.update(config, f"Rec.{key}", str(root_dir / config.Rec[key])) return config