Browse Source

feat: add ONNX configuration for thread management and integrate into table structure

myhloli 3 weeks ago
parent
commit
be2369bdd4

+ 4 - 0
mineru/model/table/rec/slanet_plus/table_structure.py

@@ -16,6 +16,7 @@ from typing import Any, Dict, List, Tuple
 
 
 import numpy as np
 import numpy as np
 
 
+from mineru.model.utils.onnx_config import get_op_num_threads
 from .table_structure_utils import (
 from .table_structure_utils import (
     OrtInferSession,
     OrtInferSession,
     TableLabelDecode,
     TableLabelDecode,
@@ -29,6 +30,9 @@ class TableStructurer:
         self.preprocess_op = TablePreprocess()
         self.preprocess_op = TablePreprocess()
         self.batch_preprocess_op = BatchTablePreprocess()
         self.batch_preprocess_op = BatchTablePreprocess()
 
 
+        config["intra_op_num_threads"] = get_op_num_threads("MINERU_INTRA_OP_NUM_THREADS")
+        config["inter_op_num_threads"] = get_op_num_threads("MINERU_INTER_OP_NUM_THREADS")
+
         self.session = OrtInferSession(config)
         self.session = OrtInferSession(config)
 
 
         self.character = self.session.get_metadata()
         self.character = self.session.get_metadata()

+ 5 - 0
mineru/model/table/rec/unet_table/table_structure_unet.py

@@ -5,6 +5,8 @@ from typing import Optional, Dict, Any, Tuple
 import cv2
 import cv2
 import numpy as np
 import numpy as np
 from skimage import measure
 from skimage import measure
+
+from mineru.model.utils.onnx_config import get_op_num_threads
 from .utils import OrtInferSession, resize_img
 from .utils import OrtInferSession, resize_img
 from .utils_table_line_rec import (
 from .utils_table_line_rec import (
     get_table_line,
     get_table_line,
@@ -28,6 +30,9 @@ class TSRUnet:
         self.inp_height = 1024
         self.inp_height = 1024
         self.inp_width = 1024
         self.inp_width = 1024
 
 
+        config["intra_op_num_threads"] = get_op_num_threads("MINERU_INTRA_OP_NUM_THREADS")
+        config["inter_op_num_threads"] = get_op_num_threads("MINERU_INTER_OP_NUM_THREADS")
+
         self.session = OrtInferSession(config)
         self.session = OrtInferSession(config)
 
 
     def __call__(
     def __call__(

+ 24 - 0
mineru/model/utils/onnx_config.py

@@ -0,0 +1,24 @@
+import os
+
+
+def get_op_num_threads(env_name: str) -> int:
+    env_value = os.getenv(env_name, None)
+    return get_op_num_threads_from_value(env_value)
+
+
+def get_op_num_threads_from_value(env_value: str) -> int:
+    if env_value is not None:
+        try:
+            num_threads = int(env_value)
+            if num_threads > 0:
+                return num_threads
+        except ValueError:
+            return -1
+    return -1
+
+
+if __name__ == '__main__':
+    print(get_op_num_threads_from_value('1'))
+    print(get_op_num_threads_from_value('0'))
+    print(get_op_num_threads_from_value('-1'))
+    print(get_op_num_threads_from_value('abc'))