Explorar o código

feat: implement streaming processing mode in document parser, allowing memory-efficient handling of large documents with enhanced output summary and configuration options

zhch158_admin hai 5 días
pai
achega
6043464fee
Modificáronse 1 ficheiros con 115 adicións e 30 borrados
  1. 115 30
      zhch/universal_doc_parser/main_v2.py

+ 115 - 30
zhch/universal_doc_parser/main_v2.py

@@ -25,6 +25,7 @@ import json
 import sys
 import sys
 import os
 import os
 from pathlib import Path
 from pathlib import Path
+from typing import Optional
 from loguru import logger
 from loguru import logger
 from datetime import datetime
 from datetime import datetime
 
 
@@ -43,10 +44,11 @@ from dotenv import load_dotenv
 load_dotenv(override=True)
 load_dotenv(override=True)
 
 
 from core.pipeline_manager_v2 import EnhancedDocPipeline
 from core.pipeline_manager_v2 import EnhancedDocPipeline
+from core.pipeline_manager_v2_streaming import StreamingDocPipeline
 from universal_doc_parser.utils import OutputFormatterV2
 from universal_doc_parser.utils import OutputFormatterV2
 
 
 
 
-def setup_logging(log_level: str = "INFO", log_file: str = None):
+def setup_logging(log_level: str = "INFO", log_file: Optional[str] = None):
     """设置日志"""
     """设置日志"""
     logger.remove()
     logger.remove()
     
     
@@ -72,8 +74,9 @@ def process_single_input(
     config_path: Path,
     config_path: Path,
     output_dir: Path,
     output_dir: Path,
     debug: bool = False,
     debug: bool = False,
-    scene: str = None,
-    page_range: str = None
+    scene: Optional[str] = None,
+    page_range: Optional[str] = None,
+    streaming: bool = False
 ) -> dict:
 ) -> dict:
     """
     """
     处理单个输入(文件或目录)
     处理单个输入(文件或目录)
