Browse Source

refactor: move PDF byte conversion functions to pdf_page_tools and simplify logic

myhloli 3 weeks ago
parent
commit
ace7f76869

+ 2 - 1
demo/demo.py

@@ -6,7 +6,8 @@ from pathlib import Path
 
 from loguru import logger
 
-from mineru.cli.common import convert_pdf_bytes_to_bytes_by_pypdfium2, prepare_env, read_fn
+from mineru.cli.common import prepare_env, read_fn
+from mineru.utils.pdf_page_tools import convert_pdf_bytes_to_bytes_by_pypdfium2
 from mineru.data.data_reader_writer import FileBasedDataWriter
 from mineru.utils.draw_bbox import draw_layout_bbox, draw_span_bbox
 from mineru.utils.enum_class import MakeMode

+ 1 - 33
mineru/cli/common.py

@@ -1,17 +1,12 @@
 # Copyright (c) Opendatalab. All rights reserved.
-import io
 import json
 import os
 import copy
-import time
-from multiprocessing import Pool
 from pathlib import Path
 
-import pypdfium2 as pdfium
 from loguru import logger
 
 from mineru.data.data_reader_writer import FileBasedDataWriter
-from mineru.utils.check_sys_env import is_windows_environment
 from mineru.utils.draw_bbox import draw_layout_bbox, draw_span_bbox, draw_line_sort_bbox
 from mineru.utils.enum_class import MakeMode
 from mineru.utils.guess_suffix_or_lang import guess_suffix_by_bytes
@@ -19,7 +14,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
+from mineru.utils.pdf_page_tools import convert_pdf_bytes_to_bytes_by_pypdfium2
 
 pdf_suffixes = ["pdf"]
 image_suffixes = ["png", "jpeg", "jp2", "webp", "gif", "bmp", "jpg", "tiff"]
@@ -48,33 +43,6 @@ def prepare_env(output_dir, pdf_file_name, parse_method):
     return local_image_dir, local_md_dir
 
 
-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:
-        end_page_id = get_end_page_id(end_page_id, len(pdf))
-
-        # 选择要导入的页面索引
-        page_indices = list(range(start_page_id, end_page_id + 1))
-
-        # 从原PDF导入页面到新PDF
-        output_pdf.import_pages(pdf, page_indices)
-
-        # 将新PDF保存到内存缓冲区
-        output_buffer = io.BytesIO()
-        output_pdf.save(output_buffer)
-
-        # 获取字节数据
-        output_bytes = output_buffer.getvalue()
-    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
-
-
 def _prepare_pdf_bytes(pdf_bytes_list, start_page_id, end_page_id):
     """准备处理PDF字节数据"""
     result = []

+ 9 - 7
mineru/utils/pdf_image_tools.py

@@ -11,7 +11,7 @@ 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
-from mineru.utils.pdf_page_id import get_end_page_id
+from mineru.utils.pdf_page_tools import get_end_page_id, convert_pdf_bytes_to_bytes_by_pypdfium2
 
 from concurrent.futures import ProcessPoolExecutor, TimeoutError as FuturesTimeoutError
 
@@ -69,7 +69,7 @@ def load_images_from_pdf(
             pdf_bytes,
             dpi,
             start_page_id,
-            end_page_id,
+            get_end_page_id(end_page_id, len(pdf_doc)),
             image_type
         ), pdf_doc
     else:
@@ -143,13 +143,15 @@ def load_images_from_pdf_core(
     image_type=ImageType.PIL,  # PIL or BASE64
 ):
     images_list = []
+    pdf_bytes = convert_pdf_bytes_to_bytes_by_pypdfium2(pdf_bytes, start_page_id, end_page_id)
     pdf_doc = pdfium.PdfDocument(pdf_bytes)
-    pdf_page_num = len(pdf_doc)
-    end_page_id = get_end_page_id(end_page_id, pdf_page_num)
+    # pdf_page_num = len(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]
+    # 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]
+    for page in pdf_doc:
         image_dict = pdf_page_to_image(page, dpi=dpi, image_type=image_type)
         images_list.append(image_dict)
 

+ 0 - 10
mineru/utils/pdf_page_id.py

@@ -1,10 +0,0 @@
-# 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

+ 40 - 0
mineru/utils/pdf_page_tools.py

@@ -0,0 +1,40 @@
+# Copyright (c) Opendatalab. All rights reserved.
+import io
+
+import pypdfium2 as pdfium
+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
+
+
+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:
+        end_page_id = get_end_page_id(end_page_id, len(pdf))
+
+        # 选择要导入的页面索引
+        page_indices = list(range(start_page_id, end_page_id + 1))
+
+        # 从原PDF导入页面到新PDF
+        output_pdf.import_pages(pdf, page_indices)
+
+        # 将新PDF保存到内存缓冲区
+        output_buffer = io.BytesIO()
+        output_pdf.save(output_buffer)
+
+        # 获取字节数据
+        output_bytes = output_buffer.getvalue()
+    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
tests/unittest/test_e2e.py

@@ -7,10 +7,10 @@ from loguru import logger
 from bs4 import BeautifulSoup
 from fuzzywuzzy import fuzz
 from mineru.cli.common import (
-    convert_pdf_bytes_to_bytes_by_pypdfium2,
     prepare_env,
     read_fn,
 )
+from mineru.utils.pdf_page_tools import convert_pdf_bytes_to_bytes_by_pypdfium2
 from mineru.data.data_reader_writer import FileBasedDataWriter
 from mineru.utils.enum_class import MakeMode
 from mineru.backend.vlm.vlm_analyze import doc_analyze as vlm_doc_analyze