Clerk2.5/paddleocr_wrapper.py

#!/usr/bin/env python3
"""PaddleOCR VL Python API 包装脚本
用于独立进程调用 PaddleOCR，避免与主进程共享显存
"""

import os
import sys
import argparse
from pathlib import Path


def main():
    parser = argparse.ArgumentParser(description='PaddleOCR VL Python API 包装脚本')
    parser.add_argument('-i', '--input', required=True, help='输入图片路径')
    parser.add_argument('-o', '--output', required=True, help='输出目录路径')
    parser.add_argument('--mode', choices=['ocr', 'doc_parser'], default='doc_parser',
                        help='处理模式: ocr(纯文本) 或 doc_parser(图表识别)')
    parser.add_argument('--device', default='gpu:0', help='设备类型，如 gpu:0 或 cpu')
    parser.add_argument('--use-chart', action='store_true', help='启用图表识别')
    parser.add_argument('--use-layout', action='store_true', help='启用版面检测')
    
    args = parser.parse_args()
    
    # 检查输入文件
    if not os.path.exists(args.input):
        print(f"错误: 输入文件不存在: {args.input}", file=sys.stderr)
        sys.exit(1)
    
    # 创建输出目录
    output_dir = Path(args.output)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    try:
        from paddleocr import PaddleOCRVL
        
        # 根据模式设置参数
        if args.mode == 'ocr':
            # 纯文本识别模式
            use_chart = False
            use_layout = False
        else:
            # 图表识别模式
            use_chart = args.use_chart
            use_layout = args.use_layout
        
        print(f"初始化 PaddleOCR VL (mode={args.mode}, chart={use_chart}, layout={use_layout})...", file=sys.stderr)
        
        # 初始化 pipeline
        pipeline = PaddleOCRVL(
            device=args.device,
            use_doc_unwarping=False,
            use_doc_orientation_classify=True,
            use_chart_recognition=use_chart,
            use_layout_detection=use_layout,
        )
        
        print(f"开始处理: {args.input}", file=sys.stderr)
        
        # 执行识别
        result = pipeline.predict(input=args.input)
        
        print(f"识别完成，结果类型: {type(result)}", file=sys.stderr)
        print(f"结果数量: {len(result) if hasattr(result, '__len__') else 'N/A'}", file=sys.stderr)
        
        # 保存结果
        image_basename = os.path.splitext(os.path.basename(args.input))[0]
        
        for idx, item in enumerate(result):
            print(f"处理结果项 {idx}: {type(item)}", file=sys.stderr)
            print(f"  - 有 save_to_markdown 方法: {hasattr(item, 'save_to_markdown')}", file=sys.stderr)
            
            if hasattr(item, 'save_to_markdown'):
                print(f"  - 保存到: {output_dir}", file=sys.stderr)
                # 使用官方 API 保存 markdown
                item.save_to_markdown(save_path=str(output_dir))
                
                # 列出保存后的文件
                if output_dir.exists():
                    files = list(output_dir.iterdir())
                    print(f"  - 保存后的文件列表: {[f.name for f in files]}", file=sys.stderr)
                
                # 递归查找所有 .md 文件
                print(f"  - 递归查找 .md 文件...", file=sys.stderr)
                md_files = list(output_dir.rglob("*.md"))
                print(f"  - 找到 {len(md_files)} 个 .md 文件: {[f.name for f in md_files]}", file=sys.stderr)
                
                if md_files:
                    print(f"SUCCESS: {md_files[0]}")
                    sys.exit(0)
        
        print("错误: 未找到 markdown 输出文件", file=sys.stderr)
        print(f"输出目录: {output_dir}", file=sys.stderr)
        print(f"保存路径: {item_save_path if 'item_save_path' in locals() else 'N/A'}", file=sys.stderr)
        sys.exit(1)
        
    except ImportError as e:
        print(f"错误: 无法导入 PaddleOCR 模块: {e}", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"错误: {type(e).__name__}: {e}", file=sys.stderr)
        import traceback
        traceback.print_exc(file=sys.stderr)
        sys.exit(1)


if __name__ == '__main__':
    main()