@@ -85,15 +88,32 @@ def process_single_input(
         debug: 是否开启debug模式
         debug: 是否开启debug模式
         scene: 场景类型覆盖
         scene: 场景类型覆盖
         page_range: 页面范围(如 "1-5,7,9-12")
         page_range: 页面范围(如 "1-5,7,9-12")
+        streaming: 是否使用流式处理模式(按页处理,立即保存,节省内存)
         
         
     Returns:
     Returns:
         处理结果和输出路径
         处理结果和输出路径
     """
     """
     try:
     try:
-        # 初始化处理流水线
-        with EnhancedDocPipeline(str(config_path)) as pipeline:
+        # 选择处理模式
+        if streaming:
+            logger.info("🔄 Using streaming processing mode (memory-efficient)")
+            pipeline_streaming = StreamingDocPipeline(str(config_path), str(output_dir))
+            use_context = False  # StreamingDocPipeline 不使用 context manager
+        else:
+            logger.info("🔄 Using batch processing mode (all pages in memory)")
+            pipeline_batch = EnhancedDocPipeline(str(config_path))
+            use_context = hasattr(pipeline_batch, '__enter__')
+            if use_context:
+                pipeline_batch = pipeline_batch.__enter__()
+        
+        try:
             
             
             # 覆盖场景设置
             # 覆盖场景设置
+            if streaming:
+                pipeline = pipeline_streaming
+            else:
+                pipeline = pipeline_batch
+            
             if scene:
             if scene:
                 pipeline.scene_name = scene
                 pipeline.scene_name = scene
                 logger.info(f"🔄 Scene overridden to: {scene}")
                 logger.info(f"🔄 Scene overridden to: {scene}")
@@ -104,38 +124,70 @@ def process_single_input(
             if page_range:
             if page_range:
                 logger.info(f"📄 页面范围: {page_range}")
                 logger.info(f"📄 页面范围: {page_range}")
             
             
-            # 处理文档
-            start_time = datetime.now()
-            results = pipeline.process_document(str(input_path), page_range=page_range)
-            process_time = (datetime.now() - start_time).total_seconds()
-            
-            logger.info(f"⏱️ 处理耗时: {process_time:.2f}秒")
-            
             # 构建输出配置
             # 构建输出配置
             output_config = {
             output_config = {
                 'save_json': True,
                 'save_json': True,
                 'save_markdown': True,
                 'save_markdown': True,
                 'save_html': True,
                 'save_html': True,
+                'save_page_json': True,
+                'save_images': True,
                 'save_layout_image': debug,
                 'save_layout_image': debug,
                 'save_ocr_image': debug,
                 'save_ocr_image': debug,
-                'draw_type_label': True,
-                'draw_bbox_number': True,
+                'normalize_numbers': True,
+                'merge_cross_page_tables': True,
+                'cleanup_temp_files': True,
             }
             }
             
             
-            # 格式化输出
-            logger.info("💾 保存结果...")
-            formatter = OutputFormatterV2(str(output_dir))
-            output_paths = formatter.save_results(results, output_config)
-            
-            # 打印摘要
-            _print_summary(results, output_paths, process_time)
+            # 处理文档
+            start_time = datetime.now()
             
             
-            return {
-                'success': True,
-                'results': results,
-                'output_paths': output_paths,
-                'process_time': process_time
-            }
+            if streaming:
+                # 流式处理模式
+                results = pipeline.process_document_streaming(  # type: ignore
+                    str(input_path),
+                    page_range=page_range,
+                    output_config=output_config
+                )
+                process_time = (datetime.now() - start_time).total_seconds()
+                
+                # 流式模式已经保存了所有结果,只需要返回摘要
+                output_paths = results.get('output_paths', {})
+                
+                # 打印摘要
+                _print_summary_streaming(results, process_time)
+                
+                return {
+                    'success': True,
+                    'results': results,
+                    'output_paths': output_paths,
+                    'process_time': process_time
+                }
+            else:
+                # 批量处理模式(原有逻辑)
+                results = pipeline.process_document(str(input_path), page_range=page_range)
+                process_time = (datetime.now() - start_time).total_seconds()
+                
+                logger.info(f"⏱️ 处理耗时: {process_time:.2f}秒")
+                
+                # 格式化输出
+                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
+                }
+        
+        finally:
+            # 关闭context manager
+            if not streaming and use_context:
+                pipeline_batch.__exit__(None, None, None)
             
             
     except Exception as e:
     except Exception as e:
         logger.error(f"❌ 处理失败: {e}")
         logger.error(f"❌ 处理失败: {e}")
@@ -187,6 +239,28 @@ def _print_summary(results: dict, output_paths: dict, process_time: float):
     print(f"{'='*60}\n")
     print(f"{'='*60}\n")
 
 
 
 
+def _print_summary_streaming(results_summary: dict, process_time: float):
+    """打印流式处理结果摘要"""
+    print(f"\n{'='*60}")
+    print(f"📊 处理摘要(流式模式)")
+    print(f"{'='*60}")
+    print(f"   📄 文档: {results_summary.get('document_path', 'N/A')}")
+    print(f"   🎯 场景: {results_summary.get('scene', 'N/A')}")
+    print(f"   📋 PDF类型: {results_summary.get('metadata', {}).get('pdf_type', 'N/A')}")
+    print(f"   📖 页面数: {results_summary.get('total_pages', 0)}")
+    print(f"   ⏱️ 耗时: {process_time:.2f}秒")
+    print(f"{'='*60}")
+    print(f"📁 输出文件:")
+    output_paths = results_summary.get('output_paths', {})
+    if output_paths.get('middle_json'):
+        print(f"   - {output_paths['middle_json']}")
+    if output_paths.get('json_pages'):
+        print(f"   - {len(output_paths['json_pages'])} 个页面JSON文件")
+    if output_paths.get('images'):
+        print(f"   - {len(output_paths['images'])} 个图片文件")
+    print(f"{'='*60}\n")
+
+
 def main():
 def main():
     parser = argparse.ArgumentParser(
     parser = argparse.ArgumentParser(
         description="金融文档处理工具 v2",
         description="金融文档处理工具 v2",
@@ -210,6 +284,9 @@ def main():
   python main_v2.py -i doc.pdf -c config.yaml -p 3,7,10   # 处理第3、7、10页
   python main_v2.py -i doc.pdf -c config.yaml -p 3,7,10   # 处理第3、7、10页
   python main_v2.py -i doc.pdf -c config.yaml -p 1-5,8-10 # 处理第1-5、8-10页
   python main_v2.py -i doc.pdf -c config.yaml -p 1-5,8-10 # 处理第1-5、8-10页
   python main_v2.py -i doc.pdf -c config.yaml -p 5-       # 从第5页到最后
   python main_v2.py -i doc.pdf -c config.yaml -p 5-       # 从第5页到最后
+  
+  # 使用流式处理模式(节省内存,适合大文档)
+  python main_v2.py -i large_doc.pdf -c config.yaml --streaming
         """
         """
     )
     )
     
     
@@ -257,6 +334,11 @@ def main():
         "--pages", "-p",
         "--pages", "-p",
         help="页面范围(PDF按页码,图片目录按排序位置),如: 1-5,7,9-12"
         help="页面范围(PDF按页码,图片目录按排序位置),如: 1-5,7,9-12"
     )
     )
+    parser.add_argument(
+        "--streaming",
+        action="store_true",
+        help="使用流式处理模式(按页处理,立即保存,节省内存,适合大文档)"
+    )
     
     
     args = parser.parse_args()
     args = parser.parse_args()
     
     
@@ -287,7 +369,8 @@ def main():
         output_dir=Path(args.output_dir),
         output_dir=Path(args.output_dir),
         debug=args.debug,
         debug=args.debug,
         scene=args.scene,
         scene=args.scene,
-        page_range=args.pages
+        page_range=args.pages,
+        streaming=args.streaming
     )
     )
     
     
     return 0 if result.get('success') else 1
     return 0 if result.get('success') else 1
@@ -332,9 +415,11 @@ if __name__ == "__main__":
             "scene": "bank_statement",
             "scene": "bank_statement",
             
             
             # 页面范围(可选)
             # 页面范围(可选)
-            # "pages": "3",  # 只处理前5
+            "pages": "4",  # 只处理前4
             # "pages": "1-3,5,7-10",  # 处理指定页面
             # "pages": "1-3,5,7-10",  # 处理指定页面
-            
+
+            "streaming": True,
+
             # Debug模式
             # Debug模式
             "debug": True,
             "debug": True,