215 lines
8.0 KiB
Python
215 lines
8.0 KiB
Python
"""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())
|