Эх сурвалжийг харах

refactor: 使用独立脚本调用 PaddleOCR,避免显存共享问题

- 创建 paddleocr_wrapper.py 独立脚本
- 支持 ocr 和 doc_parser 两种模式
- 支持 --use-chart 和 --use-layout 参数
- 修改 _call_paddleocr_api 通过 subprocess 调用独立脚本
- 超时时间增加到 120 秒
- 独立进程运行,避免与主进程共享显存
何文松 1 өдөр өмнө
parent
commit
8c60c709d6

+ 90 - 0
paddleocr_wrapper.py

@@ -0,0 +1,90 @@
+#!/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()

+ 52 - 41
pdf_converter_v2/utils/paddleocr_fallback.py

@@ -215,7 +215,7 @@ def _call_paddleocr_api(
     use_chart_recognition: bool = False,
     use_layout_detection: bool = False
 ) -> tuple[bool, Optional[str]]:
-    """使用 Python API 调用 PaddleOCR VL
+    """通过独立脚本调用 PaddleOCR VL(避免显存共享问题)
     
     Args:
         image_path: 输入图片路径
@@ -226,61 +226,72 @@ def _call_paddleocr_api(
     Returns:
         (是否成功, markdown 文件路径)
     """
-    import signal
-    
-    class TimeoutError(Exception):
-        pass
-    
-    def timeout_handler(signum, frame):
-        raise TimeoutError("PaddleOCR predict 超时")
-    
     try:
         if not os.path.exists(image_path):
-            logger.error(f"[PaddleOCR API] 图片文件不存在: {image_path}")
+            logger.error(f"[PaddleOCR Wrapper] 图片文件不存在: {image_path}")
             return False, None
         
         os.makedirs(save_path, exist_ok=True)
         
-        # 获取 pipeline
-        pipeline = _get_paddleocr_pipeline(use_chart_recognition, use_layout_detection)
+        # 获取 wrapper 脚本路径
+        current_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+        wrapper_script = os.path.join(current_dir, "paddleocr_wrapper.py")
         
-        logger.info(f"[PaddleOCR API] 开始处理: {image_path}")
-        logger.info(f"[PaddleOCR API] Pipeline 类型: {type(pipeline)}")
+        if not os.path.exists(wrapper_script):
+            logger.error(f"[PaddleOCR Wrapper] 脚本不存在: {wrapper_script}")
+            return False, None
         
-        # 执行识别(设置 60 秒超时)
-        logger.info(f"[PaddleOCR API] 调用 predict 方法...")
-        signal.signal(signal.SIGALRM, timeout_handler)
-        signal.alarm(60)  # 60 秒超时
+        # 获取 Python 解释器路径(使用 PaddleOCR 虚拟环境)
+        python_executable = "/mnt/win_d/paddle/.venv_paddleocr/bin/python"
+        if not os.path.exists(python_executable):
+            python_executable = sys.executable
+        
+        # 构建命令
+        mode = 'doc_parser' if (use_chart_recognition or use_layout_detection) else 'ocr'
+        cmd = [python_executable, wrapper_script, '-i', image_path, '-o', save_path, '--mode', mode]
+        
+        if use_chart_recognition:
+            cmd.append('--use-chart')
+        if use_layout_detection:
+            cmd.append('--use-layout')
+        
+        # 获取设备
+        device = _get_paddle_ocr_device()
+        cmd.extend(['--device', device])
+        
+        logger.info(f"[PaddleOCR Wrapper] 执行命令: {' '.join(cmd)}")
+        
+        # 执行命令
+        result = subprocess.run(
+            cmd,
+            capture_output=True,
+            text=True,
+            timeout=120,  # 2分钟超时
+            check=False,
+        )
         
-        try:
-            result = pipeline.predict(input=image_path)
-            signal.alarm(0)  # 取消超时
-            logger.info(f"[PaddleOCR API] predict 返回,结果类型: {type(result)}")
-        except TimeoutError:
-            signal.alarm(0)
-            logger.error(f"[PaddleOCR API] predict 超时(60秒),可能显存不足或模型加载失败")
+        if result.returncode != 0:
+            logger.error(f"[PaddleOCR Wrapper] 执行失败,返回码: {result.returncode}")
+            if result.stderr:
+                logger.error(f"[PaddleOCR Wrapper] 错误输出: {result.stderr}")
             return False, None
         
-        # 保存结果
-        image_basename = os.path.splitext(os.path.basename(image_path))[0]
+        # 从 stdout 解析 markdown 文件路径
+        output = result.stdout.strip()
+        if output.startswith("SUCCESS: "):
+            markdown_file = output.replace("SUCCESS: ", "")
+            if os.path.exists(markdown_file):
+                logger.info(f"[PaddleOCR Wrapper] 成功,markdown 文件: {markdown_file}")
+                return True, markdown_file
         
-        logger.info(f"[PaddleOCR API] 开始处理结果...")
-        for item in result:
-            if hasattr(item, 'save'):
-                item_save_path = os.path.join(save_path, image_basename)
-                item.save(item_save_path)
-                logger.info(f"[PaddleOCR API] 结果已保存到: {item_save_path}")
-                
-                # 查找生成的 markdown 文件
-                markdown_file = os.path.join(item_save_path, f"{image_basename}.md")
-                if os.path.exists(markdown_file):
-                    return True, markdown_file
-        
-        logger.warning(f"[PaddleOCR API] 未找到 markdown 输出文件")
+        logger.error(f"[PaddleOCR Wrapper] 未找到 markdown 输出文件")
         return False, None
         
+    except subprocess.TimeoutExpired:
+        logger.error("[PaddleOCR Wrapper] 执行超时(120秒)")
+        return False, None
     except Exception as e:
-        logger.error(f"[PaddleOCR API] 处理失败: {e}")
+        logger.error(f"[PaddleOCR Wrapper] 处理失败: {e}")
         import traceback
         logger.error(traceback.format_exc())
         return False, None