Procházet zdrojové kódy

feat(main_v2): 添加环境变量打印和参数验证功能,优化处理流水线创建逻辑

zhch158_admin před 3 dny
rodič
revize
4afb8f9fa5
1 změnil soubory, kde provedl 104 přidání a 83 odebrání
  1. 104 83
      ocr_tools/universal_doc_parser/main_v2.py

+ 104 - 83
ocr_tools/universal_doc_parser/main_v2.py

@@ -53,6 +53,76 @@ except ImportError:
     from utils import OutputFormatterV2
 
 
+# ==================== Helper Functions ====================
+
+def _print_environment_info():
+    """打印环境变量信息"""
+    env_vars = [
+        'CUDA_VISIBLE_DEVICES', 'HF_HOME', 'HF_ENDPOINT', 'HF_HUB_OFFLINE',
+        'TORCH_HOME', 'MODELSCOPE_CACHE', 'USE_MODELSCOPE_HUB', 'MINERU_MODEL_SOURCE'
+    ]
+    for var in env_vars:
+        print(f"🔧 {var}: {os.environ.get(var, 'Not set')}")
+
+
+def _validate_arguments(args: argparse.Namespace) -> bool:
+    """验证命令行参数"""
+    input_path = Path(args.input)
+    if not input_path.exists():
+        logger.error(f"❌ 输入路径不存在: {input_path}")
+        return False
+    
+    config_path = Path(args.config)
+    if not config_path.exists():
+        logger.error(f"❌ 配置文件不存在: {config_path}")
+        return False
+    
+    return True
+
+
+def _handle_dry_run(args: argparse.Namespace) -> bool:
+    """处理dry run模式"""
+    if args.dry_run:
+        if _validate_arguments(args):
+            logger.info("✅ 配置验证通过(dry run)")
+            return True
+        return False
+    return False
+
+
+def _create_pipeline(streaming: bool, config_path: str, output_dir: str):
+    """创建并初始化处理流水线"""
+    if streaming:
+        logger.info("🔄 Using streaming processing mode (memory-efficient)")
+        pipeline = StreamingDocPipeline(config_path, output_dir)
+    else:
+        logger.info("🔄 Using batch processing mode (all pages in memory)")
+        pipeline = EnhancedDocPipeline(config_path)
+    
+    return pipeline
+
+
+def _get_default_output_config(debug: bool) -> dict:
+    """获取默认输出配置"""
+    return {
+        'create_subdir': True,
+        'save_pdf_images': False,
+        'save_json': True,
+        'save_markdown': True,
+        'save_html': True,
+        'save_page_json': True,
+        'save_images': True,
+        'save_layout_image': debug,
+        'save_ocr_image': debug,
+        'draw_type_label': True,
+        'draw_bbox_number': True,
+        'save_enhanced_json': True,
+        'coordinate_precision': 2,
+        'normalize_numbers': True,
+        'merge_cross_page_tables': True,
+    }
+
+
 def setup_logging(log_level: str = "INFO", log_file: Optional[str] = None):
     """设置日志"""
     logger.remove()
