pdf_image_tools.py 7.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230
  1. # Copyright (c) Opendatalab. All rights reserved.
  2. import os
  3. from io import BytesIO
  4. import numpy as np
  5. import pypdfium2 as pdfium
  6. from loguru import logger
  7. from PIL import Image
  8. from mineru.data.data_reader_writer import FileBasedDataWriter
  9. from mineru.utils.check_sys_env import is_windows_environment
  10. from mineru.utils.os_env_config import get_load_images_timeout
  11. from mineru.utils.pdf_reader import image_to_b64str, image_to_bytes, page_to_image
  12. from mineru.utils.enum_class import ImageType
  13. from mineru.utils.hash_utils import str_sha256
  14. from mineru.utils.pdf_page_id import get_end_page_id
  15. from concurrent.futures import ProcessPoolExecutor, TimeoutError as FuturesTimeoutError
  16. def pdf_page_to_image(page: pdfium.PdfPage, dpi=200, image_type=ImageType.PIL) -> dict:
  17. """Convert pdfium.PdfDocument to image, Then convert the image to base64.
  18. Args:
  19. page (_type_): pdfium.PdfPage
  20. dpi (int, optional): reset the dpi of dpi. Defaults to 200.
  21. image_type (ImageType, optional): The type of image to return. Defaults to ImageType.PIL.
  22. Returns:
  23. dict: {'img_base64': str, 'img_pil': pil_img, 'scale': float }
  24. """
  25. pil_img, scale = page_to_image(page, dpi=dpi)
  26. image_dict = {
  27. "scale": scale,
  28. }
  29. if image_type == ImageType.BASE64:
  30. image_dict["img_base64"] = image_to_b64str(pil_img)
  31. else:
  32. image_dict["img_pil"] = pil_img
  33. return image_dict
  34. def _load_images_from_pdf_worker(pdf_bytes, dpi, start_page_id, end_page_id, image_type):
  35. """用于进程池的包装函数"""
  36. return load_images_from_pdf_core(pdf_bytes, dpi, start_page_id, end_page_id, image_type)
  37. def load_images_from_pdf(
  38. pdf_bytes: bytes,
  39. dpi=200,
  40. start_page_id=0,
  41. end_page_id=None,
  42. image_type=ImageType.PIL,
  43. timeout=None,
  44. threads=4,
  45. ):
  46. """带超时控制的 PDF 转图片函数,支持多进程加速
  47. Args:
  48. pdf_bytes (bytes): PDF 文件的 bytes
  49. dpi (int, optional): reset the dpi of dpi. Defaults to 200.
  50. start_page_id (int, optional): 起始页码. Defaults to 0.
  51. end_page_id (int | None, optional): 结束页码. Defaults to None.
  52. image_type (ImageType, optional): 图片类型. Defaults to ImageType.PIL.
  53. timeout (int): 超时时间(秒),默认 300 秒
  54. threads (int): 进程数,默认 4
  55. Raises:
  56. TimeoutError: 当转换超时时抛出
  57. """
  58. pdf_doc = pdfium.PdfDocument(pdf_bytes)
  59. if is_windows_environment():
  60. # Windows 环境下不使用多进程
  61. return load_images_from_pdf_core(
  62. pdf_bytes,
  63. dpi,
  64. start_page_id,
  65. get_end_page_id(end_page_id, len(pdf_doc)),
  66. image_type
  67. ), pdf_doc
  68. else:
  69. if timeout is None:
  70. timeout = get_load_images_timeout()
  71. end_page_id = get_end_page_id(end_page_id, len(pdf_doc))
  72. # 计算总页数
  73. total_pages = end_page_id - start_page_id + 1
  74. # 实际使用的进程数不超过总页数
  75. actual_threads = min(min(os.cpu_count() or 1, threads), total_pages)
  76. # 根据实际进程数分组页面范围
  77. pages_per_thread = max(1, total_pages // actual_threads)
  78. page_ranges = []
  79. for i in range(actual_threads):
  80. range_start = start_page_id + i * pages_per_thread
  81. if i == actual_threads - 1:
  82. # 最后一个线程处理剩余所有页面
  83. range_end = end_page_id
  84. else:
  85. range_end = start_page_id + (i + 1) * pages_per_thread - 1
  86. page_ranges.append((range_start, range_end))
  87. # logger.debug(f"PDF to images using {actual_threads} processes, page ranges: {page_ranges}")
  88. with ProcessPoolExecutor(max_workers=actual_threads) as executor:
  89. # 提交所有任务
  90. futures = []
  91. for range_start, range_end in page_ranges:
  92. future = executor.submit(
  93. _load_images_from_pdf_worker,
  94. pdf_bytes,
  95. dpi,
  96. range_start,
  97. range_end,
  98. image_type
  99. )
  100. futures.append((range_start, future))
  101. try:
  102. # 收集结果并按页码排序
  103. all_results = []
  104. for range_start, future in futures:
  105. images_list = future.result(timeout=timeout)
  106. all_results.append((range_start, images_list))
  107. # 按起始页码排序并合并结果
  108. all_results.sort(key=lambda x: x[0])
  109. images_list = []
  110. for _, imgs in all_results:
  111. images_list.extend(imgs)
  112. return images_list, pdf_doc
  113. except FuturesTimeoutError:
  114. pdf_doc.close()
  115. executor.shutdown(wait=False, cancel_futures=True)
  116. raise TimeoutError(f"PDF to images conversion timeout after {timeout}s")
  117. def load_images_from_pdf_core(
  118. pdf_bytes: bytes,
  119. dpi=200,
  120. start_page_id=0,
  121. end_page_id=None,
  122. image_type=ImageType.PIL, # PIL or BASE64
  123. ):
  124. images_list = []
  125. pdf_doc = pdfium.PdfDocument(pdf_bytes)
  126. pdf_page_num = len(pdf_doc)
  127. end_page_id = get_end_page_id(end_page_id, pdf_page_num)
  128. for index in range(start_page_id, end_page_id + 1):
  129. # logger.debug(f"Converting page {index}/{pdf_page_num} to image")
  130. page = pdf_doc[index]
  131. image_dict = pdf_page_to_image(page, dpi=dpi, image_type=image_type)
  132. images_list.append(image_dict)
  133. pdf_doc.close()
  134. return images_list
  135. def cut_image(bbox: tuple, page_num: int, page_pil_img, return_path, image_writer: FileBasedDataWriter, scale=2):
  136. """从第page_num页的page中,根据bbox进行裁剪出一张jpg图片,返回图片路径 save_path:需要同时支持s3和本地,
  137. 图片存放在save_path下,文件名是:
  138. {page_num}_{bbox[0]}_{bbox[1]}_{bbox[2]}_{bbox[3]}.jpg , bbox内数字取整。"""
  139. # 拼接文件名
  140. filename = f"{page_num}_{int(bbox[0])}_{int(bbox[1])}_{int(bbox[2])}_{int(bbox[3])}"
  141. # 老版本返回不带bucket的路径
  142. img_path = f"{return_path}_{filename}" if return_path is not None else None
  143. # 新版本生成平铺路径
  144. img_hash256_path = f"{str_sha256(img_path)}.jpg"
  145. # img_hash256_path = f'{img_path}.jpg'
  146. crop_img = get_crop_img(bbox, page_pil_img, scale=scale)
  147. img_bytes = image_to_bytes(crop_img, image_format="JPEG")
  148. image_writer.write(img_hash256_path, img_bytes)
  149. return img_hash256_path
  150. def get_crop_img(bbox: tuple, pil_img, scale=2):
  151. scale_bbox = (
  152. int(bbox[0] * scale),
  153. int(bbox[1] * scale),
  154. int(bbox[2] * scale),
  155. int(bbox[3] * scale),
  156. )
  157. return pil_img.crop(scale_bbox)
  158. def get_crop_np_img(bbox: tuple, input_img, scale=2):
  159. if isinstance(input_img, Image.Image):
  160. np_img = np.asarray(input_img)
  161. elif isinstance(input_img, np.ndarray):
  162. np_img = input_img
  163. else:
  164. raise ValueError("Input must be a pillow object or a numpy array.")
  165. scale_bbox = (
  166. int(bbox[0] * scale),
  167. int(bbox[1] * scale),
  168. int(bbox[2] * scale),
  169. int(bbox[3] * scale),
  170. )
  171. return np_img[scale_bbox[1]:scale_bbox[3], scale_bbox[0]:scale_bbox[2]]
  172. def images_bytes_to_pdf_bytes(image_bytes):
  173. # 内存缓冲区
  174. pdf_buffer = BytesIO()
  175. # 载入并转换所有图像为 RGB 模式
  176. image = Image.open(BytesIO(image_bytes)).convert("RGB")
  177. # 第一张图保存为 PDF,其余追加
  178. image.save(pdf_buffer, format="PDF", save_all=True)
  179. # 获取 PDF bytes 并重置指针(可选)
  180. pdf_bytes = pdf_buffer.getvalue()
  181. pdf_buffer.close()
  182. return pdf_bytes