| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329 |
- #!/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__":
- # 打印环境变量
- 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())
|