Browse Source

Merge pull request #3918 from myhloli/dev

Dev
Xiaomeng Zhao 3 weeks ago
parent
commit
f7225d8e17

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

@@ -99,7 +99,10 @@ def doc_analyze(
         _lang = lang_list[pdf_idx]
 
         # 收集每个数据集中的页面
+        # 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")
         all_image_lists.append(images_list)
         all_pdf_docs.append(pdf_doc)
         for page_idx in range(len(images_list)):

+ 2 - 2
mineru/backend/vlm/vlm_analyze.py

@@ -8,7 +8,7 @@ from .utils import enable_custom_logits_processors, set_default_gpu_memory_utili
 from .model_output_to_middle_json import result_to_middle_json
 from ...data.data_reader_writer import DataWriter
 from mineru.utils.pdf_image_tools import load_images_from_pdf
-from ...utils.check_mac_env import is_mac_os_version_supported
+from ...utils.check_sys_env import is_mac_os_version_supported
 from ...utils.config_reader import get_device
 
 from ...utils.enum_class import ImageType
@@ -177,7 +177,7 @@ async def aio_doc_analyze(
     images_list, pdf_doc = load_images_from_pdf(pdf_bytes, image_type=ImageType.PIL)
     images_pil_list = [image_dict["img_pil"] for image_dict in images_list]
     # load_images_time = round(time.time() - load_images_start, 2)
-    # logger.info(f"load images cost: {load_images_time}, speed: {round(len(images_base64_list)/load_images_time, 3)} images/s")
+    # logger.debug(f"load images cost: {load_images_time}, speed: {round(len(images_pil_list)/load_images_time, 3)} images/s")
 
     # infer_start = time.time()
     results = await predictor.aio_batch_two_step_extract(images=images_pil_list)

+ 1 - 1
mineru/cli/client.py

@@ -4,7 +4,7 @@ import click
 from pathlib import Path
 from loguru import logger
 
-from mineru.utils.check_mac_env import is_mac_os_version_supported
+from mineru.utils.check_sys_env import is_mac_os_version_supported
 from mineru.utils.cli_parser import arg_parse
 from mineru.utils.config_reader import get_device
 from mineru.utils.guess_suffix_or_lang import guess_suffix_by_path

+ 8 - 15
mineru/cli/common.py

@@ -5,8 +5,8 @@ import os
 import copy
 from pathlib import Path
 
-import pypdfium2 as pdfium
 from loguru import logger
+import pypdfium2 as pdfium
 
 from mineru.data.data_reader_writer import FileBasedDataWriter
 from mineru.utils.draw_bbox import draw_layout_bbox, draw_span_bbox, draw_line_sort_bbox
@@ -16,10 +16,12 @@ 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"]
 
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
 def read_fn(path):
     if not isinstance(path, Path):
@@ -44,18 +46,10 @@ def prepare_env(output_dir, pdf_file_name, parse_method):
 
 
 def convert_pdf_bytes_to_bytes_by_pypdfium2(pdf_bytes, start_page_id=0, end_page_id=None):
+    pdf = pdfium.PdfDocument(pdf_bytes)
+    output_pdf = pdfium.PdfDocument.new()
     try:
-        # 从字节数据加载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
-
-        # 创建一个新的PDF文档
-        output_pdf = pdfium.PdfDocument.new()
+        end_page_id = get_end_page_id(end_page_id, len(pdf))
 
         # 选择要导入的页面索引
         page_indices = list(range(start_page_id, end_page_id + 1))
@@ -69,13 +63,12 @@ def convert_pdf_bytes_to_bytes_by_pypdfium2(pdf_bytes, start_page_id=0, end_page
 
         # 获取字节数据
         output_bytes = output_buffer.getvalue()
-
-        pdf.close()  # 关闭原PDF文档以释放资源
-        output_pdf.close()  # 关闭新PDF文档以释放资源
     except Exception as e:
         logger.warning(f"Error in converting PDF bytes: {e}, Using original PDF bytes.")
         output_bytes = pdf_bytes
 
+    pdf.close()
+    output_pdf.close()
     return output_bytes
 
 

+ 1 - 1
mineru/cli/gradio_app.py

@@ -13,7 +13,7 @@ from gradio_pdf import PDF
 from loguru import logger
 
 from mineru.cli.common import prepare_env, read_fn, aio_do_parse, pdf_suffixes, image_suffixes
-from mineru.utils.check_mac_env import is_mac_os_version_supported
+from mineru.utils.check_sys_env import is_mac_os_version_supported
 from mineru.utils.cli_parser import arg_parse
 from mineru.utils.hash_utils import str_sha256
 

+ 4 - 0
mineru/utils/check_mac_env.py → mineru/utils/check_sys_env.py

@@ -4,6 +4,10 @@ import platform
 from packaging import version
 
 
+def is_windows_environment() -> bool:
+    return platform.system() == "Windows"
+
+
 # Detect if the current environment is a Mac computer
 def is_mac_environment() -> bool:
     return platform.system() == "Darwin"

+ 113 - 14
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
@@ -7,9 +8,13 @@ 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 .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
 
 
 def pdf_page_to_image(page: pdfium.PdfPage, dpi=200, image_type=ImageType.PIL) -> dict:
@@ -35,7 +40,102 @@ def pdf_page_to_image(page: pdfium.PdfPage, dpi=200, image_type=ImageType.PIL) -
     return image_dict
 
 
+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)
+
+
 def load_images_from_pdf(
+        pdf_bytes: bytes,
+        dpi=200,
+        start_page_id=0,
+        end_page_id=None,
+        image_type=ImageType.PIL,
+        timeout=300,
+        threads=4,
+):
+    """带超时控制的 PDF 转图片函数,支持多进程加速
+
+    Args:
+        timeout (int): 超时时间(秒),默认 300 秒
+        threads (int): 进程数,默认 4
+
+    Raises:
+        TimeoutError: 当转换超时时抛出
+    """
+    pdf_doc = pdfium.PdfDocument(pdf_bytes)
+    if is_windows_environment():
+        # Windows 环境下不使用多进程
+        return load_images_from_pdf_core(
+            pdf_bytes,
+            dpi,
+            start_page_id,
+            get_end_page_id(end_page_id, len(pdf_doc)),
+            image_type
+        ), pdf_doc
+    else:
+        # 根据进程数调整超时时间
+        threads = min(os.cpu_count() or 1, 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(
     pdf_bytes: bytes,
     dpi=200,
     start_page_id=0,
@@ -45,18 +145,17 @@ def load_images_from_pdf(
     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
-
-    for index in range(0, pdf_page_num):
-        if start_page_id <= index <= end_page_id:
-            page = pdf_doc[index]
-            image_dict = pdf_page_to_image(page, dpi=dpi, image_type=image_type)
-            images_list.append(image_dict)
-
-    return images_list, pdf_doc
+    end_page_id = get_end_page_id(end_page_id, pdf_page_num)
+
+    for index in range(start_page_id, end_page_id + 1):
+        # logger.debug(f"Converting page {index}/{pdf_page_num} to image")
+        page = pdf_doc[index]
+        image_dict = pdf_page_to_image(page, dpi=dpi, image_type=image_type)
+        images_list.append(image_dict)
+
+    pdf_doc.close()
+
+    return images_list
 
 
 def cut_image(bbox: tuple, page_num: int, page_pil_img, return_path, image_writer: FileBasedDataWriter, scale=2):

+ 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