pdf_image_tools.py 7.3 KB

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