Browse Source

Merge pull request #3931 from myhloli/dev

Dev
Xiaomeng Zhao 3 weeks ago
parent
commit
f1ebf5a7f0

+ 11 - 0
docs/en/usage/cli_tools.md

@@ -100,3 +100,14 @@ Here are the environment variables and their descriptions:
     * Used to enable table merging functionality
     * Default is `true`, can be set to `false` via environment variable to disable table merging functionality.
 
+- `MINERU_PDF_LOAD_IMAGES_TIMEOUT`:
+    * Used to set the timeout period (in seconds) for rendering PDF to images
+    * Default is `300` seconds, can be set to other values via environment variable to adjust the image rendering timeout.
+
+- `MINERU_INTRA_OP_NUM_THREADS`:
+    * Used to set the intra_op thread count for ONNX models, affects the computation speed of individual operators
+    * Default is `-1` (auto-select), can be set to other values via environment variable to adjust the thread count.
+
+- `MINERU_INTER_OP_NUM_THREADS`:
+    * Used to set the inter_op thread count for ONNX models, affects the parallel execution of multiple operators
+    * Default is `-1` (auto-select), can be set to other values via environment variable to adjust the thread count.

+ 12 - 0
docs/zh/usage/cli_tools.md

@@ -94,3 +94,15 @@ MinerU命令行工具的某些参数存在相同功能的环境变量配置,
 - `MINERU_TABLE_MERGE_ENABLE`:
     * 用于启用表格合并功能
     * 默认为`true`,可通过环境变量设置为`false`来禁用表格合并功能。
+
+- `MINERU_PDF_LOAD_IMAGES_TIMEOUT`:
+    * 用于设置将PDF渲染为图片的超时时间(秒)
+    * 默认为`300`秒,可通过环境变量设置为其他值以调整渲染图片的超时时间。
+
+- `MINERU_INTRA_OP_NUM_THREADS`:
+    * 用于设置onnx模型的intra_op线程数,影响单个算子的计算速度
+    * 默认为`-1`(自动选择),可通过环境变量设置为其他值以调整线程数。
+
+- `MINERU_INTER_OP_NUM_THREADS`:
+    * 用于设置onnx模型的inter_op线程数,影响多个算子的并行执行
+    * 默认为`-1`(自动选择),可通过环境变量设置为其他值以调整线程数。

+ 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
 
+from mineru.utils.os_env_config import get_op_num_threads
 from .table_structure_utils import (
     OrtInferSession,
     TableLabelDecode,
@@ -29,6 +30,9 @@ class TableStructurer:
         self.preprocess_op = TablePreprocess()
         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.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 numpy as np
 from skimage import measure
+
+from mineru.utils.os_env_config import get_op_num_threads
 from .utils import OrtInferSession, resize_img
 from .utils_table_line_rec import (
     get_table_line,
@@ -28,6 +30,9 @@ class TSRUnet:
         self.inp_height = 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)
 
     def __call__(

+ 30 - 0
mineru/utils/os_env_config.py

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

+ 11 - 6
mineru/utils/pdf_image_tools.py

@@ -9,6 +9,7 @@ from PIL import Image
 
 from mineru.data.data_reader_writer import FileBasedDataWriter
 from mineru.utils.check_sys_env import is_windows_environment
+from mineru.utils.os_env_config import get_load_images_timeout
 from mineru.utils.pdf_reader import image_to_b64str, image_to_bytes, page_to_image
 from mineru.utils.enum_class import ImageType
 from mineru.utils.hash_utils import str_sha256
@@ -51,13 +52,18 @@ def load_images_from_pdf(
         start_page_id=0,
         end_page_id=None,
         image_type=ImageType.PIL,
-        timeout=300,
+        timeout=None,
         threads=4,
 ):
     """带超时控制的 PDF 转图片函数,支持多进程加速
 
     Args:
-        timeout (int): 超时时间(秒),默认 300 秒
+        pdf_bytes (bytes): PDF 文件的 bytes
+        dpi (int, optional): reset the dpi of dpi. Defaults to 200.
+        start_page_id (int, optional): 起始页码. Defaults to 0.
+        end_page_id (int | None, optional): 结束页码. Defaults to None.
+        image_type (ImageType, optional): 图片类型. Defaults to ImageType.PIL.
+        timeout (int | None, optional): 超时时间(秒)。如果为 None,则从环境变量 MINERU_PDF_LOAD_IMAGES_TIMEOUT 读取,若未设置则默认为 300 秒。
         threads (int): 进程数,默认 4
 
     Raises:
@@ -74,15 +80,15 @@ def load_images_from_pdf(
             image_type
         ), pdf_doc
     else:
-        # 根据进程数调整超时时间
-        threads = min(os.cpu_count() or 1, threads)
+        if timeout is None:
+            timeout = get_load_images_timeout()
         end_page_id = get_end_page_id(end_page_id, len(pdf_doc))
 
         # 计算总页数
         total_pages = end_page_id - start_page_id + 1
 
         # 实际使用的进程数不超过总页数
-        actual_threads = min(threads, total_pages)
+        actual_threads = min(min(os.cpu_count() or 1, threads), total_pages)
 
         # 根据实际进程数分组页面范围
         pages_per_thread = max(1, total_pages // actual_threads)
@@ -129,7 +135,6 @@ def load_images_from_pdf(
 
                 return images_list, pdf_doc
             except FuturesTimeoutError:
-                logger.error(f"PDF conversion timeout after {timeout}s")
                 pdf_doc.close()
                 executor.shutdown(wait=False, cancel_futures=True)
                 raise TimeoutError(f"PDF to images conversion timeout after {timeout}s")