183 lines
6.7 KiB
Python
183 lines
6.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 time
|
||
from pathlib import Path
|
||
|
||
from ocr_logging import default_log_path, setup_run_logger
|
||
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 最大输入像素参数")
|
||
parser.add_argument("--log-file", type=Path, default=None, help="日志文件路径")
|
||
parser.add_argument("--verbose", action="store_true", help="输出详细日志")
|
||
return parser.parse_args()
|
||
|
||
|
||
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="cpu",
|
||
)
|
||
logger = setup_run_logger("ocr.pdf.cpu", log_file, verbose=args.verbose)
|
||
logger.info(
|
||
"PROGRAM_STARTED input=%s output=%s pages=%s dpi=%d resume=%s overwrite=%s keep_rendered=%s fail_fast=%s",
|
||
args.pdf,
|
||
args.output,
|
||
args.pages or "all",
|
||
args.dpi,
|
||
args.resume,
|
||
args.overwrite,
|
||
args.keep_rendered,
|
||
args.fail_fast,
|
||
)
|
||
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:
|
||
logger.error("INVALID_ARGUMENT threads=%d", threads)
|
||
return 2
|
||
|
||
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
|
||
except Exception as exc:
|
||
logger.error("PREFLIGHT_FAILED type=%s error=%s", type(exc).__name__, exc, exc_info=args.verbose)
|
||
return 1
|
||
|
||
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"],
|
||
)
|
||
|
||
import_started = time.perf_counter()
|
||
from paddle import core
|
||
from paddleocr import PaddleOCRVL
|
||
import_seconds = time.perf_counter() - import_started
|
||
|
||
core.set_num_threads(threads)
|
||
logger.info(
|
||
"RUNTIME_READY import_seconds=%.3f threads=%d total_cores=%d reserved_cores=%d",
|
||
import_seconds,
|
||
threads,
|
||
total_cores,
|
||
max(0, total_cores - threads),
|
||
)
|
||
logger.info("MODEL_INITIALIZATION_STARTED pipeline_version=v1.6 device=cpu")
|
||
init_started = time.perf_counter()
|
||
pipeline = PaddleOCRVL(pipeline_version="v1.6", device="cpu")
|
||
init_seconds = time.perf_counter() - init_started
|
||
logger.info("MODEL_INITIALIZED seconds=%.3f pipeline_version=v1.6 device=cpu", init_seconds)
|
||
|
||
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",
|
||
"preflight_seconds": round(preflight_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,
|
||
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())
|