@@ -99,26 +169,15 @@ def process_single_input(
         处理结果和输出路径
     """
     try:
-        # 选择处理模式
-        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__()
+        # 创建流水线
+        pipeline = _create_pipeline(streaming, str(config_path), str(output_dir))
+        output_config = pipeline.config.get('output', {}) or _get_default_output_config(debug)
+        
+        use_context = not streaming and hasattr(pipeline, '__enter__')
+        if use_context:
+            pipeline = pipeline.__enter__()
         
         try:
-            
-            # 覆盖场景设置
-            if streaming:
-                pipeline = pipeline_streaming
-            else:
-                pipeline = pipeline_batch
-            
             if scene:
                 pipeline.scene_name = scene
                 logger.info(f"🔄 Scene overridden to: {scene}")
@@ -129,62 +188,37 @@ def process_single_input(
             if page_range:
                 logger.info(f"📄 页面范围: {page_range}")
             
-            # 构建输出配置
-            output_config = {
-                'save_json': True,
-                'save_markdown': True,
-                'save_html': True,
-                'save_page_json': True,
-                'save_images': True,
-                'save_layout_image': debug,
-                'save_ocr_image': debug,
-                'normalize_numbers': True,
-                'merge_cross_page_tables': True,
-                'cleanup_temp_files': True,
-            }
-            
-            # 处理文档
             start_time = datetime.now()
             
             if streaming:
                 # 流式处理模式
-                results = pipeline.process_document_streaming(  # type: ignore
+                results = pipeline.process_document_streaming(
                     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,
+                    'output_paths': results.get('output_paths', {}),
                     'process_time': process_time
                 }
             else:
-                # 批量处理模式(原有逻辑)
-                # 批量处理模式(原有逻辑)
+                # 批量处理模式
                 results = pipeline.process_document(
                     str(input_path), 
                     page_range=page_range,
                     output_dir=str(output_dir)
                 )
                 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 {
@@ -195,9 +229,8 @@ def process_single_input(
                 }
         
         finally:
-            # 关闭context manager
-            if not streaming and use_context:
-                pipeline_batch.__exit__(None, None, None)
+            if use_context:
+                pipeline.__exit__(None, None, None)
             
     except Exception as e:
         logger.error(f"❌ 处理失败: {e}")
@@ -352,30 +385,17 @@ def main():
     
     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
+    if _handle_dry_run(args):
+        return 0
     
-    # 验证配置文件
-    config_path = Path(args.config)
-    if not config_path.exists():
-        logger.error(f"❌ 配置文件不存在: {config_path}")
+    if not _validate_arguments(args):
         return 1
     
-    # 仅验证模式
-    if args.dry_run:
-        logger.info("✅ 配置验证通过(dry run)")
-        return 0
-    
-    # 处理文档
     result = process_single_input(
-        input_path=input_path,
-        config_path=config_path,
+        input_path=Path(args.input),
+        config_path=Path(args.config),
         output_dir=Path(args.output_dir),
         debug=args.debug,
         scene=args.scene,
@@ -387,21 +407,12 @@ def main():
 
 
 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')}")
+    _print_environment_info()
 
     if len(sys.argv) == 1:
-        # 没有命令行参数时,使用默认配置运行
         print("ℹ️  未提供命令行参数,使用默认配置运行...")
         
-        # 默认配置
+        # 默认配置(用于开发测试)
         default_config = {
             # 测试输入
             # "input": "/Users/zhch158/workspace/data/流水分析/康强_北京农村商业银行.pdf",
@@ -414,15 +425,15 @@ if __name__ == "__main__":
             # "output_dir": "/Users/zhch158/workspace/data/流水分析/B用户_扫描流水/bank_statement_yusys_v2",
 
             # "input": "/Users/zhch158/workspace/repository.git/ocr_platform/ocr_tools/universal_doc_parser/tests/2023年度报告母公司_page_005.png",
-            # "input": "/Users/zhch158/workspace/repository.git/ocr_platform/ocr_tools/universal_doc_parser/tests/2023年度报告母公司_page_003_270.png",
+            # "input": "/Users/zhch158/workspace/repository.git/ocr_platform/ocr_tools/universal_doc_parser/tests/2023年度报告母公司_page_003.png",
             # "input": "/Users/zhch158/workspace/repository.git/ocr_platform/ocr_tools/universal_doc_parser/tests/2023年度报告母公司_page_003_270_skew(-0.4).png",
-            "input": "/Users/zhch158/workspace/data/流水分析/2023年度报告母公司.pdf",
-            "output_dir": "./output/2023年度报告母公司/bank_statement_wired_unet",
+            # "input": "/Users/zhch158/workspace/data/流水分析/2023年度报告母公司.pdf",
+            # "output_dir": "./output/2023年度报告母公司/bank_statement_wired_unet",
 
             # "input": "/Users/zhch158/workspace/data/流水分析/2023年度报告母公司.pdf",
             # "output_dir": "/Users/zhch158/workspace/data/流水分析/2023年度报告母公司/bank_statement_yusys_v2",
 
-            # "input": "/Users/zhch158/workspace/data/流水分析/A用户_单元格扫描流水.pdf",
+            # # "input": "/Users/zhch158/workspace/data/流水分析/A用户_单元格扫描流水.pdf",
             # "output_dir": "/Users/zhch158/workspace/data/流水分析/A用户_单元格扫描流水/bank_statement_wired_unet",
 
             # "input": "/Users/zhch158/workspace/repository.git/ocr_platform/ocr_tools/universal_doc_parser/tests/600916_中国黄金_2022年报_page_096.png",
@@ -430,6 +441,10 @@ if __name__ == "__main__":
             # "input": "/Users/zhch158/workspace/data/流水分析/600916_中国黄金_2022年报.pdf",
             # "output_dir": "./output/600916_中国黄金_2022年报/bank_statement_wired_unet",
 
+
+            # "input": "/Users/zhch158/workspace/repository.git/ocr_platform/ocr_tools/universal_doc_parser/tests/提取自赤峰黄金2023年报.pdf",
+            # "output_dir": "./output/提取自赤峰黄金2023年报/bank_statement_wired_unet",
+
             # "input": "/Users/zhch158/workspace/data/流水分析/施博深.pdf",
             # "output_dir": "/Users/zhch158/workspace/data/流水分析/施博深/bank_statement_yusys_v2",
 
@@ -439,6 +454,9 @@ if __name__ == "__main__":
             # "input": "/Users/zhch158/workspace/data/流水分析/施博深.wiredtable",
             # "output_dir": "/Users/zhch158/workspace/data/流水分析/施博深/bank_statement_wired_unet",
 
+            "input": "/Users/zhch158/workspace/data/流水分析/山西云集科技有限公司.pdf",
+            "output_dir": "/Users/zhch158/workspace/data/流水分析/山西云集科技有限公司/bank_statement_wired_unet",
+
             # 配置文件
             "config": "./config/bank_statement_wired_unet.yaml",
             # "config": "./config/bank_statement_yusys_v2.yaml",
@@ -448,7 +466,7 @@ if __name__ == "__main__":
             "scene": "bank_statement",
             
             # 页面范围(可选)
-            # "pages": "6",  # 只处理前1页
+            # "pages": "12,26,27",  # 只处理前1页
             # "pages": "1-3,5,7-10",  # 处理指定页面
             # "pages": "83-109",  # 处理指定页面
 
@@ -459,6 +477,9 @@ if __name__ == "__main__":
             
             # 日志级别
             "log_level": "DEBUG",
+
+            # 日志文件
+            "log_file": "./logs/bank_statement_wired_unet/process.log",
         }
         
         # 构造参数