瀏覽代碼

feat: add PDF page extraction utility for selective page saving, including command-line interface and error handling

zhch158_admin 4 天之前
父節點
當前提交
8131b73341
共有 1 個文件被更改,包括 246 次插入0 次删除
  1. 246 0
      zhch/universal_doc_parser/utils/pdf_extractor.py

+ 246 - 0
zhch/universal_doc_parser/utils/pdf_extractor.py

@@ -0,0 +1,246 @@
+#!/usr/bin/env python3
+"""
+PDF页面提取工具
+
+从PDF文件中提取指定页面并保存为新PDF文件。
+
+使用方法:
+    python pdf_extractor.py input.pdf --pages "1-5,7,9-12" --output output.pdf
+    python pdf_extractor.py input.pdf --pages "1-" --output output.pdf  # 提取第1页到最后
+    python pdf_extractor.py input.pdf --pages "-10" --output output.pdf  # 提取前10页
+"""
+import argparse
+import sys
+from pathlib import Path
+from typing import List
+import io
+
+try:
+    import pypdfium2 as pdfium
+    PDFIUM_AVAILABLE = True
+except ImportError:
+    PDFIUM_AVAILABLE = False
+    pdfium = None
+
+from loguru import logger
+from core.pdf_utils import PDFUtils
+
+
+def extract_pdf_pages(
+    input_pdf_path: Path,
+    page_indices: List[int],
+    output_pdf_path: Path
+) -> bool:
+    """
+    从PDF中提取指定页面并保存为新PDF
+    
+    Args:
+        input_pdf_path: 输入PDF文件路径
+        page_indices: 要提取的页面索引列表(0-based)
+        output_pdf_path: 输出PDF文件路径
+        
+    Returns:
+        是否成功
+    """
+    if not PDFIUM_AVAILABLE:
+        logger.error("❌ pypdfium2 未安装,请先安装: pip install pypdfium2")
+        return False
+    
+    if not input_pdf_path.exists():
+        logger.error(f"❌ 输入文件不存在: {input_pdf_path}")
+        return False
+    
+    if not input_pdf_path.suffix.lower() == '.pdf':
+        logger.error(f"❌ 输入文件不是PDF格式: {input_pdf_path}")
+        return False
+    
+    try:
+        # 读取PDF文件
+        with open(input_pdf_path, 'rb') as f:
+            pdf_bytes = f.read()
+        
+        # 加载PDF文档
+        pdf = pdfium.PdfDocument(pdf_bytes)
+        total_pages = len(pdf)
+        
+        if total_pages == 0:
+            logger.error("❌ PDF文件为空")
+            pdf.close()
+            return False
+        
+        # 验证页面索引
+        valid_indices = []
+        for idx in sorted(set(page_indices)):  # 去重并排序
+            if 0 <= idx < total_pages:
+                valid_indices.append(idx)
+            else:
+                logger.warning(f"⚠️  页面索引 {idx + 1} 超出范围(总页数: {total_pages}),已跳过")
+        
+        if not valid_indices:
+            logger.error("❌ 没有有效的页面可提取")
+            pdf.close()
+            return False
+        
+        # 创建新PDF文档
+        output_pdf = pdfium.PdfDocument.new()
+        
+        # 导入指定页面
+        success_count = 0
+        for page_idx in valid_indices:
+            try:
+                output_pdf.import_pages(pdf, pages=[page_idx])
+                success_count += 1
+            except Exception as e:
+                logger.warning(f"⚠️  导入第 {page_idx + 1} 页失败: {e},已跳过")
+                continue
+        
+        if success_count == 0:
+            logger.error("❌ 没有成功导入任何页面")
+            pdf.close()
+            output_pdf.close()
+            return False
+        
+        # 保存到文件
+        output_pdf_path.parent.mkdir(parents=True, exist_ok=True)
+        
+        # 保存到内存缓冲区
+        output_buffer = io.BytesIO()
+        output_pdf.save(output_buffer)
+        output_bytes = output_buffer.getvalue()
+        
+        # 写入文件
+        with open(output_pdf_path, 'wb') as f:
+            f.write(output_bytes)
+        
+        # 清理资源
+        pdf.close()
+        output_pdf.close()
+        
+        logger.info(f"✅ 成功提取 {success_count} 页到: {output_pdf_path}")
+        logger.info(f"   提取的页面: {', '.join([str(idx + 1) for idx in valid_indices])}")
+        return True
+        
+    except Exception as e:
+        logger.error(f"❌ 提取PDF页面时出错: {e}")
+        import traceback
+        logger.debug(traceback.format_exc())
+        return False
+
+
+def main():
+    """命令行入口"""
+    parser = argparse.ArgumentParser(
+        description='从PDF文件中提取指定页面并保存为新PDF文件',
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+示例:
+  # 提取第1-5页和第7页
+  python pdf_extractor.py --input input.pdf --pages "1-5,7" --output output.pdf
+  
+  # 提取第1页到最后
+  python pdf_extractor.py --input input.pdf --pages "1-" --output output.pdf
+  
+  # 提取前10页
+  python pdf_extractor.py --input input.pdf --pages "-10" --output output.pdf
+  
+  # 提取单页
+  python pdf_extractor.py --input input.pdf --pages "3" --output output.pdf
+        """
+    )
+    
+    parser.add_argument(
+        '--input', '-i',
+        type=str,
+        help='输入PDF文件路径'
+    )
+    
+    parser.add_argument(
+        '--pages', '-p',
+        type=str,
+        required=True,
+        help='要提取的页面范围,支持格式: "1-5,7,9-12", "1-", "-10", "3"'
+    )
+    
+    parser.add_argument(
+        '--output', '-o',
+        type=str,
+        required=True,
+        help='输出PDF文件路径'
+    )
+    
+    args = parser.parse_args()
+    
+    # 解析输入路径
+    input_path = Path(args.input).resolve()
+    
+    # 解析输出路径
+    output_path = Path(args.output).resolve()
+    
+    # 如果输出路径是目录,自动生成文件名
+    if output_path.is_dir() or not output_path.suffix.lower() == '.pdf':
+        output_path = output_path / f"{input_path.stem}_extracted.pdf"
+    
+    # 检查 pypdfium2 是否可用
+    if not PDFIUM_AVAILABLE:
+        logger.error("❌ pypdfium2 未安装,请先安装: pip install pypdfium2")
+        sys.exit(1)
+    
+    # 先加载PDF获取总页数(用于验证页面范围)
+    try:
+        with open(input_path, 'rb') as f:
+            pdf_bytes = f.read()
+        pdf = pdfium.PdfDocument(pdf_bytes)
+        total_pages = len(pdf)
+        pdf.close()
+    except Exception as e:
+        logger.error(f"❌ 无法读取PDF文件: {e}")
+        sys.exit(1)
+    
+    # 解析页面范围
+    page_set = PDFUtils.parse_page_range(args.pages, total_pages)
+    page_indices = sorted(list(page_set))
+    
+    if not page_indices:
+        logger.error(f"❌ 页面范围 '{args.pages}' 没有匹配到任何有效页面(总页数: {total_pages})")
+        sys.exit(1)
+    
+    logger.info(f"📋 PDF总页数: {total_pages}")
+    logger.info(f"📋 要提取的页面: {args.pages} → {len(page_indices)} 页")
+    
+    # 执行提取
+    success = extract_pdf_pages(input_path, page_indices, output_path)
+    
+    if success:
+        logger.info(f"✅ 提取完成!输出文件: {output_path}")
+        sys.exit(0)
+    else:
+        logger.error("❌ 提取失败")
+        sys.exit(1)
+
+
+if __name__ == '__main__':
+    if len(sys.argv) == 1:
+        # 没有命令行参数时,使用默认配置运行
+        print("ℹ️  未提供命令行参数,使用默认配置运行...")
+        
+        # 默认配置
+        default_config = {
+            "input": "/Users/zhch158/workspace/data/流水分析/施博深.pdf",
+            "output": "/Users/zhch158/workspace/data/流水分析/施博深_extracted.pdf",
+
+            # 页面范围(可选)
+            "pages": "1-20",  # 只处理前20页
+            # "pages": "1-3,5,7-10",  # 处理指定页面
+        }
+        
+        # 构造参数
+        sys.argv = [sys.argv[0]]
+        for key, value in default_config.items():
+            if isinstance(value, bool):
+                if value:
+                    sys.argv.append(f"--{key}")
+            else:
+                sys.argv.extend([f"--{key}", str(value)])
+    
+    sys.exit(main())
+