pdf_image_tools.py 5.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168
  1. # Copyright (c) Opendatalab. All rights reserved.
  2. from io import BytesIO
  3. import numpy as np
  4. import pypdfium2 as pdfium
  5. from loguru import logger
  6. from PIL import Image
  7. from mineru.data.data_reader_writer import FileBasedDataWriter
  8. from mineru.utils.pdf_reader import image_to_b64str, image_to_bytes, page_to_image
  9. from .enum_class import ImageType
  10. from .hash_utils import str_sha256
  11. from concurrent.futures import ProcessPoolExecutor, TimeoutError as FuturesTimeoutError
  12. def pdf_page_to_image(page: pdfium.PdfPage, dpi=200, image_type=ImageType.PIL) -> dict:
  13. """Convert pdfium.PdfDocument to image, Then convert the image to base64.
  14. Args:
  15. page (_type_): pdfium.PdfPage
  16. dpi (int, optional): reset the dpi of dpi. Defaults to 200.
  17. image_type (ImageType, optional): The type of image to return. Defaults to ImageType.PIL.
  18. Returns:
  19. dict: {'img_base64': str, 'img_pil': pil_img, 'scale': float }
  20. """
  21. pil_img, scale = page_to_image(page, dpi=dpi)
  22. image_dict = {
  23. "scale": scale,
  24. }
  25. if image_type == ImageType.BASE64:
  26. image_dict["img_base64"] = image_to_b64str(pil_img)
  27. else:
  28. image_dict["img_pil"] = pil_img
  29. return image_dict
  30. #@todo 通过多进程渲染pdf提速
  31. def _load_images_from_pdf_worker(pdf_bytes, dpi, start_page_id, end_page_id, image_type):
  32. """用于进程池的包装函数"""
  33. return load_images_from_pdf_core(pdf_bytes, dpi, start_page_id, end_page_id, image_type)
  34. def load_images_from_pdf(
  35. pdf_bytes: bytes,
  36. dpi=200,
  37. start_page_id=0,
  38. end_page_id=None,
  39. image_type=ImageType.PIL,
  40. timeout=300,
  41. ):
  42. """带超时控制的 PDF 转图片函数
  43. Args:
  44. timeout (int): 超时时间(秒),默认 300 秒
  45. Raises:
  46. TimeoutError: 当转换超时时抛出
  47. """
  48. with ProcessPoolExecutor(max_workers=1) as executor:
  49. future = executor.submit(
  50. _load_images_from_pdf_worker,
  51. pdf_bytes,
  52. dpi,
  53. start_page_id,
  54. end_page_id,
  55. image_type
  56. )
  57. try:
  58. images_list = future.result(timeout=timeout)
  59. pdf_doc = pdfium.PdfDocument(pdf_bytes)
  60. return images_list, pdf_doc
  61. except FuturesTimeoutError:
  62. logger.error(f"PDF conversion timeout after {timeout}s")
  63. executor.shutdown(wait=False, cancel_futures=True)
  64. raise TimeoutError(f"PDF to images conversion timeout after {timeout}s")
  65. def load_images_from_pdf_core(
  66. pdf_bytes: bytes,
  67. dpi=200,
  68. start_page_id=0,
  69. end_page_id=None,
  70. image_type=ImageType.PIL, # PIL or BASE64
  71. ):
  72. images_list = []
  73. pdf_doc = pdfium.PdfDocument(pdf_bytes)
  74. pdf_page_num = len(pdf_doc)
  75. end_page_id = end_page_id if end_page_id is not None and end_page_id >= 0 else pdf_page_num - 1
  76. if end_page_id > pdf_page_num - 1:
  77. logger.warning("end_page_id is out of range, use images length")
  78. end_page_id = pdf_page_num - 1
  79. for index in range(0, pdf_page_num):
  80. if start_page_id <= index <= end_page_id:
  81. page = pdf_doc[index]
  82. image_dict = pdf_page_to_image(page, dpi=dpi, image_type=image_type)
  83. images_list.append(image_dict)
  84. return images_list
  85. def cut_image(bbox: tuple, page_num: int, page_pil_img, return_path, image_writer: FileBasedDataWriter, scale=2):
  86. """从第page_num页的page中,根据bbox进行裁剪出一张jpg图片,返回图片路径 save_path:需要同时支持s3和本地,
  87. 图片存放在save_path下,文件名是:
  88. {page_num}_{bbox[0]}_{bbox[1]}_{bbox[2]}_{bbox[3]}.jpg , bbox内数字取整。"""
  89. # 拼接文件名
  90. filename = f"{page_num}_{int(bbox[0])}_{int(bbox[1])}_{int(bbox[2])}_{int(bbox[3])}"
  91. # 老版本返回不带bucket的路径
  92. img_path = f"{return_path}_{filename}" if return_path is not None else None
  93. # 新版本生成平铺路径
  94. img_hash256_path = f"{str_sha256(img_path)}.jpg"
  95. # img_hash256_path = f'{img_path}.jpg'
  96. crop_img = get_crop_img(bbox, page_pil_img, scale=scale)
  97. img_bytes = image_to_bytes(crop_img, image_format="JPEG")
  98. image_writer.write(img_hash256_path, img_bytes)
  99. return img_hash256_path
  100. def get_crop_img(bbox: tuple, pil_img, scale=2):
  101. scale_bbox = (
  102. int(bbox[0] * scale),
  103. int(bbox[1] * scale),
  104. int(bbox[2] * scale),
  105. int(bbox[3] * scale),
  106. )
  107. return pil_img.crop(scale_bbox)
  108. def get_crop_np_img(bbox: tuple, input_img, scale=2):
  109. if isinstance(input_img, Image.Image):
  110. np_img = np.asarray(input_img)
  111. elif isinstance(input_img, np.ndarray):
  112. np_img = input_img
  113. else:
  114. raise ValueError("Input must be a pillow object or a numpy array.")
  115. scale_bbox = (
  116. int(bbox[0] * scale),
  117. int(bbox[1] * scale),
  118. int(bbox[2] * scale),
  119. int(bbox[3] * scale),
  120. )
  121. return np_img[scale_bbox[1]:scale_bbox[3], scale_bbox[0]:scale_bbox[2]]
  122. def images_bytes_to_pdf_bytes(image_bytes):
  123. # 内存缓冲区
  124. pdf_buffer = BytesIO()
  125. # 载入并转换所有图像为 RGB 模式
  126. image = Image.open(BytesIO(image_bytes)).convert("RGB")
  127. # 第一张图保存为 PDF,其余追加
  128. image.save(pdf_buffer, format="PDF", save_all=True)
  129. # 获取 PDF bytes 并重置指针(可选)
  130. pdf_bytes = pdf_buffer.getvalue()
  131. pdf_buffer.close()
  132. return pdf_bytes