Forráskód Böngészése

feat: 添加 DotsOCR 和 PaddleOCR 合并程序,支持单文件和批量处理,输出为统一的MinerU格式

zhch158_admin 1 hete
szülő
commit
d451e66d4c
1 módosított fájl, 313 hozzáadás és 0 törlés
  1. 313 0
      merger/merge_dotsocr_paddleocr.py

+ 313 - 0
merger/merge_dotsocr_paddleocr.py

@@ -0,0 +1,313 @@
+"""
+合并 DotsOCR 和 PaddleOCR 的结果
+主程序入口
+"""
+import json
+import argparse
+from pathlib import Path
+
+try:
+    from .dotsocr_merger import DotsOCRMerger
+except ImportError:
+    from dotsocr_merger import DotsOCRMerger
+
+
+def merge_single_file(dotsocr_file: Path, paddle_file: Path, output_dir: Path, 
+                     output_type: str, merger: DotsOCRMerger) -> bool:
+    """
+    合并单个文件
+    
+    Args:
+        dotsocr_file: DotsOCR JSON 文件路径
+        paddle_file: PaddleOCR JSON 文件路径
+        output_dir: 输出目录
+        output_type: 输出格式
+        merger: 合并器实例
+    
+    Returns:
+        是否成功
+    """
+    print(f"📄 处理: {dotsocr_file.name}")
+    
+    # 输出文件路径
+    merged_md_path = output_dir / f"{dotsocr_file.stem}.md"
+    merged_json_path = output_dir / f"{dotsocr_file.stem}.json"
+    
+    try:
+        # ✅ 合并数据 (统一输出为MinerU格式)
+        merged_data = merger.merge_table_with_bbox(
+            str(dotsocr_file),
+            str(paddle_file),
+            data_format='mineru'  # 强制使用MinerU格式
+        )
+        
+        # ✅ 生成 Markdown (基于MinerU格式)
+        if output_type in ['markdown', 'both']:
+            markdown = merger.generate_enhanced_markdown(
+                merged_data, 
+                str(merged_md_path), 
+                str(dotsocr_file),
+                data_format='mineru'  # 强制使用MinerU格式
+            )
+        
+        # ✅ 保存 JSON (MinerU格式)
+        if output_type in ['json', 'both']:
+            with open(merged_json_path, 'w', encoding='utf-8') as f:
+                json.dump(merged_data, f, ensure_ascii=False, indent=2)
+
+        print(f"  ✅ 合并完成 (MinerU格式)")
+        print(f"  📊 共处理了 {len(merged_data)} 个对象")
+        print(f"  💾 输出文件:")
+        if output_type in ['markdown', 'both']:
+            print(f"    - {merged_md_path.name}")
+        if output_type in ['json', 'both']:
+            print(f"    - {merged_json_path.name}")
+
+        return True
+        
+    except Exception as e:
+        print(f"  ❌ 处理失败: {e}")
+        import traceback
+        traceback.print_exc()
+        return False
+
+
+def merge_dotsocr_batch(dotsocr_dir: str, paddle_dir: str, output_dir: str,
+                       output_type: str = 'both',
+                       look_ahead_window: int = 10, 
+                       similarity_threshold: int = 80):
+    """
+    批量合并 DotsOCR 和 PaddleOCR 的结果
+    
+    Args:
+        dotsocr_dir: DotsOCR 结果目录
+        paddle_dir: PaddleOCR 结果目录
+        output_dir: 输出目录
+        output_type: 输出格式
+        look_ahead_window: 向前查找窗口大小
+        similarity_threshold: 相似度阈值
+    """
+    dotsocr_path = Path(dotsocr_dir)
+    paddle_path = Path(paddle_dir)
+    output_path = Path(output_dir)
+    output_path.mkdir(parents=True, exist_ok=True)
+    
+    merger = DotsOCRMerger(look_ahead_window, similarity_threshold)
+    
+    # 查找所有 DotsOCR 的 JSON 文件
+    dotsocr_files = list(dotsocr_path.glob('*_page_*[0-9].json'))
+    dotsocr_files.sort()
+    
+    print(f"\n🔍 找到 {len(dotsocr_files)} 个 DotsOCR 文件")
+    print(f"📂 DotsOCR 目录: {dotsocr_dir}")
+    print(f"📂 PaddleOCR 目录: {paddle_dir}")
+    print(f"📂 输出目录: {output_dir}")
+    print(f"⚙️  查找窗口: {look_ahead_window}")
+    print(f"⚙️  相似度阈值: {similarity_threshold}%\n")
+    
+    success_count = 0
+    failed_count = 0
+    
+    for dotsocr_file in dotsocr_files:
+        # 查找对应的 PaddleOCR 文件
+        paddle_file = paddle_path / dotsocr_file.name
+        
+        if not paddle_file.exists():
+            print(f"⚠️  跳过: 未找到对应的 PaddleOCR 文件: {paddle_file.name}\n")
+            failed_count += 1
+            continue
+
+        if merge_single_file(dotsocr_file, paddle_file, output_path, output_type, merger):
+            success_count += 1
+        else:
+            failed_count += 1
+        
+        print()
+    
+    print("=" * 60)
+    print(f"✅ 处理完成!")
+    print(f"📊 统计信息:")
+    print(f"  - 总文件数: {len(dotsocr_files)}")
+    print(f"  - 成功: {success_count}")
+    print(f"  - 失败: {failed_count}")
+    print("=" * 60)
+
+
+def main():
+    """主函数"""
+    parser = argparse.ArgumentParser(
+        description='合并 DotsOCR 和 PaddleOCR 的识别结果,统一输出为MinerU格式',
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+示例用法:
+
+  1. 批量处理整个目录:
+     python merge_dotsocr_paddleocr.py \\
+         --dotsocr-dir /path/to/dotsocr/results \\
+         --paddle-dir /path/to/paddle/results \\
+         --output-dir /path/to/output
+
+  2. 处理单个文件:
+     python merge_dotsocr_paddleocr.py \\
+         --dotsocr-file /path/to/file_page_001.json \\
+         --paddle-file /path/to/file_page_001.json \\
+         --output-dir /path/to/output
+
+  3. 自定义参数:
+     python merge_dotsocr_paddleocr.py \\
+         --dotsocr-dir /path/to/dotsocr \\
+         --paddle-dir /path/to/paddle \\
+         --output-dir /path/to/output \\
+         --window 15 \\
+         --threshold 85
+        
+输出格式说明:
+  - JSON: 统一的MinerU格式JSON文件
+  - Markdown: 基于MinerU格式生成的Markdown文件
+        """
+    )
+    
+    # 文件/目录参数
+    file_group = parser.add_argument_group('文件参数')
+    file_group.add_argument(
+        '--dotsocr-file', 
+        type=str,
+        help='DotsOCR 输出的 JSON 文件路径(单文件模式)'
+    )
+    file_group.add_argument(
+        '--paddle-file', 
+        type=str,
+        help='PaddleOCR 输出的 JSON 文件路径(单文件模式)'
+    )
+    
+    dir_group = parser.add_argument_group('目录参数')
+    dir_group.add_argument(
+        '--dotsocr-dir', 
+        type=str,
+        help='DotsOCR 结果目录(批量模式)'
+    )
+    dir_group.add_argument(
+        '--paddle-dir', 
+        type=str,
+        help='PaddleOCR 结果目录(批量模式)'
+    )
+    
+    # 输出参数
+    output_group = parser.add_argument_group('输出参数')
+    output_group.add_argument(
+        '-o', '--output-dir',
+        type=str,
+        required=True,
+        help='输出目录(必需)'
+    )
+    output_group.add_argument(
+        '-f', '--output-type', 
+        choices=['json', 'markdown', 'both'], 
+        default='both', 
+        help='输出格式'
+    )
+
+    # 算法参数
+    algo_group = parser.add_argument_group('算法参数')
+    algo_group.add_argument(
+        '-w', '--window',
+        type=int,
+        default=15,
+        help='向前查找的窗口大小(默认: 15)'
+    )
+    algo_group.add_argument(
+        '-t', '--threshold',
+        type=int,
+        default=80,
+        help='文本相似度阈值(0-100,默认: 80)'
+    )
+    
+    args = parser.parse_args()
+    output_type = args.output_type.lower()
+    
+    # 验证参数
+    if args.dotsocr_file and args.paddle_file:
+        # 单文件模式
+        dotsocr_file = Path(args.dotsocr_file)
+        paddle_file = Path(args.paddle_file)
+        output_dir = Path(args.output_dir)
+        
+        if not dotsocr_file.exists():
+            print(f"❌ 错误: DotsOCR 文件不存在: {dotsocr_file}")
+            return
+        
+        if not paddle_file.exists():
+            print(f"❌ 错误: PaddleOCR 文件不存在: {paddle_file}")
+            return
+        
+        output_dir.mkdir(parents=True, exist_ok=True)
+        
+        print("\n🔧 单文件处理模式")
+        print(f"📄 DotsOCR 文件: {dotsocr_file}")
+        print(f"📄 PaddleOCR 文件: {paddle_file}")
+        print(f"📂 输出目录: {output_dir}\n")
+        
+        merger = DotsOCRMerger(
+            look_ahead_window=args.window,
+            similarity_threshold=args.threshold
+        )
+        
+        success = merge_single_file(dotsocr_file, paddle_file, output_dir, output_type, merger)
+        
+        if success:
+            print("\n✅ 处理完成!")
+        else:
+            print("\n❌ 处理失败!")
+    
+    elif args.dotsocr_dir and args.paddle_dir:
+        # 批量模式
+        if not Path(args.dotsocr_dir).exists():
+            print(f"❌ 错误: DotsOCR 目录不存在: {args.dotsocr_dir}")
+            return
+        
+        if not Path(args.paddle_dir).exists():
+            print(f"❌ 错误: PaddleOCR 目录不存在: {args.paddle_dir}")
+            return
+        
+        print("\n🔧 批量处理模式")
+        
+        merge_dotsocr_batch(
+            args.dotsocr_dir,
+            args.paddle_dir,
+            args.output_dir,
+            output_type=output_type,
+            look_ahead_window=args.window,
+            similarity_threshold=args.threshold
+        )
+    
+    else:
+        parser.print_help()
+        print("\n❌ 错误: 请指定单文件模式或批量模式的参数")
+
+
+if __name__ == "__main__":
+    print("🚀 启动 DotsOCR + PaddleOCR 合并程序 (统一输出MinerU格式)...")
+    
+    import sys
+    
+    if len(sys.argv) == 1:
+        # 默认配置
+        default_config = {
+            "dotsocr-file": "/Users/zhch158/workspace/data/流水分析/A用户_单元格扫描流水/dotsocr_vllm_results/A用户_单元格扫描流水_page_002.json",
+            "paddle-file": "/Users/zhch158/workspace/data/流水分析/A用户_单元格扫描流水/ppstructurev3_client_results/A用户_单元格扫描流水_page_002.json",
+            "output-dir": "/Users/zhch158/workspace/data/流水分析/A用户_单元格扫描流水/dotsocr_vllm_results_cell_bbox",
+            "output-type": "both",
+            "window": "15",
+            "threshold": "85"
+        }
+        
+        print("ℹ️  未提供命令行参数,使用默认配置运行...")
+        print("⚙️  默认参数:")
+        for key, value in default_config.items():
+            print(f"  --{key}: {value}")
+        
+        sys.argv = [sys.argv[0]]
+        for key, value in default_config.items():
+            sys.argv.extend([f"--{key}", str(value)])
+    
+    sys.exit(main())