"""GPU entry point for page-by-page PaddleOCR-VL PDF recognition.""" from __future__ import annotations import argparse import platform import sys import time from pathlib import Path GPU_DIR = Path(__file__).resolve().parent PROJECT_ROOT = GPU_DIR.parent if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) from ocr_logging import default_log_path, setup_run_logger from pdf_ocr_core import preflight_pdf, process_pdf DEFAULT_OUTPUT = PROJECT_ROOT / "outputs" def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="PaddleOCR-VL-1.6 GPU PDF OCR（逐页、可恢复）", formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) parser.add_argument("pdf", type=Path, help="输入 PDF 文件") parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT, help="输出根目录") parser.add_argument("--pages", help="一页或多个页码范围，例如 1-5,8,10-") parser.add_argument("--dpi", type=int, default=144, help="PDF 页面渲染 DPI") parser.add_argument("--password", help="加密 PDF 密码") parser.add_argument("--device-id", type=int, default=0, help="CUDA GPU 编号") parser.add_argument("--resume", action="store_true", help="跳过已完成页，继续现有任务") parser.add_argument("--overwrite", action="store_true", help="删除已有输出并重新处理") parser.add_argument("--keep-rendered", action="store_true", help="保留逐页渲染 PNG") parser.add_argument("--fail-fast", action="store_true", help="任一页失败后立即停止") parser.add_argument("--max-new-tokens", type=int, default=None, help="限制每个文本块最大生成 token") parser.add_argument("--min-pixels", type=int, default=None, help="VLM 最小输入像素参数") parser.add_argument("--max-pixels", type=int, default=None, help="VLM 最大输入像素参数") parser.add_argument("--log-file", type=Path, default=None, help="日志文件路径") parser.add_argument("--verbose", action="store_true", help="输出详细日志") return parser.parse_args() def configure_cuda(device_id: int): try: import paddle except ImportError as exc: raise RuntimeError("未安装 GPU 子项目依赖，请先运行 gpu/setup_env.py。") from exc if not paddle.is_compiled_with_cuda(): raise RuntimeError("当前 PaddlePaddle 不是 CUDA 构建；本程序不会回退到 CPU。") try: device_count = paddle.device.cuda.device_count() except Exception as exc: raise RuntimeError(f"无法查询 CUDA 设备: {exc}") from exc if device_count < 1: raise RuntimeError("未检测到 NVIDIA CUDA GPU；本程序不会回退到 CPU。") if device_id < 0 or device_id >= device_count: raise RuntimeError(f"GPU {device_id} 不存在，当前检测到 {device_count} 个设备。") device = f"gpu:{device_id}" paddle.set_device(device) paddle.device.cuda.synchronize(device_id) try: name = paddle.device.cuda.get_device_name(device_id) except Exception: name = "unknown" return paddle, device, name def main() -> int: program_started = time.perf_counter() args = parse_args() log_file = args.log_file or default_log_path( PROJECT_ROOT, "pdf", args.pdf.stem, device=f"gpu{args.device_id}", ) logger = setup_run_logger("ocr.pdf.gpu", log_file, verbose=args.verbose) logger.info( "PROGRAM_STARTED input=%s output=%s pages=%s dpi=%d device_id=%d resume=%s overwrite=%s keep_rendered=%s fail_fast=%s", args.pdf, args.output, args.pages or "all", args.dpi, args.device_id, args.resume, args.overwrite, args.keep_rendered, args.fail_fast, ) try: preflight_started = time.perf_counter() preflight = preflight_pdf( pdf_path=args.pdf, output_root=args.output, pages=args.pages, dpi=args.dpi, password=args.password, resume=args.resume, overwrite=args.overwrite, ) preflight_seconds = time.perf_counter() - preflight_started cuda_started = time.perf_counter() paddle, device, device_name = configure_cuda(args.device_id) cuda_setup_seconds = time.perf_counter() - cuda_started except Exception as exc: logger.error("PREFLIGHT_OR_CUDA_FAILED type=%s error=%s", type(exc).__name__, exc, exc_info=args.verbose) return 1 import_started = time.perf_counter() from paddleocr import PaddleOCRVL import_seconds = time.perf_counter() - import_started logger.info( "PREFLIGHT_COMPLETED seconds=%.3f page_count=%d selected_pages=%d document_dir=%s", preflight_seconds, preflight["page_count"], len(preflight["selected_pages"]), preflight["document_dir"], ) logger.info( "RUNTIME_READY cuda_setup_seconds=%.3f import_seconds=%.3f device=%s device_name=%s paddle_version=%s", cuda_setup_seconds, import_seconds, device, device_name, paddle.__version__, ) logger.info("MODEL_INITIALIZATION_STARTED pipeline_version=v1.6 device=%s", device) init_started = time.perf_counter() pipeline = PaddleOCRVL(pipeline_version="v1.6", device=device) paddle.device.cuda.synchronize(args.device_id) init_seconds = time.perf_counter() - init_started logger.info("MODEL_INITIALIZED seconds=%.3f pipeline_version=v1.6 device=%s", init_seconds, device) predict_kwargs = { key: value for key, value in { "max_new_tokens": args.max_new_tokens, "min_pixels": args.min_pixels, "max_pixels": args.max_pixels, }.items() if value is not None } metadata = { "device": device, "device_name": device_name, "python_version": platform.python_version(), "platform": platform.platform(), "paddle_version": paddle.__version__, "model_init_seconds": round(init_seconds, 3), "pipeline_version": "v1.6", "preflight_seconds": round(preflight_seconds, 3), "cuda_setup_seconds": round(cuda_setup_seconds, 3), "runtime_import_seconds": round(import_seconds, 3), "log_file": str(log_file.resolve()), } try: summary = process_pdf( pipeline=pipeline, pdf_path=args.pdf, output_root=args.output, pages=args.pages, dpi=args.dpi, password=args.password, resume=args.resume, overwrite=args.overwrite, keep_rendered=args.keep_rendered, fail_fast=args.fail_fast, run_metadata=metadata, predict_kwargs=predict_kwargs, synchronize=lambda: paddle.device.cuda.synchronize(args.device_id), logger=logger, ) except KeyboardInterrupt: logger.warning( "PROGRAM_INTERRUPTED total_seconds=%.3f resume_hint=--resume", time.perf_counter() - program_started, ) return 130 except Exception as exc: logger.exception( "PROGRAM_FAILED type=%s error=%s total_seconds=%.3f", type(exc).__name__, exc, time.perf_counter() - program_started, ) return 1 program_total = time.perf_counter() - program_started timing = summary.get("timing", {}) logger.info( "PROGRAM_COMPLETED status=%s completed_pages=%d selected_pages=%d failed_pages=%s model_init_seconds=%.3f pdf_task_seconds=%.3f program_total_seconds=%.3f output=%s log=%s", summary["status"], summary["completed_pages"], summary["selected_pages"], summary["failed_pages"], init_seconds, timing.get("task_total_seconds", 0.0), program_total, summary["document_dir"], log_file.resolve(), ) return 0 if not summary["failed_pages"] else 3 if __name__ == "__main__": raise SystemExit(main())