ocr_platform/ocr_tools/ocr_merger/merge_dotsocr_paddleocr.py

"""
合并 DotsOCR 和 PaddleOCR 的结果
主程序入口
"""
import json
import argparse
from pathlib import Path

import sys
from pathlib import Path

# 添加 ocr_platform 根目录到 Python 路径
ocr_platform_root = Path(__file__).parents[2]  # ocr_merger -> ocr_tools -> ocr_platform -> repository.git
if str(ocr_platform_root) not in sys.path:
    sys.path.insert(0, str(ocr_platform_root))

try:
    from ocr_tools.ocr_merger.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())