main_v2.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329
  1. #!/usr/bin/env python3
  2. """
  3. 金融文档处理统一入口 v2
  4. 支持完整的处理流程:
  5. 1. PDF分类(扫描件/数字原生PDF)
  6. 2. 页面方向识别
  7. 3. Layout检测
  8. 4. 并行处理:文本OCR + 表格VLM识别
  9. 5. 单元格坐标匹配
  10. 6. 多格式输出(JSON、Markdown、HTML、可视化图片)
  11. 使用方法:
  12. # 处理单个PDF
  13. python main_v2.py -i /path/to/document.pdf -c ./config/bank_statement_mineru_vl.yaml
  14. # 处理图片目录
  15. python main_v2.py -i /path/to/images/ -c ./config/bank_statement_paddle_vl.yaml
  16. # 开启debug模式(输出可视化图片)
  17. python main_v2.py -i /path/to/doc.pdf -c ./config/xxx.yaml --debug
  18. """
  19. import argparse
  20. import json
  21. import sys
  22. import os
  23. from pathlib import Path
  24. from loguru import logger
  25. from datetime import datetime
  26. # 添加项目根目录到 Python 路径
  27. project_root = Path(__file__).parent
  28. if str(project_root) not in sys.path:
  29. sys.path.insert(0, str(project_root))
  30. # 添加 MinerU 根目录
  31. # 添加项目根目录到 Python 路径
  32. project_root = Path(__file__).parents[1]
  33. if str(project_root) not in sys.path:
  34. sys.path.insert(0, str(project_root))
  35. from dotenv import load_dotenv
  36. load_dotenv(override=True)
  37. from core.pipeline_manager_v2 import EnhancedDocPipeline
  38. from universal_doc_parser.utils import OutputFormatterV2
  39. def setup_logging(log_level: str = "INFO", log_file: str = None):
  40. """设置日志"""
  41. logger.remove()
  42. # 控制台输出
  43. logger.add(
  44. sys.stdout,
  45. level=log_level,
  46. 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>"
  47. )
  48. # 文件输出
  49. if log_file:
  50. logger.add(
  51. log_file,
  52. level="DEBUG",
  53. format="{time:YYYY-MM-DD HH:mm:ss} | {level: <8} | {name}:{function}:{line} - {message}",
  54. rotation="10 MB"
  55. )
  56. def process_single_input(
  57. input_path: Path,
  58. config_path: Path,
  59. output_dir: Path,
  60. debug: bool = False,
  61. scene: str = None
  62. ) -> dict:
  63. """
  64. 处理单个输入(文件或目录)
  65. Args:
  66. input_path: 输入路径
  67. config_path: 配置文件路径
  68. output_dir: 输出目录
  69. debug: 是否开启debug模式
  70. scene: 场景类型覆盖
  71. Returns:
  72. 处理结果和输出路径
  73. """
  74. try:
  75. # 初始化处理流水线
  76. with EnhancedDocPipeline(str(config_path)) as pipeline:
  77. # 覆盖场景设置
  78. if scene:
  79. pipeline.scene_name = scene
  80. logger.info(f"🔄 Scene overridden to: {scene}")
  81. logger.info(f"🚀 开始处理: {input_path}")
  82. logger.info(f"📋 场景配置: {pipeline.scene_name}")
  83. logger.info(f"📁 输出目录: {output_dir}")
  84. # 处理文档
  85. start_time = datetime.now()
  86. results = pipeline.process_document(str(input_path))
  87. process_time = (datetime.now() - start_time).total_seconds()
  88. logger.info(f"⏱️ 处理耗时: {process_time:.2f}秒")
  89. # 构建输出配置
  90. output_config = {
  91. 'save_json': True,
  92. 'save_markdown': True,
  93. 'save_html': True,
  94. 'save_layout_image': debug,
  95. 'save_ocr_image': debug,
  96. 'draw_type_label': True,
  97. 'draw_bbox_number': True,
  98. }
  99. # 格式化输出
  100. logger.info("💾 保存结果...")
  101. formatter = OutputFormatterV2(str(output_dir))
  102. output_paths = formatter.save_results(results, output_config)
  103. # 打印摘要
  104. _print_summary(results, output_paths, process_time)
  105. return {
  106. 'success': True,
  107. 'results': results,
  108. 'output_paths': output_paths,
  109. 'process_time': process_time
  110. }
  111. except Exception as e:
  112. logger.error(f"❌ 处理失败: {e}")
  113. import traceback
  114. traceback.print_exc()
  115. return {
  116. 'success': False,
  117. 'error': str(e)
  118. }
  119. def _print_summary(results: dict, output_paths: dict, process_time: float):
  120. """打印处理结果摘要"""
  121. total_pages = len(results.get('pages', []))
  122. total_tables = 0
  123. total_text_blocks = 0
  124. total_cells = 0
  125. for page in results.get('pages', []):
  126. for element in page.get('elements', []):
  127. elem_type = element.get('type', '')
  128. if elem_type in ['table', 'table_body']:
  129. total_tables += 1
  130. cells = element.get('content', {}).get('cells', [])
  131. total_cells += len(cells)
  132. elif elem_type in ['text', 'title', 'ocr_text', 'ref_text']:
  133. total_text_blocks += 1
  134. print(f"\n{'='*60}")
  135. print(f"📊 处理摘要")
  136. print(f"{'='*60}")
  137. print(f" 📄 文档: {results.get('document_path', 'N/A')}")
  138. print(f" 🎯 场景: {results.get('scene', 'N/A')}")
  139. print(f" 📋 PDF类型: {results.get('metadata', {}).get('pdf_type', 'N/A')}")
  140. print(f" 📖 页面数: {total_pages}")
  141. print(f" 📋 表格数: {total_tables}")
  142. print(f" 📝 文本块: {total_text_blocks}")
  143. print(f" 🔢 单元格: {total_cells} (带坐标)")
  144. print(f" ⏱️ 耗时: {process_time:.2f}秒")
  145. print(f"{'='*60}")
  146. print(f"📁 输出文件:")
  147. for key, path in output_paths.items():
  148. if isinstance(path, list):
  149. for p in path:
  150. print(f" - {p}")
  151. else:
  152. print(f" - {path}")
  153. print(f"{'='*60}\n")
  154. def main():
  155. parser = argparse.ArgumentParser(
  156. description="金融文档处理工具 v2",
  157. formatter_class=argparse.RawDescriptionHelpFormatter,
  158. epilog="""
  159. 示例:
  160. # 处理单个PDF文件
  161. python main_v2.py -i document.pdf -c config/bank_statement_mineru_vl.yaml
  162. # 处理图片目录
  163. python main_v2.py -i ./images/ -c config/bank_statement_paddle_vl.yaml
  164. # 开启debug模式(输出可视化图片)
  165. python main_v2.py -i doc.pdf -c config.yaml --debug
  166. # 指定输出目录
  167. python main_v2.py -i doc.pdf -c config.yaml -o ./my_output/
  168. """
  169. )
  170. parser.add_argument(
  171. "--input", "-i",
  172. required=True,
  173. help="输入路径(PDF文件、图片文件或图片目录)"
  174. )
  175. parser.add_argument(
  176. "--config", "-c",
  177. required=True,
  178. help="配置文件路径"
  179. )
  180. parser.add_argument(
  181. "--output_dir", "-o",
  182. default="./output",
  183. help="输出目录(默认: ./output)"
  184. )
  185. parser.add_argument(
  186. "--scene", "-s",
  187. choices=["bank_statement", "financial_report"],
  188. help="场景类型(覆盖配置文件设置)"
  189. )
  190. parser.add_argument(
  191. "--debug",
  192. action="store_true",
  193. help="开启debug模式(输出layout和OCR可视化图片)"
  194. )
  195. parser.add_argument(
  196. "--log_level",
  197. default="INFO",
  198. choices=["DEBUG", "INFO", "WARNING", "ERROR"],
  199. help="日志级别(默认: INFO)"
  200. )
  201. parser.add_argument(
  202. "--log_file",
  203. help="日志文件路径"
  204. )
  205. parser.add_argument(
  206. "--dry_run",
  207. action="store_true",
  208. help="仅验证配置,不执行处理"
  209. )
  210. args = parser.parse_args()
  211. # 设置日志
  212. setup_logging(args.log_level, args.log_file)
  213. # 验证输入
  214. input_path = Path(args.input)
  215. if not input_path.exists():
  216. logger.error(f"❌ 输入路径不存在: {input_path}")
  217. return 1
  218. # 验证配置文件
  219. config_path = Path(args.config)
  220. if not config_path.exists():
  221. logger.error(f"❌ 配置文件不存在: {config_path}")
  222. return 1
  223. # 仅验证模式
  224. if args.dry_run:
  225. logger.info("✅ 配置验证通过(dry run)")
  226. return 0
  227. # 处理文档
  228. result = process_single_input(
  229. input_path=input_path,
  230. config_path=config_path,
  231. output_dir=Path(args.output_dir),
  232. debug=args.debug,
  233. scene=args.scene
  234. )
  235. return 0 if result.get('success') else 1
  236. if __name__ == "__main__":
  237. # 打印环境变量
  238. print(f"🔧 CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES', 'Not set')}")
  239. print(f"🔧 HF_HOME: {os.environ.get('HF_HOME', 'Not set')}")
  240. print(f"🔧 HF_ENDPOINT: {os.environ.get('HF_ENDPOINT', 'Not set')}")
  241. print(f"🔧 HF_HUB_OFFLINE: {os.environ.get('HF_HUB_OFFLINE', 'Not set')}")
  242. print(f"🔧 TORCH_HOME: {os.environ.get('TORCH_HOME', 'Not set')}")
  243. print(f"🔧 MODELSCOPE_CACHE: {os.environ.get('MODELSCOPE_CACHE', 'Not set')}")
  244. print(f"🔧 USE_MODELSCOPE_HUB: {os.environ.get('USE_MODELSCOPE_HUB', 'Not set')}")
  245. print(f"🔧 MINERU_MODEL_SOURCE: {os.environ.get('MINERU_MODEL_SOURCE', 'Not set')}")
  246. if len(sys.argv) == 1:
  247. # 没有命令行参数时,使用默认配置运行
  248. print("ℹ️ 未提供命令行参数,使用默认配置运行...")
  249. # 默认配置
  250. default_config = {
  251. # 测试输入
  252. # "input": "/Users/zhch158/workspace/data/流水分析/康强_北京农村商业银行.pdf",
  253. # "output_dir": "./output/康强_北京农村商业银行_bank_statement_v2",
  254. "input": "/Users/zhch158/workspace/data/流水分析/2023年度报告母公司/mineru_vllm_results/2023年度报告母公司/2023年度报告母公司_page_003.png",
  255. "output_dir": "./output/2023年度报告母公司_bank_statement_v2",
  256. # 配置文件
  257. "config": "./config/bank_statement_v2.yaml",
  258. # "config": "./config/bank_statement_paddle_vl.yaml",
  259. # 场景
  260. "scene": "bank_statement",
  261. # Debug模式
  262. "debug": True,
  263. # 日志级别
  264. "log_level": "DEBUG",
  265. }
  266. # 构造参数
  267. sys.argv = [sys.argv[0]]
  268. for key, value in default_config.items():
  269. if isinstance(value, bool):
  270. if value:
  271. sys.argv.append(f"--{key}")
  272. else:
  273. sys.argv.extend([f"--{key}", str(value)])
  274. sys.exit(main())