Browse Source

feat: refactor PDF image conversion to use get_end_page_id utility function and add multi-threading support

myhloli 3 weeks ago
parent
commit
66d5f3dfd2
3 changed files with 75 additions and 24 deletions
  1. 2 5
      mineru/cli/common.py
  2. 63 19
      mineru/utils/pdf_image_tools.py
  3. 10 0
      mineru/utils/pdf_page_id.py

+ 2 - 5
mineru/cli/common.py

@@ -16,6 +16,7 @@ from mineru.utils.pdf_image_tools import images_bytes_to_pdf_bytes
 from mineru.backend.vlm.vlm_middle_json_mkcontent import union_make as vlm_union_make
 from mineru.backend.vlm.vlm_analyze import doc_analyze as vlm_doc_analyze
 from mineru.backend.vlm.vlm_analyze import aio_doc_analyze as aio_vlm_doc_analyze
+from mineru.utils.pdf_page_id import get_end_page_id
 
 pdf_suffixes = ["pdf"]
 image_suffixes = ["png", "jpeg", "jp2", "webp", "gif", "bmp", "jpg", "tiff"]
@@ -49,11 +50,7 @@ def convert_pdf_bytes_to_bytes_by_pypdfium2(pdf_bytes, start_page_id=0, end_page
         # 从字节数据加载PDF
         pdf = pdfium.PdfDocument(pdf_bytes)
 
-        # 确定结束页
-        end_page_id = end_page_id if end_page_id is not None and end_page_id >= 0 else len(pdf) - 1
-        if end_page_id > len(pdf) - 1:
-            logger.warning("end_page_id is out of range, use pdf_docs length")
-            end_page_id = len(pdf) - 1
+        end_page_id = get_end_page_id(end_page_id, len(pdf))
 
         # 创建一个新的PDF文档
         output_pdf = pdfium.PdfDocument.new()

+ 63 - 19
mineru/utils/pdf_image_tools.py

@@ -8,8 +8,9 @@ from PIL import Image
 
 from mineru.data.data_reader_writer import FileBasedDataWriter
 from mineru.utils.pdf_reader import image_to_b64str, image_to_bytes, page_to_image
-from .enum_class import ImageType
-from .hash_utils import str_sha256
+from mineru.utils.enum_class import ImageType
+from mineru.utils.hash_utils import str_sha256
+from mineru.utils.pdf_page_id import get_end_page_id
 
 from concurrent.futures import ProcessPoolExecutor, TimeoutError as FuturesTimeoutError
 
@@ -36,7 +37,7 @@ def pdf_page_to_image(page: pdfium.PdfPage, dpi=200, image_type=ImageType.PIL) -
 
     return image_dict
 
-#@todo 通过多进程渲染pdf提速
+
 def _load_images_from_pdf_worker(pdf_bytes, dpi, start_page_id, end_page_id, image_type):
     """用于进程池的包装函数"""
     return load_images_from_pdf_core(pdf_bytes, dpi, start_page_id, end_page_id, image_type)
@@ -49,30 +50,74 @@ def load_images_from_pdf(
         end_page_id=None,
         image_type=ImageType.PIL,
         timeout=300,
+        threads=4,
 ):
-    """带超时控制的 PDF 转图片函数
+    """带超时控制的 PDF 转图片函数,支持多进程加速
 
     Args:
         timeout (int): 超时时间(秒),默认 300 秒
+        threads (int): 进程数,默认 4
 
     Raises:
         TimeoutError: 当转换超时时抛出
     """
-    with ProcessPoolExecutor(max_workers=1) as executor:
-        future = executor.submit(
-            _load_images_from_pdf_worker,
-            pdf_bytes,
-            dpi,
-            start_page_id,
-            end_page_id,
-            image_type
-        )
+    # 根据进程数调整超时时间
+    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 = min(start_page_id + (i + 1) * pages_per_thread - 1, end_page_id)
+
+        page_ranges.append((range_start, range_end))
+
+    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:
-            images_list = future.result(timeout=timeout)
-            pdf_doc = pdfium.PdfDocument(pdf_bytes)
+            # 收集结果并按页码排序
+            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")
 
@@ -87,10 +132,7 @@ def load_images_from_pdf_core(
     images_list = []
     pdf_doc = pdfium.PdfDocument(pdf_bytes)
     pdf_page_num = len(pdf_doc)
-    end_page_id = end_page_id if end_page_id is not None and end_page_id >= 0 else pdf_page_num - 1
-    if end_page_id > pdf_page_num - 1:
-        logger.warning("end_page_id is out of range, use images length")
-        end_page_id = pdf_page_num - 1
+    end_page_id = get_end_page_id(end_page_id, pdf_page_num)
 
     for index in range(0, pdf_page_num):
         if start_page_id <= index <= end_page_id:
@@ -98,6 +140,8 @@ def load_images_from_pdf_core(
             image_dict = pdf_page_to_image(page, dpi=dpi, image_type=image_type)
             images_list.append(image_dict)
 
+    pdf_doc.close()
+
     return images_list
 
 

+ 10 - 0
mineru/utils/pdf_page_id.py

@@ -0,0 +1,10 @@
+# Copyright (c) Opendatalab. All rights reserved.
+from loguru import logger
+
+
+def get_end_page_id(end_page_id, pdf_page_num):
+    end_page_id = end_page_id if end_page_id is not None and end_page_id >= 0 else pdf_page_num - 1
+    if end_page_id > pdf_page_num - 1:
+        logger.warning("end_page_id is out of range, use images length")
+        end_page_id = pdf_page_num - 1
+    return end_page_id