Browse Source

refactor: enhance PDF image loading by removing multiprocessing for Windows environment and improving logging

myhloli 4 weeks ago
parent
commit
5349fd7ccd
2 changed files with 74 additions and 63 deletions
  1. 3 3
      mineru/backend/pipeline/pipeline_analyze.py
  2. 71 60
      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)):

+ 71 - 60
mineru/utils/pdf_image_tools.py

@@ -7,6 +7,7 @@ from loguru import logger
 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.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
@@ -61,67 +62,77 @@ def load_images_from_pdf(
     Raises:
         TimeoutError: 当转换超时时抛出
     """
-    # 根据进程数调整超时时间
-    timeout = timeout // threads
-
     pdf_doc = pdfium.PdfDocument(pdf_bytes)
-    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)
-
-    # 根据实际进程数分组页面范围
-    pages_per_thread = max(1, total_pages // actual_threads)
-    page_ranges = []
-
-    for i in range(actual_threads):
-        range_start = start_page_id + i * pages_per_thread
-        if i == actual_threads - 1:
-            # 最后一个线程处理剩余所有页面
-            range_end = end_page_id
-        else:
-            range_end = start_page_id + (i + 1) * pages_per_thread - 1
-
-        page_ranges.append((range_start, range_end))
-
-    # logger.debug(f"PDF to images using {actual_threads} processes, page ranges: {page_ranges}")
-
-    with ProcessPoolExecutor(max_workers=actual_threads) as executor:
-        # 提交所有任务
-        futures = []
-        for range_start, range_end in page_ranges:
-            future = executor.submit(
-                _load_images_from_pdf_worker,
-                pdf_bytes,
-                dpi,
-                range_start,
-                range_end,
-                image_type
-            )
-            futures.append((range_start, future))
-
-        try:
-            # 收集结果并按页码排序
-            all_results = []
-            for range_start, future in futures:
-                images_list = future.result(timeout=timeout)
-                all_results.append((range_start, images_list))
-
-            # 按起始页码排序并合并结果
-            all_results.sort(key=lambda x: x[0])
-            images_list = []
-            for _, imgs in all_results:
-                images_list.extend(imgs)
-
-            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")
+    if is_windows_environment():
+        # Windows 环境下不使用多进程
+        return load_images_from_pdf_core(
+            pdf_bytes,
+            dpi,
+            start_page_id,
+            end_page_id,
+            image_type
+        ), pdf_doc
+    else:
+        # 根据进程数调整超时时间
+        timeout = timeout // threads
+
+        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)
+
+        # 根据实际进程数分组页面范围
+        pages_per_thread = max(1, total_pages // actual_threads)
+        page_ranges = []
+
+        for i in range(actual_threads):
+            range_start = start_page_id + i * pages_per_thread
+            if i == actual_threads - 1:
+                # 最后一个线程处理剩余所有页面
+                range_end = end_page_id
+            else:
+                range_end = start_page_id + (i + 1) * pages_per_thread - 1
+
+            page_ranges.append((range_start, range_end))
+
+        # logger.debug(f"PDF to images using {actual_threads} processes, page ranges: {page_ranges}")
+
+        with ProcessPoolExecutor(max_workers=actual_threads) as executor:
+            # 提交所有任务
+            futures = []
+            for range_start, range_end in page_ranges:
+                future = executor.submit(
+                    _load_images_from_pdf_worker,
+                    pdf_bytes,
+                    dpi,
+                    range_start,
+                    range_end,
+                    image_type
+                )
+                futures.append((range_start, future))
+
+            try:
+                # 收集结果并按页码排序
+                all_results = []
+                for range_start, future in futures:
+                    images_list = future.result(timeout=timeout)
+                    all_results.append((range_start, images_list))
+
+                # 按起始页码排序并合并结果
+                all_results.sort(key=lambda x: x[0])
+                images_list = []
+                for _, imgs in all_results:
+                    images_list.extend(imgs)
+
+                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")
 
 
 def load_images_from_pdf_core(