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)
        
        # 保存结果
        image_basename = os.path.splitext(os.path.basename(args.input))[0]
        
        for item in result:
            if hasattr(item, 'save'):
                item_save_path = output_dir / image_basename
                item.save(str(item_save_path))
                
                # 查找生成的 markdown 文件
                markdown_file = item_save_path / f"{image_basename}.md"
                if markdown_file.exists():
                    print(f"SUCCESS: {markdown_file}")
                    sys.exit(0)
        
        print("错误: 未找到 markdown 输出文件", 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()