""" PDF处理工具模块 提供PDF相关处理功能: - PDF加载与分类 - PDF文本提取 - 跨页表格合并 - 页面范围解析与过滤 """ from typing import Dict, List, Any, Optional, Tuple, Set from pathlib import Path from PIL import Image from loguru import logger import re # 导入页面范围解析函数(不依赖 MinerU) from .file_utils import parse_page_range # 导入 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: raise ImportError("MinerU components not available for PDF processing") class PDFUtils: """PDF处理工具类""" @staticmethod def parse_page_range(page_range: Optional[str], total_pages: int) -> Set[int]: """ 解析页面范围字符串(向后兼容包装函数) 此方法是对 file_utils.parse_page_range 的包装,保持向后兼容性。 新代码应直接使用 file_utils.parse_page_range。 支持格式: - "1-5" → {0, 1, 2, 3, 4}(页码从1开始,内部转为0-based索引) - "3" → {2} - "1-5,7,9-12" → {0, 1, 2, 3, 4, 6, 8, 9, 10, 11} - "1-" → 从第1页到最后 - "-5" → 从第1页到第5页 Args: page_range: 页面范围字符串(页码从1开始) total_pages: 总页数 Returns: 页面索引集合(0-based) """ return parse_page_range(page_range, total_pages) @staticmethod def load_and_classify_document( document_path: Path, dpi: int = 200, page_range: Optional[str] = None ) -> Tuple[List[Dict], str, Optional[Any]]: """ 加载文档并分类,支持页面范围过滤 Args: document_path: 文档路径 dpi: PDF渲染DPI page_range: 页面范围字符串,如 "1-5,7,9-12" - PDF:按页码(从1开始) - 图片目录:按文件名排序后的位置(从1开始) Returns: (images_list, pdf_type, pdf_doc) - images_list: 图像列表,每个元素包含 {'img_pil': PIL.Image, 'scale': float, 'page_idx': int} - pdf_type: 'ocr' 或 'txt' - pdf_doc: PDF文档对象(如果是PDF) """ pdf_doc = None pdf_type = 'ocr' # 默认使用OCR模式 all_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 ]) # 解析页面范围 total_pages = len(image_files) selected_pages = parse_page_range(page_range, total_pages) if page_range: logger.info(f"📋 图片目录共 {total_pages} 张,选择处理 {len(selected_pages)} 张") for idx, img_file in enumerate(image_files): if idx not in selected_pages: continue img = Image.open(img_file) if img.mode != 'RGB': img = img.convert('RGB') all_images.append({ 'img_pil': img, 'scale': 1.0, 'source_path': str(img_file), 'page_idx': idx, # 原始索引 'page_name': img_file.stem # 文件名(不含扩展名) }) 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 ) # 解析页面范围 total_pages = len(images_list) selected_pages = parse_page_range(page_range, total_pages) if page_range: logger.info(f"📋 PDF 共 {total_pages} 页,选择处理 {len(selected_pages)} 页") for idx, img_dict in enumerate(images_list): if idx not in selected_pages: continue all_images.append({ 'img_pil': img_dict['img_pil'], 'scale': img_dict.get('scale', dpi / 72), 'source_path': str(document_path), 'page_idx': idx, # 原始页码索引 'page_name': f"{document_path.stem}_page_{idx + 1:03d}" }) 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') all_images.append({ 'img_pil': img, 'scale': 1.0, 'source_path': str(document_path), 'page_idx': 0, 'page_name': document_path.stem }) pdf_type = 'ocr' else: raise ValueError(f"Unsupported file format: {document_path.suffix}") return all_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