#!/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())