pdf_utils.py 6.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214
  1. """
  2. PDF处理工具模块
  3. 提供PDF相关处理功能:
  4. - PDF加载与分类
  5. - PDF文本提取
  6. - 跨页表格合并
  7. """
  8. from typing import Dict, List, Any, Optional, Tuple
  9. from pathlib import Path
  10. from PIL import Image
  11. from loguru import logger
  12. # 导入 MinerU 组件
  13. try:
  14. from mineru.utils.pdf_classify import classify as pdf_classify
  15. from mineru.utils.pdf_image_tools import load_images_from_pdf
  16. from mineru.utils.enum_class import ImageType
  17. from mineru.utils.pdf_text_tool import get_page as pdf_get_page_text
  18. MINERU_AVAILABLE = True
  19. except ImportError:
  20. MINERU_AVAILABLE = False
  21. pdf_classify = None
  22. load_images_from_pdf = None
  23. ImageType = None
  24. pdf_get_page_text = None
  25. class PDFUtils:
  26. """PDF处理工具类"""
  27. @staticmethod
  28. def load_and_classify_document(
  29. document_path: Path,
  30. dpi: int = 200
  31. ) -> Tuple[List[Dict], str, Optional[Any]]:
  32. """
  33. 加载文档并分类
  34. Args:
  35. document_path: 文档路径
  36. dpi: PDF渲染DPI
  37. Returns:
  38. (images_list, pdf_type, pdf_doc)
  39. - images_list: 图像列表,每个元素包含 {'img_pil': PIL.Image, 'scale': float}
  40. - pdf_type: 'ocr' 或 'txt'
  41. - pdf_doc: PDF文档对象(如果是PDF)
  42. """
  43. pdf_doc = None
  44. pdf_type = 'ocr' # 默认使用OCR模式
  45. images = []
  46. if document_path.is_dir():
  47. # 处理目录:遍历所有图片
  48. image_extensions = {'.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.tif'}
  49. image_files = sorted([
  50. f for f in document_path.iterdir()
  51. if f.suffix.lower() in image_extensions
  52. ])
  53. for img_file in image_files:
  54. img = Image.open(img_file)
  55. if img.mode != 'RGB':
  56. img = img.convert('RGB')
  57. images.append({
  58. 'img_pil': img,
  59. 'scale': 1.0,
  60. 'source_path': str(img_file)
  61. })
  62. pdf_type = 'ocr' # 图片目录始终使用OCR模式
  63. elif document_path.suffix.lower() == '.pdf':
  64. # 处理PDF文件
  65. if not MINERU_AVAILABLE:
  66. raise RuntimeError("MinerU components not available for PDF processing")
  67. with open(document_path, 'rb') as f:
  68. pdf_bytes = f.read()
  69. # PDF分类
  70. pdf_type = pdf_classify(pdf_bytes)
  71. logger.info(f"📋 PDF classified as: {pdf_type}")
  72. # 加载图像
  73. images_list, pdf_doc = load_images_from_pdf(
  74. pdf_bytes,
  75. dpi=dpi,
  76. image_type=ImageType.PIL
  77. )
  78. for img_dict in images_list:
  79. images.append({
  80. 'img_pil': img_dict['img_pil'],
  81. 'scale': img_dict.get('scale', dpi / 72),
  82. 'source_path': str(document_path)
  83. })
  84. elif document_path.suffix.lower() in ['.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.tif']:
  85. # 处理单个图片
  86. img = Image.open(document_path)
  87. if img.mode != 'RGB':
  88. img = img.convert('RGB')
  89. images.append({
  90. 'img_pil': img,
  91. 'scale': 1.0,
  92. 'source_path': str(document_path)
  93. })
  94. pdf_type = 'ocr'
  95. else:
  96. raise ValueError(f"Unsupported file format: {document_path.suffix}")
  97. return images, pdf_type, pdf_doc
  98. @staticmethod
  99. def extract_text_from_pdf(
  100. pdf_doc: Any,
  101. page_idx: int,
  102. bbox: List[float],
  103. scale: float
  104. ) -> Tuple[str, bool]:
  105. """
  106. 从PDF直接提取文本(使用 MinerU 的 pypdfium2 方式)
  107. Args:
  108. pdf_doc: pypdfium2 的 PdfDocument 对象
  109. page_idx: 页码索引
  110. bbox: 目标区域的bbox(图像坐标)
  111. scale: 图像与PDF的缩放比例
  112. Returns:
  113. (text, success)
  114. """
  115. if not MINERU_AVAILABLE or pdf_get_page_text is None:
  116. logger.debug("MinerU pdf_text_tool not available")
  117. return "", False
  118. try:
  119. page = pdf_doc[page_idx]
  120. # 将图像坐标转换为PDF坐标
  121. pdf_bbox = [
  122. bbox[0] / scale,
  123. bbox[1] / scale,
  124. bbox[2] / scale,
  125. bbox[3] / scale
  126. ]
  127. # 使用 MinerU 的方式获取页面文本信息
  128. page_dict = pdf_get_page_text(page)
  129. # 从 blocks 中提取与 bbox 重叠的文本
  130. text_parts = []
  131. for block in page_dict.get('blocks', []):
  132. for line in block.get('lines', []):
  133. line_bbox = line.get('bbox')
  134. if line_bbox and hasattr(line_bbox, 'bbox'):
  135. line_bbox = line_bbox.bbox # pdftext 的 BBox 对象
  136. elif isinstance(line_bbox, (list, tuple)) and len(line_bbox) >= 4:
  137. line_bbox = list(line_bbox)
  138. else:
  139. continue
  140. # 检查 line 是否与目标 bbox 重叠
  141. if PDFUtils._bbox_overlap(pdf_bbox, line_bbox):
  142. for span in line.get('spans', []):
  143. span_text = span.get('text', '')
  144. if span_text:
  145. text_parts.append(span_text)
  146. text = ' '.join(text_parts)
  147. return text.strip(), bool(text.strip())
  148. except Exception as e:
  149. import traceback
  150. logger.debug(f"PDF text extraction error: {e}")
  151. logger.debug(traceback.format_exc())
  152. return "", False
  153. @staticmethod
  154. def _bbox_overlap(bbox1: List[float], bbox2: List[float]) -> bool:
  155. """检查两个 bbox 是否重叠"""
  156. if len(bbox1) < 4 or len(bbox2) < 4:
  157. return False
  158. x1_1, y1_1, x2_1, y2_1 = bbox1[:4]
  159. x1_2, y1_2, x2_2, y2_2 = bbox2[:4]
  160. if x2_1 < x1_2 or x2_2 < x1_1:
  161. return False
  162. if y2_1 < y1_2 or y2_2 < y1_1:
  163. return False
  164. return True
  165. @staticmethod
  166. def merge_cross_page_tables(results: Dict[str, Any]) -> Dict[str, Any]:
  167. """
  168. 合并跨页表格
  169. TODO: 实现跨页表格合并逻辑
  170. 可以参考 MinerU 的 cross_page_table_merge 实现
  171. Args:
  172. results: 处理结果字典
  173. Returns:
  174. 合并后的结果
  175. """
  176. # TODO: 实现跨页表格合并逻辑
  177. return results