pdf_utils.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299
  1. """
  2. PDF处理工具模块
  3. 提供PDF相关处理功能:
  4. - PDF加载与分类
  5. - PDF文本提取
  6. - 跨页表格合并
  7. - 页面范围解析与过滤
  8. """
  9. from typing import Dict, List, Any, Optional, Tuple, Set
  10. from pathlib import Path
  11. from PIL import Image
  12. from loguru import logger
  13. import re
  14. # 导入 MinerU 组件
  15. try:
  16. from mineru.utils.pdf_classify import classify as pdf_classify
  17. from mineru.utils.pdf_image_tools import load_images_from_pdf
  18. from mineru.utils.enum_class import ImageType
  19. from mineru.utils.pdf_text_tool import get_page as pdf_get_page_text
  20. MINERU_AVAILABLE = True
  21. except ImportError:
  22. MINERU_AVAILABLE = False
  23. pdf_classify = None
  24. load_images_from_pdf = None
  25. ImageType = None
  26. pdf_get_page_text = None
  27. class PDFUtils:
  28. """PDF处理工具类"""
  29. @staticmethod
  30. def parse_page_range(page_range: Optional[str], total_pages: int) -> Set[int]:
  31. """
  32. 解析页面范围字符串
  33. 支持格式:
  34. - "1-5" → {0, 1, 2, 3, 4}(页码从1开始,内部转为0-based索引)
  35. - "3" → {2}
  36. - "1-5,7,9-12" → {0, 1, 2, 3, 4, 6, 8, 9, 10, 11}
  37. - "1-" → 从第1页到最后
  38. - "-5" → 从第1页到第5页
  39. Args:
  40. page_range: 页面范围字符串(页码从1开始)
  41. total_pages: 总页数
  42. Returns:
  43. 页面索引集合(0-based)
  44. """
  45. if not page_range or not page_range.strip():
  46. return set(range(total_pages))
  47. pages = set()
  48. parts = page_range.replace(' ', '').split(',')
  49. for part in parts:
  50. part = part.strip()
  51. if not part:
  52. continue
  53. if '-' in part:
  54. # 范围格式
  55. match = re.match(r'^(\d*)-(\d*)$', part)
  56. if match:
  57. start_str, end_str = match.groups()
  58. start = int(start_str) if start_str else 1
  59. end = int(end_str) if end_str else total_pages
  60. # 转换为 0-based 索引
  61. start = max(0, start - 1)
  62. end = min(total_pages, end)
  63. pages.update(range(start, end))
  64. else:
  65. # 单页
  66. try:
  67. page_num = int(part)
  68. if 1 <= page_num <= total_pages:
  69. pages.add(page_num - 1) # 转换为 0-based 索引
  70. except ValueError:
  71. logger.warning(f"Invalid page number: {part}")
  72. return pages
  73. @staticmethod
  74. def load_and_classify_document(
  75. document_path: Path,
  76. dpi: int = 200,
  77. page_range: Optional[str] = None
  78. ) -> Tuple[List[Dict], str, Optional[Any]]:
  79. """
  80. 加载文档并分类,支持页面范围过滤
  81. Args:
  82. document_path: 文档路径
  83. dpi: PDF渲染DPI
  84. page_range: 页面范围字符串,如 "1-5,7,9-12"
  85. - PDF:按页码(从1开始)
  86. - 图片目录:按文件名排序后的位置(从1开始)
  87. Returns:
  88. (images_list, pdf_type, pdf_doc)
  89. - images_list: 图像列表,每个元素包含 {'img_pil': PIL.Image, 'scale': float, 'page_idx': int}
  90. - pdf_type: 'ocr' 或 'txt'
  91. - pdf_doc: PDF文档对象(如果是PDF)
  92. """
  93. pdf_doc = None
  94. pdf_type = 'ocr' # 默认使用OCR模式
  95. all_images = []
  96. if document_path.is_dir():
  97. # 处理目录:遍历所有图片
  98. image_extensions = {'.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.tif'}
  99. image_files = sorted([
  100. f for f in document_path.iterdir()
  101. if f.suffix.lower() in image_extensions
  102. ])
  103. # 解析页面范围
  104. total_pages = len(image_files)
  105. selected_pages = PDFUtils.parse_page_range(page_range, total_pages)
  106. if page_range:
  107. logger.info(f"📋 图片目录共 {total_pages} 张,选择处理 {len(selected_pages)} 张")
  108. for idx, img_file in enumerate(image_files):
  109. if idx not in selected_pages:
  110. continue
  111. img = Image.open(img_file)
  112. if img.mode != 'RGB':
  113. img = img.convert('RGB')
  114. all_images.append({
  115. 'img_pil': img,
  116. 'scale': 1.0,
  117. 'source_path': str(img_file),
  118. 'page_idx': idx, # 原始索引
  119. 'page_name': img_file.stem # 文件名(不含扩展名)
  120. })
  121. pdf_type = 'ocr' # 图片目录始终使用OCR模式
  122. elif document_path.suffix.lower() == '.pdf':
  123. # 处理PDF文件
  124. if not MINERU_AVAILABLE:
  125. raise RuntimeError("MinerU components not available for PDF processing")
  126. with open(document_path, 'rb') as f:
  127. pdf_bytes = f.read()
  128. # PDF分类
  129. pdf_type = pdf_classify(pdf_bytes)
  130. logger.info(f"📋 PDF classified as: {pdf_type}")
  131. # 加载图像
  132. images_list, pdf_doc = load_images_from_pdf(
  133. pdf_bytes,
  134. dpi=dpi,
  135. image_type=ImageType.PIL
  136. )
  137. # 解析页面范围
  138. total_pages = len(images_list)
  139. selected_pages = PDFUtils.parse_page_range(page_range, total_pages)
  140. if page_range:
  141. logger.info(f"📋 PDF 共 {total_pages} 页,选择处理 {len(selected_pages)} 页")
  142. for idx, img_dict in enumerate(images_list):
  143. if idx not in selected_pages:
  144. continue
  145. all_images.append({
  146. 'img_pil': img_dict['img_pil'],
  147. 'scale': img_dict.get('scale', dpi / 72),
  148. 'source_path': str(document_path),
  149. 'page_idx': idx # 原始页码索引
  150. })
  151. elif document_path.suffix.lower() in ['.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.tif']:
  152. # 处理单个图片
  153. img = Image.open(document_path)
  154. if img.mode != 'RGB':
  155. img = img.convert('RGB')
  156. all_images.append({
  157. 'img_pil': img,
  158. 'scale': 1.0,
  159. 'source_path': str(document_path),
  160. 'page_idx': 0,
  161. 'page_name': document_path.stem
  162. })
  163. pdf_type = 'ocr'
  164. else:
  165. raise ValueError(f"Unsupported file format: {document_path.suffix}")
  166. return all_images, pdf_type, pdf_doc
  167. @staticmethod
  168. def extract_text_from_pdf(
  169. pdf_doc: Any,
  170. page_idx: int,
  171. bbox: List[float],
  172. scale: float
  173. ) -> Tuple[str, bool]:
  174. """
  175. 从PDF直接提取文本(使用 MinerU 的 pypdfium2 方式)
  176. Args:
  177. pdf_doc: pypdfium2 的 PdfDocument 对象
  178. page_idx: 页码索引
  179. bbox: 目标区域的bbox(图像坐标)
  180. scale: 图像与PDF的缩放比例
  181. Returns:
  182. (text, success)
  183. """
  184. if not MINERU_AVAILABLE or pdf_get_page_text is None:
  185. logger.debug("MinerU pdf_text_tool not available")
  186. return "", False
  187. try:
  188. page = pdf_doc[page_idx]
  189. # 将图像坐标转换为PDF坐标
  190. pdf_bbox = [
  191. bbox[0] / scale,
  192. bbox[1] / scale,
  193. bbox[2] / scale,
  194. bbox[3] / scale
  195. ]
  196. # 使用 MinerU 的方式获取页面文本信息
  197. page_dict = pdf_get_page_text(page)
  198. # 从 blocks 中提取与 bbox 重叠的文本
  199. text_parts = []
  200. for block in page_dict.get('blocks', []):
  201. for line in block.get('lines', []):
  202. line_bbox = line.get('bbox')
  203. if line_bbox and hasattr(line_bbox, 'bbox'):
  204. line_bbox = line_bbox.bbox # pdftext 的 BBox 对象
  205. elif isinstance(line_bbox, (list, tuple)) and len(line_bbox) >= 4:
  206. line_bbox = list(line_bbox)
  207. else:
  208. continue
  209. # 检查 line 是否与目标 bbox 重叠
  210. if PDFUtils._bbox_overlap(pdf_bbox, line_bbox):
  211. for span in line.get('spans', []):
  212. span_text = span.get('text', '')
  213. if span_text:
  214. text_parts.append(span_text)
  215. text = ' '.join(text_parts)
  216. return text.strip(), bool(text.strip())
  217. except Exception as e:
  218. import traceback
  219. logger.debug(f"PDF text extraction error: {e}")
  220. logger.debug(traceback.format_exc())
  221. return "", False
  222. @staticmethod
  223. def _bbox_overlap(bbox1: List[float], bbox2: List[float]) -> bool:
  224. """检查两个 bbox 是否重叠"""
  225. if len(bbox1) < 4 or len(bbox2) < 4:
  226. return False
  227. x1_1, y1_1, x2_1, y2_1 = bbox1[:4]
  228. x1_2, y1_2, x2_2, y2_2 = bbox2[:4]
  229. if x2_1 < x1_2 or x2_2 < x1_1:
  230. return False
  231. if y2_1 < y1_2 or y2_2 < y1_1:
  232. return False
  233. return True
  234. @staticmethod
  235. def merge_cross_page_tables(results: Dict[str, Any]) -> Dict[str, Any]:
  236. """
  237. 合并跨页表格
  238. TODO: 实现跨页表格合并逻辑
  239. 可以参考 MinerU 的 cross_page_table_merge 实现
  240. Args:
  241. results: 处理结果字典
  242. Returns:
  243. 合并后的结果
  244. """
  245. # TODO: 实现跨页表格合并逻辑
  246. return results