main_v2.py 16 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443
  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 typing import Optional
  25. from loguru import logger
  26. from datetime import datetime
  27. # 添加项目根目录到 Python 路径
  28. project_root = Path(__file__).parent
  29. if str(project_root) not in sys.path:
  30. sys.path.insert(0, str(project_root))
  31. # 添加 MinerU 根目录
  32. # 添加项目根目录到 Python 路径
  33. project_root = Path(__file__).parents[1]
  34. if str(project_root) not in sys.path:
  35. sys.path.insert(0, str(project_root))
  36. from dotenv import load_dotenv
  37. load_dotenv(override=True)
  38. from core.pipeline_manager_v2 import EnhancedDocPipeline
  39. from core.pipeline_manager_v2_streaming import StreamingDocPipeline
  40. from universal_doc_parser.utils import OutputFormatterV2
  41. def setup_logging(log_level: str = "INFO", log_file: Optional[str] = None):
  42. """设置日志"""
  43. logger.remove()
  44. # 控制台输出
  45. logger.add(
  46. sys.stdout,
  47. level=log_level,
  48. 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>"
  49. )
  50. # 文件输出
  51. if log_file:
  52. logger.add(
  53. log_file,
  54. level="DEBUG",
  55. format="{time:YYYY-MM-DD HH:mm:ss} | {level: <8} | {name}:{function}:{line} - {message}",
  56. rotation="10 MB"
  57. )
  58. def process_single_input(
  59. input_path: Path,
  60. config_path: Path,
  61. output_dir: Path,
  62. debug: bool = False,
  63. scene: Optional[str] = None,
  64. page_range: Optional[str] = None,
  65. streaming: bool = False
  66. ) -> dict:
  67. """
  68. 处理单个输入(文件或目录)
  69. Args:
  70. input_path: 输入路径
  71. config_path: 配置文件路径
  72. output_dir: 输出目录
  73. debug: 是否开启debug模式
  74. scene: 场景类型覆盖
  75. page_range: 页面范围(如 "1-5,7,9-12")
  76. streaming: 是否使用流式处理模式(按页处理,立即保存,节省内存)
  77. Returns:
  78. 处理结果和输出路径
  79. """
  80. try:
  81. # 选择处理模式
  82. if streaming:
  83. logger.info("🔄 Using streaming processing mode (memory-efficient)")
  84. pipeline_streaming = StreamingDocPipeline(str(config_path), str(output_dir))
  85. use_context = False # StreamingDocPipeline 不使用 context manager
  86. else:
  87. logger.info("🔄 Using batch processing mode (all pages in memory)")
  88. pipeline_batch = EnhancedDocPipeline(str(config_path))
  89. use_context = hasattr(pipeline_batch, '__enter__')
  90. if use_context:
  91. pipeline_batch = pipeline_batch.__enter__()
  92. try:
  93. # 覆盖场景设置
  94. if streaming:
  95. pipeline = pipeline_streaming
  96. else:
  97. pipeline = pipeline_batch
  98. if scene:
  99. pipeline.scene_name = scene
  100. logger.info(f"🔄 Scene overridden to: {scene}")
  101. logger.info(f"🚀 开始处理: {input_path}")
  102. logger.info(f"📋 场景配置: {pipeline.scene_name}")
  103. logger.info(f"📁 输出目录: {output_dir}")
  104. if page_range:
  105. logger.info(f"📄 页面范围: {page_range}")
  106. # 构建输出配置
  107. output_config = {
  108. 'save_json': True,
  109. 'save_markdown': True,
  110. 'save_html': True,
  111. 'save_page_json': True,
  112. 'save_images': True,
  113. 'save_layout_image': debug,
  114. 'save_ocr_image': debug,
  115. 'normalize_numbers': True,
  116. 'merge_cross_page_tables': True,
  117. 'cleanup_temp_files': True,
  118. }
  119. # 处理文档
  120. start_time = datetime.now()
  121. if streaming:
  122. # 流式处理模式
  123. results = pipeline.process_document_streaming( # type: ignore
  124. str(input_path),
  125. page_range=page_range,
  126. output_config=output_config
  127. )
  128. process_time = (datetime.now() - start_time).total_seconds()
  129. # 流式模式已经保存了所有结果,只需要返回摘要
  130. output_paths = results.get('output_paths', {})
  131. # 打印摘要
  132. _print_summary_streaming(results, process_time)
  133. return {
  134. 'success': True,
  135. 'results': results,
  136. 'output_paths': output_paths,
  137. 'process_time': process_time
  138. }
  139. else:
  140. # 批量处理模式(原有逻辑)
  141. results = pipeline.process_document(str(input_path), page_range=page_range)
  142. process_time = (datetime.now() - start_time).total_seconds()
  143. logger.info(f"⏱️ 处理耗时: {process_time:.2f}秒")
  144. # 格式化输出
  145. logger.info("💾 保存结果...")
  146. formatter = OutputFormatterV2(str(output_dir))
  147. output_paths = formatter.save_results(results, output_config)
  148. # 打印摘要
  149. _print_summary(results, output_paths, process_time)
  150. return {
  151. 'success': True,
  152. 'results': results,
  153. 'output_paths': output_paths,
  154. 'process_time': process_time
  155. }
  156. finally:
  157. # 关闭context manager
  158. if not streaming and use_context:
  159. pipeline_batch.__exit__(None, None, None)
  160. except Exception as e:
  161. logger.error(f"❌ 处理失败: {e}")
  162. import traceback
  163. traceback.print_exc()
  164. return {
  165. 'success': False,
  166. 'error': str(e)
  167. }
  168. def _print_summary(results: dict, output_paths: dict, process_time: float):
  169. """打印处理结果摘要"""
  170. total_pages = len(results.get('pages', []))
  171. total_tables = 0
  172. total_text_blocks = 0
  173. total_cells = 0
  174. for page in results.get('pages', []):
  175. for element in page.get('elements', []):
  176. elem_type = element.get('type', '')
  177. if elem_type in ['table', 'table_body']:
  178. total_tables += 1
  179. cells = element.get('content', {}).get('cells', [])
  180. total_cells += len(cells)
  181. elif elem_type in ['text', 'title', 'ocr_text', 'ref_text']:
  182. total_text_blocks += 1
  183. print(f"\n{'='*60}")
  184. print(f"📊 处理摘要")
  185. print(f"{'='*60}")
  186. print(f" 📄 文档: {results.get('document_path', 'N/A')}")
  187. print(f" 🎯 场景: {results.get('scene', 'N/A')}")
  188. print(f" 📋 PDF类型: {results.get('metadata', {}).get('pdf_type', 'N/A')}")
  189. print(f" 📖 页面数: {total_pages}")
  190. print(f" 📋 表格数: {total_tables}")
  191. print(f" 📝 文本块: {total_text_blocks}")
  192. print(f" 🔢 单元格: {total_cells} (带坐标)")
  193. print(f" ⏱️ 耗时: {process_time:.2f}秒")
  194. print(f"{'='*60}")
  195. print(f"📁 输出文件:")
  196. for key, path in output_paths.items():
  197. if isinstance(path, list):
  198. for p in path:
  199. print(f" - {p}")
  200. else:
  201. print(f" - {path}")
  202. print(f"{'='*60}\n")
  203. def _print_summary_streaming(results_summary: dict, process_time: float):
  204. """打印流式处理结果摘要"""
  205. print(f"\n{'='*60}")
  206. print(f"📊 处理摘要(流式模式)")
  207. print(f"{'='*60}")
  208. print(f" 📄 文档: {results_summary.get('document_path', 'N/A')}")
  209. print(f" 🎯 场景: {results_summary.get('scene', 'N/A')}")
  210. print(f" 📋 PDF类型: {results_summary.get('metadata', {}).get('pdf_type', 'N/A')}")
  211. print(f" 📖 页面数: {results_summary.get('total_pages', 0)}")
  212. print(f" ⏱️ 耗时: {process_time:.2f}秒")
  213. print(f"{'='*60}")
  214. print(f"📁 输出文件:")
  215. output_paths = results_summary.get('output_paths', {})
  216. if output_paths.get('middle_json'):
  217. print(f" - {output_paths['middle_json']}")
  218. if output_paths.get('json_pages'):
  219. print(f" - {len(output_paths['json_pages'])} 个页面JSON文件")
  220. if output_paths.get('images'):
  221. print(f" - {len(output_paths['images'])} 个图片文件")
  222. print(f"{'='*60}\n")
  223. def main():
  224. parser = argparse.ArgumentParser(
  225. description="金融文档处理工具 v2",
  226. formatter_class=argparse.RawDescriptionHelpFormatter,
  227. epilog="""
  228. 示例:
  229. # 处理单个PDF文件
  230. python main_v2.py -i document.pdf -c config/bank_statement_mineru_vl.yaml
  231. # 处理图片目录
  232. python main_v2.py -i ./images/ -c config/bank_statement_paddle_vl.yaml
  233. # 开启debug模式(输出可视化图片)
  234. python main_v2.py -i doc.pdf -c config.yaml --debug
  235. # 指定输出目录
  236. python main_v2.py -i doc.pdf -c config.yaml -o ./my_output/
  237. # 指定页面范围(PDF按页码,图片目录按排序位置)
  238. python main_v2.py -i doc.pdf -c config.yaml -p 1-5 # 处理第1-5页
  239. python main_v2.py -i doc.pdf -c config.yaml -p 3,7,10 # 处理第3、7、10页
  240. python main_v2.py -i doc.pdf -c config.yaml -p 1-5,8-10 # 处理第1-5、8-10页
  241. python main_v2.py -i doc.pdf -c config.yaml -p 5- # 从第5页到最后
  242. # 使用流式处理模式(节省内存,适合大文档)
  243. python main_v2.py -i large_doc.pdf -c config.yaml --streaming
  244. """
  245. )
  246. parser.add_argument(
  247. "--input", "-i",
  248. required=True,
  249. help="输入路径(PDF文件、图片文件或图片目录)"
  250. )
  251. parser.add_argument(
  252. "--config", "-c",
  253. required=True,
  254. help="配置文件路径"
  255. )
  256. parser.add_argument(
  257. "--output_dir", "-o",
  258. default="./output",
  259. help="输出目录(默认: ./output)"
  260. )
  261. parser.add_argument(
  262. "--scene", "-s",
  263. choices=["bank_statement", "financial_report"],
  264. help="场景类型(覆盖配置文件设置)"
  265. )
  266. parser.add_argument(
  267. "--debug",
  268. action="store_true",
  269. help="开启debug模式(输出layout和OCR可视化图片)"
  270. )
  271. parser.add_argument(
  272. "--log_level",
  273. default="INFO",
  274. choices=["DEBUG", "INFO", "WARNING", "ERROR"],
  275. help="日志级别(默认: INFO)"
  276. )
  277. parser.add_argument(
  278. "--log_file",
  279. help="日志文件路径"
  280. )
  281. parser.add_argument(
  282. "--dry_run",
  283. action="store_true",
  284. help="仅验证配置,不执行处理"
  285. )
  286. parser.add_argument(
  287. "--pages", "-p",
  288. help="页面范围(PDF按页码,图片目录按排序位置),如: 1-5,7,9-12"
  289. )
  290. parser.add_argument(
  291. "--streaming",
  292. action="store_true",
  293. help="使用流式处理模式(按页处理,立即保存,节省内存,适合大文档)"
  294. )
  295. args = parser.parse_args()
  296. # 设置日志
  297. setup_logging(args.log_level, args.log_file)
  298. # 验证输入
  299. input_path = Path(args.input)
  300. if not input_path.exists():
  301. logger.error(f"❌ 输入路径不存在: {input_path}")
  302. return 1
  303. # 验证配置文件
  304. config_path = Path(args.config)
  305. if not config_path.exists():
  306. logger.error(f"❌ 配置文件不存在: {config_path}")
  307. return 1
  308. # 仅验证模式
  309. if args.dry_run:
  310. logger.info("✅ 配置验证通过(dry run)")
  311. return 0
  312. # 处理文档
  313. result = process_single_input(
  314. input_path=input_path,
  315. config_path=config_path,
  316. output_dir=Path(args.output_dir),
  317. debug=args.debug,
  318. scene=args.scene,
  319. page_range=args.pages,
  320. streaming=args.streaming
  321. )
  322. return 0 if result.get('success') else 1
  323. if __name__ == "__main__":
  324. # 打印环境变量
  325. print(f"🔧 CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES', 'Not set')}")
  326. print(f"🔧 HF_HOME: {os.environ.get('HF_HOME', 'Not set')}")
  327. print(f"🔧 HF_ENDPOINT: {os.environ.get('HF_ENDPOINT', 'Not set')}")
  328. print(f"🔧 HF_HUB_OFFLINE: {os.environ.get('HF_HUB_OFFLINE', 'Not set')}")
  329. print(f"🔧 TORCH_HOME: {os.environ.get('TORCH_HOME', 'Not set')}")
  330. print(f"🔧 MODELSCOPE_CACHE: {os.environ.get('MODELSCOPE_CACHE', 'Not set')}")
  331. print(f"🔧 USE_MODELSCOPE_HUB: {os.environ.get('USE_MODELSCOPE_HUB', 'Not set')}")
  332. print(f"🔧 MINERU_MODEL_SOURCE: {os.environ.get('MINERU_MODEL_SOURCE', 'Not set')}")
  333. if len(sys.argv) == 1:
  334. # 没有命令行参数时,使用默认配置运行
  335. print("ℹ️ 未提供命令行参数,使用默认配置运行...")
  336. # 默认配置
  337. default_config = {
  338. # 测试输入
  339. # "input": "/Users/zhch158/workspace/data/流水分析/康强_北京农村商业银行.pdf",
  340. # "output_dir": "./output/康强_北京农村商业银行_bank_statement_v2",
  341. # "input": "/Users/zhch158/workspace/data/流水分析/2023年度报告母公司/mineru_vllm_results/2023年度报告母公司/2023年度报告母公司_page_003.png",
  342. # "output_dir": "./output/2023年度报告母公司_bank_statement_v2",
  343. # "input": "/Users/zhch158/workspace/data/流水分析/B用户_扫描流水.pdf",
  344. # "output_dir": "/Users/zhch158/workspace/data/流水分析/B用户_扫描流水/bank_statement_yusys_v2",
  345. # "input": "/Users/zhch158/workspace/data/流水分析/2023年度报告母公司.pdf",
  346. # "output_dir": "/Users/zhch158/workspace/data/流水分析/2023年度报告母公司/bank_statement_yusys_v2",
  347. "input": "/Users/zhch158/workspace/data/流水分析/施博深.pdf",
  348. "output_dir": "/Users/zhch158/workspace/data/流水分析/施博深/bank_statement_yusys_v2",
  349. # 配置文件
  350. "config": "./config/bank_statement_yusys_v2.yaml",
  351. # "config": "./config/bank_statement_v2.yaml",
  352. # "config": "./config/bank_statement_paddle_vl.yaml",
  353. # 场景
  354. "scene": "bank_statement",
  355. # 页面范围(可选)
  356. "pages": "1", # 只处理前4页
  357. # "pages": "1-3,5,7-10", # 处理指定页面
  358. "streaming": True,
  359. # Debug模式
  360. "debug": True,
  361. # 日志级别
  362. "log_level": "DEBUG",
  363. }
  364. # 构造参数
  365. sys.argv = [sys.argv[0]]
  366. for key, value in default_config.items():
  367. if isinstance(value, bool):
  368. if value:
  369. sys.argv.append(f"--{key}")
  370. else:
  371. sys.argv.extend([f"--{key}", str(value)])
  372. sys.exit(main())