Browse Source

feat: introduce main_v2 script for enhanced financial document processing, supporting PDF classification, layout detection, parallel OCR and table recognition, and multi-format output options

zhch158_admin 3 days ago
parent
commit
a3f8091998
1 changed files with 319 additions and 0 deletions
  1. 319 0
      zhch/universal_doc_parser/main_v2.py

+ 319 - 0
zhch/universal_doc_parser/main_v2.py

@@ -0,0 +1,319 @@
+#!/usr/bin/env python3
+"""
+金融文档处理统一入口 v2
+支持完整的处理流程:
+1. PDF分类(扫描件/数字原生PDF)
+2. 页面方向识别
+3. Layout检测
+4. 并行处理:文本OCR + 表格VLM识别
+5. 单元格坐标匹配
+6. 多格式输出(JSON、Markdown、HTML、可视化图片)
+
+使用方法:
+    # 处理单个PDF
+    python main_v2.py -i /path/to/document.pdf -c ./config/bank_statement_mineru_vl.yaml
+    
+    # 处理图片目录
+    python main_v2.py -i /path/to/images/ -c ./config/bank_statement_paddle_vl.yaml
+    
+    # 开启debug模式(输出可视化图片)
+    python main_v2.py -i /path/to/doc.pdf -c ./config/xxx.yaml --debug
+"""
+
+import argparse
+import json
+import sys
+import os
+from pathlib import Path
+from loguru import logger
+from datetime import datetime
+
+# 添加项目根目录到 Python 路径
+project_root = Path(__file__).parent
+if str(project_root) not in sys.path:
+    sys.path.insert(0, str(project_root))
+
+# 添加 MinerU 根目录
+# 添加项目根目录到 Python 路径
+project_root = Path(__file__).parents[1]
+if str(project_root) not in sys.path:
+    sys.path.insert(0, str(project_root))
+
+from dotenv import load_dotenv
+load_dotenv(override=True)
+
+from core.pipeline_manager_v2 import EnhancedDocPipeline
+from universal_doc_parser.utils import OutputFormatterV2
+
+
+def setup_logging(log_level: str = "INFO", log_file: str = None):
+    """设置日志"""
+    logger.remove()
+    
+    # 控制台输出
+    logger.add(
+        sys.stdout,
+        level=log_level,
+        format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>"
+    )
+    
+    # 文件输出
+    if log_file:
+        logger.add(
+            log_file,
+            level="DEBUG",
+            format="{time:YYYY-MM-DD HH:mm:ss} | {level: <8} | {name}:{function}:{line} - {message}",
+            rotation="10 MB"
+        )
+
+
+def process_single_input(
+    input_path: Path,
+    config_path: Path,
+    output_dir: Path,
+    debug: bool = False,
+    scene: str = None
+) -> dict:
+    """
+    处理单个输入(文件或目录)
+    
+    Args:
+        input_path: 输入路径
+        config_path: 配置文件路径
+        output_dir: 输出目录
+        debug: 是否开启debug模式
+        scene: 场景类型覆盖
+        
+    Returns:
+        处理结果和输出路径
+    """
+    try:
+        # 初始化处理流水线
+        with EnhancedDocPipeline(str(config_path)) as pipeline:
+            
+            # 覆盖场景设置
+            if scene:
+                pipeline.scene_name = scene
+                logger.info(f"🔄 Scene overridden to: {scene}")
+            
+            logger.info(f"🚀 开始处理: {input_path}")
+            logger.info(f"📋 场景配置: {pipeline.scene_name}")
+            logger.info(f"📁 输出目录: {output_dir}")
+            
+            # 处理文档
+            start_time = datetime.now()
+            results = pipeline.process_document(str(input_path))
+            process_time = (datetime.now() - start_time).total_seconds()
+            
+            logger.info(f"⏱️ 处理耗时: {process_time:.2f}秒")
+            
+            # 构建输出配置
+            output_config = {
+                'save_json': True,
+                'save_markdown': True,
+                'save_html': True,
+                'save_layout_image': debug,
+                'save_ocr_image': debug,
+                'draw_type_label': True,
+                'draw_bbox_number': True,
+            }
+            
+            # 格式化输出
+            logger.info("💾 保存结果...")
+            formatter = OutputFormatterV2(str(output_dir))
+            output_paths = formatter.save_results(results, output_config)
+            
+            # 打印摘要
+            _print_summary(results, output_paths, process_time)
+            
+            return {
+                'success': True,
+                'results': results,
+                'output_paths': output_paths,
+                'process_time': process_time
+            }
+            
+    except Exception as e:
+        logger.error(f"❌ 处理失败: {e}")
+        import traceback
+        traceback.print_exc()
+        return {
+            'success': False,
+            'error': str(e)
+        }
+
+
+def _print_summary(results: dict, output_paths: dict, process_time: float):
+    """打印处理结果摘要"""
+    total_pages = len(results.get('pages', []))
+    
+    total_tables = 0
+    total_text_blocks = 0
+    total_cells = 0
+    
+    for page in results.get('pages', []):
+        for element in page.get('elements', []):
+            elem_type = element.get('type', '')
+            if elem_type in ['table', 'table_body']:
+                total_tables += 1
+                cells = element.get('content', {}).get('cells', [])
+                total_cells += len(cells)
+            elif elem_type in ['text', 'title', 'ocr_text', 'ref_text']:
+                total_text_blocks += 1
+    
+    print(f"\n{'='*60}")
+    print(f"📊 处理摘要")
+    print(f"{'='*60}")
+    print(f"   📄 文档: {results.get('document_path', 'N/A')}")
+    print(f"   🎯 场景: {results.get('scene', 'N/A')}")
+    print(f"   📋 PDF类型: {results.get('metadata', {}).get('pdf_type', 'N/A')}")
+    print(f"   📖 页面数: {total_pages}")
+    print(f"   📋 表格数: {total_tables}")
+    print(f"   📝 文本块: {total_text_blocks}")
+    print(f"   🔢 单元格: {total_cells} (带坐标)")
+    print(f"   ⏱️ 耗时: {process_time:.2f}秒")
+    print(f"{'='*60}")
+    print(f"📁 输出文件:")
+    for key, path in output_paths.items():
+        if isinstance(path, list):
+            for p in path:
+                print(f"   - {p}")
+        else:
+            print(f"   - {path}")
+    print(f"{'='*60}\n")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="金融文档处理工具 v2",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+示例:
+  # 处理单个PDF文件
+  python main_v2.py -i document.pdf -c config/bank_statement_mineru_vl.yaml
+  
+  # 处理图片目录
+  python main_v2.py -i ./images/ -c config/bank_statement_paddle_vl.yaml
+  
+  # 开启debug模式(输出可视化图片)
+  python main_v2.py -i doc.pdf -c config.yaml --debug
+  
+  # 指定输出目录
+  python main_v2.py -i doc.pdf -c config.yaml -o ./my_output/
+        """
+    )
+    
+    parser.add_argument(
+        "--input", "-i",
+        required=True,
+        help="输入路径(PDF文件、图片文件或图片目录)"
+    )
+    parser.add_argument(
+        "--config", "-c",
+        required=True,
+        help="配置文件路径"
+    )
+    parser.add_argument(
+        "--output_dir", "-o",
+        default="./output",
+        help="输出目录(默认: ./output)"
+    )
+    parser.add_argument(
+        "--scene", "-s",
+        choices=["bank_statement", "financial_report"],
+        help="场景类型(覆盖配置文件设置)"
+    )
+    parser.add_argument(
+        "--debug",
+        action="store_true",
+        help="开启debug模式(输出layout和OCR可视化图片)"
+    )
+    parser.add_argument(
+        "--log_level",
+        default="INFO",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
+        help="日志级别(默认: INFO)"
+    )
+    parser.add_argument(
+        "--log_file",
+        help="日志文件路径"
+    )
+    parser.add_argument(
+        "--dry_run",
+        action="store_true",
+        help="仅验证配置,不执行处理"
+    )
+    
+    args = parser.parse_args()
+    
+    # 设置日志
+    setup_logging(args.log_level, args.log_file)
+    
+    # 验证输入
+    input_path = Path(args.input)
+    if not input_path.exists():
+        logger.error(f"❌ 输入路径不存在: {input_path}")
+        return 1
+    
+    # 验证配置文件
+    config_path = Path(args.config)
+    if not config_path.exists():
+        logger.error(f"❌ 配置文件不存在: {config_path}")
+        return 1
+    
+    # 仅验证模式
+    if args.dry_run:
+        logger.info("✅ 配置验证通过(dry run)")
+        return 0
+    
+    # 处理文档
+    result = process_single_input(
+        input_path=input_path,
+        config_path=config_path,
+        output_dir=Path(args.output_dir),
+        debug=args.debug,
+        scene=args.scene
+    )
+    
+    return 0 if result.get('success') else 1
+
+
+if __name__ == "__main__":
+    if len(sys.argv) == 1:
+        # 没有命令行参数时,使用默认配置运行
+        print("ℹ️  未提供命令行参数,使用默认配置运行...")
+        
+        # 默认配置
+        default_config = {
+            # 测试输入
+            "input": "/Users/zhch158/workspace/data/流水分析/2023年度报告母公司/mineru_vllm_results/2023年度报告母公司/2023年度报告母公司_page_003.png",
+            # "input": "/Users/zhch158/workspace/data/流水分析/B用户_扫描流水/B用户_扫描流水_page_022.png",
+            
+            # 配置文件
+            "config": "./config/bank_statement_v2.yaml",
+            # "config": "./config/bank_statement_paddle_vl.yaml",
+            
+            # 输出目录
+            "output_dir": "./output/test_v2",
+            
+            # 场景
+            "scene": "bank_statement",
+            
+            # Debug模式
+            "debug": True,
+            
+            # 日志级别
+            "log_level": "DEBUG",
+        }
+        
+        # 构造参数
+        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())
+