fix: fixed issue #114

SWHL · SWHL · commit 7a871ef0b846 · 2025-07-16T09:31:59.000+08:00
diff --git a/demo.py b/demo.py
@@ -1,11 +1,17 @@
 # -*- encoding: utf-8 -*-
 # @Author: SWHL
 # @Contact: liekkaskono@163.com
-from rapidocr import RapidOCR
+from rapidocr import EngineType, RapidOCR
 
 from rapid_table import ModelType, RapidTable, RapidTableInput
 
-ocr_engine = RapidOCR()
+ocr_engine = RapidOCR(
+    params={
+        "Det.engine_type": EngineType.TORCH,
+        "Cls.engine_type": EngineType.TORCH,
+        "Rec.engine_type": EngineType.TORCH,
+    }
+)
 
 input_args = RapidTableInput(model_type=ModelType.UNITABLE)
 table_engine = RapidTable(input_args)
diff --git a/rapid_table/main.py b/rapid_table/main.py
@@ -5,7 +5,7 @@
 import time
 from dataclasses import asdict
 from pathlib import Path
-from typing import List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import numpy as np
 
@@ -15,7 +15,6 @@
     LoadImage,
     Logger,
     ModelType,
-    EngineType,
     RapidTableInput,
     RapidTableOutput,
     get_boxes_recs,
@@ -39,32 +38,21 @@ def __init__(self, cfg: Optional[RapidTableInput] = None):
 
         self.ocr_engine = None
         if cfg.use_ocr:
-            self.ocr_engine = self._init_ocr_engine()
+            self.ocr_engine = self._init_ocr_engine(self.cfg.ocr_params)
 
         self.table_matcher = TableMatch()
         self.load_img = LoadImage()
 
-    def _init_ocr_engine(self):
-        try:
-            rapidocr_ = import_package("rapidocr")
-            if rapidocr_ is None:
-                raise ModuleNotFoundError("rapidocr package is not installed")
-            if self.cfg.engine_type == EngineType.TORCH:
-                EngineType_RapidOCR = rapidocr_.EngineType
-                return rapidocr_.RapidOCR(
-                    params={
-                        "Det.engine_type": EngineType_RapidOCR.TORCH,
-                        "Cls.engine_type": EngineType_RapidOCR.TORCH,
-                        "Rec.engine_type": EngineType_RapidOCR.TORCH,
-                        "EngineConfig.torch.use_cuda": True,  # Use torch GPU to infer
-                        "EngineConfig.torch.gpu_id": 0  # Specify GPU id
-                    }
-                )
-            return rapidocr_.RapidOCR()
-        except ModuleNotFoundError:
+    def _init_ocr_engine(self, params: Dict[Any, Any]):
+        rapidocr_ = import_package("rapidocr")
+        if rapidocr_ is None:
             logger.warning("rapidocr package is not installed, only table rec")
             return None
 
+        if not params:
+            return rapidocr_.RapidOCR()
+        return rapidocr_.RapidOCR(params=params)
+
     def _init_table_structer(self):
         if self.cfg.model_type == ModelType.UNITABLE:
             from .table_structure.unitable import UniTableStructure
diff --git a/rapid_table/utils/typings.py b/rapid_table/utils/typings.py
@@ -29,11 +29,12 @@ class RapidTableInput:
     model_type: Optional[ModelType] = ModelType.SLANETPLUS
     model_dir_or_path: Union[str, Path, None, Dict[str, str]] = None
 
-    use_ocr: bool = True
-
     engine_type: Optional[EngineType] = None
     engine_cfg: dict = field(default_factory=dict)
 
+    use_ocr: bool = True
+    ocr_params: dict = field(default_factory=dict)
+
 
 @dataclass
 class RapidTableOutput:
