| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214 |
- """
- PDF处理工具模块
- 提供PDF相关处理功能:
- - PDF加载与分类
- - PDF文本提取
- - 跨页表格合并
- """
- from typing import Dict, List, Any, Optional, Tuple
- from pathlib import Path
- from PIL import Image
- from loguru import logger
- # 导入 MinerU 组件
- try:
- from mineru.utils.pdf_classify import classify as pdf_classify
- from mineru.utils.pdf_image_tools import load_images_from_pdf
- from mineru.utils.enum_class import ImageType
- from mineru.utils.pdf_text_tool import get_page as pdf_get_page_text
- MINERU_AVAILABLE = True
- except ImportError:
- MINERU_AVAILABLE = False
- pdf_classify = None
- load_images_from_pdf = None
- ImageType = None
- pdf_get_page_text = None
- class PDFUtils:
- """PDF处理工具类"""
-
- @staticmethod
- def load_and_classify_document(
- document_path: Path,
- dpi: int = 200
- ) -> Tuple[List[Dict], str, Optional[Any]]:
- """
- 加载文档并分类
-
- Args:
- document_path: 文档路径
- dpi: PDF渲染DPI
-
- Returns:
- (images_list, pdf_type, pdf_doc)
- - images_list: 图像列表,每个元素包含 {'img_pil': PIL.Image, 'scale': float}
- - pdf_type: 'ocr' 或 'txt'
- - pdf_doc: PDF文档对象(如果是PDF)
- """
- pdf_doc = None
- pdf_type = 'ocr' # 默认使用OCR模式
- images = []
-
- if document_path.is_dir():
- # 处理目录:遍历所有图片
- image_extensions = {'.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.tif'}
- image_files = sorted([
- f for f in document_path.iterdir()
- if f.suffix.lower() in image_extensions
- ])
-
- for img_file in image_files:
- img = Image.open(img_file)
- if img.mode != 'RGB':
- img = img.convert('RGB')
- images.append({
- 'img_pil': img,
- 'scale': 1.0,
- 'source_path': str(img_file)
- })
-
- pdf_type = 'ocr' # 图片目录始终使用OCR模式
-
- elif document_path.suffix.lower() == '.pdf':
- # 处理PDF文件
- if not MINERU_AVAILABLE:
- raise RuntimeError("MinerU components not available for PDF processing")
-
- with open(document_path, 'rb') as f:
- pdf_bytes = f.read()
-
- # PDF分类
- pdf_type = pdf_classify(pdf_bytes)
- logger.info(f"📋 PDF classified as: {pdf_type}")
-
- # 加载图像
- images_list, pdf_doc = load_images_from_pdf(
- pdf_bytes,
- dpi=dpi,
- image_type=ImageType.PIL
- )
-
- for img_dict in images_list:
- images.append({
- 'img_pil': img_dict['img_pil'],
- 'scale': img_dict.get('scale', dpi / 72),
- 'source_path': str(document_path)
- })
-
- elif document_path.suffix.lower() in ['.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.tif']:
- # 处理单个图片
- img = Image.open(document_path)
- if img.mode != 'RGB':
- img = img.convert('RGB')
- images.append({
- 'img_pil': img,
- 'scale': 1.0,
- 'source_path': str(document_path)
- })
- pdf_type = 'ocr'
-
- else:
- raise ValueError(f"Unsupported file format: {document_path.suffix}")
-
- return images, pdf_type, pdf_doc
-
- @staticmethod
- def extract_text_from_pdf(
- pdf_doc: Any,
- page_idx: int,
- bbox: List[float],
- scale: float
- ) -> Tuple[str, bool]:
- """
- 从PDF直接提取文本(使用 MinerU 的 pypdfium2 方式)
-
- Args:
- pdf_doc: pypdfium2 的 PdfDocument 对象
- page_idx: 页码索引
- bbox: 目标区域的bbox(图像坐标)
- scale: 图像与PDF的缩放比例
-
- Returns:
- (text, success)
- """
- if not MINERU_AVAILABLE or pdf_get_page_text is None:
- logger.debug("MinerU pdf_text_tool not available")
- return "", False
-
- try:
- page = pdf_doc[page_idx]
-
- # 将图像坐标转换为PDF坐标
- pdf_bbox = [
- bbox[0] / scale,
- bbox[1] / scale,
- bbox[2] / scale,
- bbox[3] / scale
- ]
-
- # 使用 MinerU 的方式获取页面文本信息
- page_dict = pdf_get_page_text(page)
-
- # 从 blocks 中提取与 bbox 重叠的文本
- text_parts = []
- for block in page_dict.get('blocks', []):
- for line in block.get('lines', []):
- line_bbox = line.get('bbox')
- if line_bbox and hasattr(line_bbox, 'bbox'):
- line_bbox = line_bbox.bbox # pdftext 的 BBox 对象
- elif isinstance(line_bbox, (list, tuple)) and len(line_bbox) >= 4:
- line_bbox = list(line_bbox)
- else:
- continue
-
- # 检查 line 是否与目标 bbox 重叠
- if PDFUtils._bbox_overlap(pdf_bbox, line_bbox):
- for span in line.get('spans', []):
- span_text = span.get('text', '')
- if span_text:
- text_parts.append(span_text)
-
- text = ' '.join(text_parts)
- return text.strip(), bool(text.strip())
-
- except Exception as e:
- import traceback
- logger.debug(f"PDF text extraction error: {e}")
- logger.debug(traceback.format_exc())
- return "", False
-
- @staticmethod
- def _bbox_overlap(bbox1: List[float], bbox2: List[float]) -> bool:
- """检查两个 bbox 是否重叠"""
- if len(bbox1) < 4 or len(bbox2) < 4:
- return False
-
- x1_1, y1_1, x2_1, y2_1 = bbox1[:4]
- x1_2, y1_2, x2_2, y2_2 = bbox2[:4]
-
- if x2_1 < x1_2 or x2_2 < x1_1:
- return False
- if y2_1 < y1_2 or y2_2 < y1_1:
- return False
-
- return True
-
- @staticmethod
- def merge_cross_page_tables(results: Dict[str, Any]) -> Dict[str, Any]:
- """
- 合并跨页表格
-
- TODO: 实现跨页表格合并逻辑
- 可以参考 MinerU 的 cross_page_table_merge 实现
-
- Args:
- results: 处理结果字典
-
- Returns:
- 合并后的结果
- """
- # TODO: 实现跨页表格合并逻辑
- return results
|