Browse Source

refactor: adjust thread count based on CPU cores and comment out image loading time logging

myhloli 3 weeks ago
parent
commit
2079395774
2 changed files with 5 additions and 3 deletions
  1. 3 3
      mineru/backend/pipeline/pipeline_analyze.py
  2. 2 0
      mineru/utils/pdf_image_tools.py

+ 3 - 3
mineru/backend/pipeline/pipeline_analyze.py

@@ -99,10 +99,10 @@ def doc_analyze(
         _lang = lang_list[pdf_idx]
 
         # 收集每个数据集中的页面
-        load_images_start = time.time()
+        # load_images_start = time.time()
         images_list, pdf_doc = load_images_from_pdf(pdf_bytes, image_type=ImageType.PIL)
-        load_images_time = round(time.time() - load_images_start, 2)
-        logger.debug(f"load images cost: {load_images_time}, speed: {round(len(images_list) / load_images_time, 3)} images/s")
+        # load_images_time = round(time.time() - load_images_start, 2)
+        # logger.debug(f"load images cost: {load_images_time}, speed: {round(len(images_list) / load_images_time, 3)} images/s")
         all_image_lists.append(images_list)
         all_pdf_docs.append(pdf_doc)
         for page_idx in range(len(images_list)):

+ 2 - 0
mineru/utils/pdf_image_tools.py

@@ -1,4 +1,5 @@
 # Copyright (c) Opendatalab. All rights reserved.
+import os
 from io import BytesIO
 
 import numpy as np
@@ -74,6 +75,7 @@ def load_images_from_pdf(
         ), pdf_doc
     else:
         # 根据进程数调整超时时间
+        threads = min(os.cpu_count(), threads)
         timeout = timeout // threads
 
         end_page_id = get_end_page_id(end_page_id, len(pdf_doc))