128 lines
4.7 KiB
Python
128 lines
4.7 KiB
Python
"""CPU entry point for page-by-page PaddleOCR-VL PDF recognition."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import os
|
||
import platform
|
||
import sys
|
||
import time
|
||
from pathlib import Path
|
||
|
||
from pdf_ocr_core import preflight_pdf, process_pdf
|
||
|
||
PROJECT_ROOT = Path(__file__).resolve().parent
|
||
DEFAULT_OUTPUT = PROJECT_ROOT / "outputs"
|
||
|
||
|
||
def parse_args() -> argparse.Namespace:
|
||
parser = argparse.ArgumentParser(
|
||
description="PaddleOCR-VL-1.6 CPU 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("--threads", type=int, default=None, help="Paddle CPU 线程数")
|
||
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 main() -> int:
|
||
args = parse_args()
|
||
total_cores = os.cpu_count() or 4
|
||
safe_default_threads = max(1, total_cores - 2)
|
||
threads = args.threads or int(os.environ.get("PADDLE_THREADS", safe_default_threads))
|
||
if threads < 1:
|
||
print("[ERROR] --threads 必须大于等于 1", file=sys.stderr)
|
||
return 2
|
||
|
||
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,
|
||
)
|
||
except Exception as exc:
|
||
print(f"[ERROR] {type(exc).__name__}: {exc}", file=sys.stderr)
|
||
return 1
|
||
|
||
print(
|
||
f"[PDF] {preflight['page_count']} pages, selected: "
|
||
f"{len(preflight['selected_pages'])}, output: {preflight['document_dir']}"
|
||
)
|
||
|
||
from paddle import core
|
||
from paddleocr import PaddleOCRVL
|
||
|
||
core.set_num_threads(threads)
|
||
print(f"[CPU] Threads: {threads} / {total_cores} (reserved: {max(0, total_cores - threads)})")
|
||
print("Loading PaddleOCR-VL-1.6...")
|
||
init_started = time.perf_counter()
|
||
pipeline = PaddleOCRVL(pipeline_version="v1.6", device="cpu")
|
||
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": "cpu",
|
||
"cpu_threads": threads,
|
||
"python_version": platform.python_version(),
|
||
"platform": platform.platform(),
|
||
"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,
|
||
)
|
||
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())
|