Merge pull request #1582 from myhloli/dev

fix(ocr): improve ONNX model initialization and error handling

magic_pdf/model/sub_modules/ocr/paddleocr/ocr_utils.py

@@ -7,6 +7,8 @@ import base64
 from magic_pdf.libs.boxbase import __is_overlaps_y_exceeds_threshold
 from magic_pdf.pre_proc.ocr_dict_merge import merge_spans_to_line
 
+import importlib.resources
+from paddleocr import PaddleOCR
 from ppocr.utils.utility import check_and_read
 
 
@@ -327,11 +329,12 @@ class ONNXModelSingleton:
         return self._models[key]
 
 def onnx_model_init(key):
+    if len(key) < 4:
+        logger.error('Invalid key length, expected at least 4 elements')
+        exit(1)
 
-    import importlib.resources
-
-    with importlib.resources.path('rapidocr_onnxruntime.models','') as resource_path:
-        onnx_model = None
+    try:
+        with importlib.resources.path('rapidocr_onnxruntime.models', '') as resource_path:
             additional_ocr_params = {
                 "use_onnx": True,
                 "det_model_dir": f'{resource_path}/ch_PP-OCRv4_det_infer.onnx',
@@ -341,12 +344,12 @@ def onnx_model_init(key):
                 "use_dilation": key[2],
                 "det_db_unclip_ratio": key[3],
             }
-        # logger.info(f"additional_ocr_params: {additional_ocr_params}")
 
             if key[0] is not None:
                 additional_ocr_params["lang"] = key[0]
 
-    from paddleocr import PaddleOCR
+            # logger.info(f"additional_ocr_params: {additional_ocr_params}")
+
             onnx_model = PaddleOCR(**additional_ocr_params)
 
             if onnx_model is None:
@@ -354,3 +357,7 @@ def onnx_model_init(key):
                 exit(1)
             else:
                 return onnx_model
+
+    except Exception as e:
+        logger.exception(f'Error initializing model: {e}')
+        exit(1)
\ No newline at end of file