chore: update files

RapidAI · Feb 11, 2025 · 8bde98f · 8bde98f
1 parent d6cd486
commit 8bde98f
Show file tree

Hide file tree

Showing 6 changed files with 30 additions and 271 deletions.
diff --git a/python/rapidocr/config.yaml b/python/rapidocr/config.yaml
@@ -4,29 +4,34 @@ Global:
     use_det: true
     use_cls: true
     use_rec: true
+
     print_verbose: false
     min_height: 30
     width_height_ratio: 8
     max_side_len: 2000
     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
 
-Det:
-    intra_op_num_threads: *intra_nums
-    inter_op_num_threads: *inter_nums
+InferEngineConfig:
+    onnxruntime:
+        intra_op_num_threads: -1
+        inter_op_num_threads: -1
+        use_cuda: false
+        use_dml: false
 
-    use_cuda: false
-    use_dml: false
+    paddlepaddle:
+        cpu_math_library_num_threads: -1
+        gpu_id: 0
+        gpu_mem: 500
 
+Det:
     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,27 +46,17 @@ 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
     cls_thresh: 0.9
     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
@@ -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
@@ -33,13 +33,13 @@ def __init__(self, config: Dict[str, Any]):
         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
@@ -18,7 +18,7 @@ def __init__(self, config, mode: Optional[str] = None) -> None:
         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 @@ def __call__(self, img: np.ndarray):
             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
@@ -58,15 +58,18 @@ def __init__(self, config_path: Optional[str] = None, **kwargs):
         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()