PaddleOCR-VL-1.6_Demo/gpu/pdf_ocr.py

"""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 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 最大输入像素参数")
    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:
    args = parse_args()
    try:
        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,
        )
        paddle, device, device_name = configure_cuda(args.device_id)
    except Exception as exc:
        print(f"[ERROR] {type(exc).__name__}: {exc}", file=sys.stderr)
        return 1

    from paddleocr import PaddleOCRVL

    print(
        f"[PDF] {preflight['page_count']} pages, selected: "
        f"{len(preflight['selected_pages'])}, output: {preflight['document_dir']}"
    )
    print(f"[GPU] Device: {device} ({device_name})")
    print("Loading PaddleOCR-VL-1.6...")
    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
    print(f"Model init: {init_seconds:.2f}s")

    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",
    }

    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),
        )
    except KeyboardInterrupt:
        print("\n[INTERRUPTED] 已保存当前进度；使用 --resume 继续。", file=sys.stderr)
        return 130
    except Exception as exc:
        print(f"[ERROR] {type(exc).__name__}: {exc}", file=sys.stderr)
        return 1

    print("\n[PDF OCR Summary]")
    print(f"Status:    {summary['status']}")
    print(f"Completed: {summary['completed_pages']} / {summary['selected_pages']}")
    print(f"Failed:    {summary['failed_pages']}")
    print(f"Output:    {summary['document_dir']}")
    return 0 if not summary["failed_pages"] else 3


if __name__ == "__main__":
    raise SystemExit(main())