#!/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="{time:YYYY-MM-DD HH:mm:ss} | {level: <8} | {name}:{function}:{line} - {message}" ) # 文件输出 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__": # 打印环境变量 print(f"🔧 CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES', 'Not set')}") print(f"🔧 HF_HOME: {os.environ.get('HF_HOME', 'Not set')}") print(f"🔧 HF_ENDPOINT: {os.environ.get('HF_ENDPOINT', 'Not set')}") print(f"🔧 HF_HUB_OFFLINE: {os.environ.get('HF_HUB_OFFLINE', 'Not set')}") print(f"🔧 TORCH_HOME: {os.environ.get('TORCH_HOME', 'Not set')}") print(f"🔧 MODELSCOPE_CACHE: {os.environ.get('MODELSCOPE_CACHE', 'Not set')}") print(f"🔧 USE_MODELSCOPE_HUB: {os.environ.get('USE_MODELSCOPE_HUB', 'Not set')}") print(f"🔧 MINERU_MODEL_SOURCE: {os.environ.get('MINERU_MODEL_SOURCE', 'Not set')}") if len(sys.argv) == 1: # 没有命令行参数时,使用默认配置运行 print("ℹ️ 未提供命令行参数,使用默认配置运行...") # 默认配置 default_config = { # 测试输入 # "input": "/Users/zhch158/workspace/data/流水分析/康强_北京农村商业银行.pdf", # "output_dir": "./output/康强_北京农村商业银行_bank_statement_v2", "input": "/Users/zhch158/workspace/data/流水分析/2023年度报告母公司/mineru_vllm_results/2023年度报告母公司/2023年度报告母公司_page_003.png", "output_dir": "./output/2023年度报告母公司_bank_statement_v2", # 配置文件 "config": "./config/bank_statement_v2.yaml", # "config": "./config/bank_statement_paddle_vl.yaml", # 场景 "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())