|
|
@@ -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())
|
|
|
+
|