Kaynağa Gözat

feat: add page range parsing and filtering to PDFUtils for enhanced document processing, allowing users to specify which pages to load and process

zhch158_admin 6 gün önce
ebeveyn
işleme
e0bd11a14e
1 değiştirilmiş dosya ile 99 ekleme ve 14 silme
  1. 99 14
      zhch/universal_doc_parser/core/pdf_utils.py

+ 99 - 14
zhch/universal_doc_parser/core/pdf_utils.py

@@ -5,11 +5,13 @@ PDF处理工具模块
 - PDF加载与分类
 - PDF文本提取
 - 跨页表格合并
+- 页面范围解析与过滤
 """
-from typing import Dict, List, Any, Optional, Tuple
+from typing import Dict, List, Any, Optional, Tuple, Set
 from pathlib import Path
 from PIL import Image
 from loguru import logger
+import re
 
 # 导入 MinerU 组件
 try:
@@ -30,26 +32,84 @@ class PDFUtils:
     """PDF处理工具类"""
     
     @staticmethod
+    def parse_page_range(page_range: Optional[str], total_pages: int) -> Set[int]:
+        """
+        解析页面范围字符串
+        
+        支持格式:
+        - "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)
+        """
+        if not page_range or not page_range.strip():
+            return set(range(total_pages))
+        
+        pages = set()
+        parts = page_range.replace(' ', '').split(',')
+        
+        for part in parts:
+            part = part.strip()
+            if not part:
+                continue
+            
+            if '-' in part:
+                # 范围格式
+                match = re.match(r'^(\d*)-(\d*)$', part)
+                if match:
+                    start_str, end_str = match.groups()
+                    start = int(start_str) if start_str else 1
+                    end = int(end_str) if end_str else total_pages
+                    
+                    # 转换为 0-based 索引
+                    start = max(0, start - 1)
+                    end = min(total_pages, end)
+                    
+                    pages.update(range(start, end))
+            else:
+                # 单页
+                try:
+                    page_num = int(part)
+                    if 1 <= page_num <= total_pages:
+                        pages.add(page_num - 1)  # 转换为 0-based 索引
+                except ValueError:
+                    logger.warning(f"Invalid page number: {part}")
+        
+        return pages
+    
+    @staticmethod
     def load_and_classify_document(
         document_path: Path,
-        dpi: int = 200
+        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}
+            - 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模式
-        images = []
+        all_images = []
         
         if document_path.is_dir():
             # 处理目录:遍历所有图片
@@ -59,14 +119,26 @@ class PDFUtils:
                 if f.suffix.lower() in image_extensions
             ])
             
-            for img_file in image_files:
+            # 解析页面范围
+            total_pages = len(image_files)
+            selected_pages = PDFUtils.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')
-                images.append({
+                all_images.append({
                     'img_pil': img,
                     'scale': 1.0,
-                    'source_path': str(img_file)
+                    'source_path': str(img_file),
+                    'page_idx': idx,  # 原始索引
+                    'page_name': img_file.stem  # 文件名(不含扩展名)
                 })
             
             pdf_type = 'ocr'  # 图片目录始终使用OCR模式
@@ -90,11 +162,22 @@ class PDFUtils:
                 image_type=ImageType.PIL
             )
             
-            for img_dict in images_list:
-                images.append({
+            # 解析页面范围
+            total_pages = len(images_list)
+            selected_pages = PDFUtils.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)
+                    'source_path': str(document_path),
+                    'page_idx': idx  # 原始页码索引
                 })
                 
         elif document_path.suffix.lower() in ['.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.tif']:
@@ -102,17 +185,19 @@ class PDFUtils:
             img = Image.open(document_path)
             if img.mode != 'RGB':
                 img = img.convert('RGB')
-            images.append({
+            all_images.append({
                 'img_pil': img,
                 'scale': 1.0,
-                'source_path': str(document_path)
+                '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 images, pdf_type, pdf_doc
+        return all_images, pdf_type, pdf_doc
     
     @staticmethod
     def extract_text_from_pdf(