PaddleOCR-VL-1.6_Demo/main.py

"""CPU single-image OCR benchmark with structured timing logs."""

from __future__ import annotations

import argparse
import os
import statistics
import time
from pathlib import Path

from ocr_logging import default_log_path, setup_run_logger

PROJECT_ROOT = Path(__file__).resolve().parent
DEFAULT_IMAGE = PROJECT_ROOT / "images" / "名片02.jpg"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="PaddleOCR-VL-1.6 CPU 单图 Benchmark",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("image", nargs="?", type=Path, default=DEFAULT_IMAGE, help="输入图片")
    parser.add_argument("--threads", type=int, default=None, help="Paddle CPU 线程数")
    parser.add_argument("--warmup", type=int, default=0, help="预热轮数")
    parser.add_argument("--rounds", type=int, default=1, help="正式测试轮数")
    parser.add_argument("--log-file", type=Path, default=None, help="日志文件路径")
    parser.add_argument("--verbose", action="store_true", help="输出详细日志")
    parser.add_argument("--no-result", action="store_true", help="不输出识别文本")
    return parser.parse_args()


def main() -> int:
    program_started = time.perf_counter()
    args = parse_args()
    image_path = args.image.expanduser().resolve()
    log_file = args.log_file or default_log_path(PROJECT_ROOT, "single", image_path.stem, device="cpu")
    logger = setup_run_logger("ocr.single.cpu", log_file, verbose=args.verbose)

    if not image_path.is_file():
        logger.error("INPUT_NOT_FOUND path=%s", image_path)
        return 1
    if args.warmup < 0 or args.rounds < 1:
        logger.error("INVALID_ARGUMENT warmup=%d rounds=%d", args.warmup, args.rounds)
        return 2

    total_cores = os.cpu_count() or 4
    threads = args.threads or int(os.environ.get("PADDLE_THREADS", total_cores))
    if threads < 1:
        logger.error("INVALID_ARGUMENT threads=%d", threads)
        return 2

    logger.info(
        "PROGRAM_STARTED image=%s size_bytes=%d threads=%d total_cores=%d warmup=%d rounds=%d",
        image_path,
        image_path.stat().st_size,
        threads,
        total_cores,
        args.warmup,
        args.rounds,
    )

    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", import_seconds, 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", init_seconds)

    warmup_times: list[float] = []
    for index in range(args.warmup):
        started = time.perf_counter()
        pipeline.predict(str(image_path))
        elapsed = time.perf_counter() - started
        warmup_times.append(elapsed)
        logger.info("WARMUP_COMPLETED round=%d/%d seconds=%.3f", index + 1, args.warmup, elapsed)

    result = None
    inference_times: list[float] = []
    for index in range(args.rounds):
        started = time.perf_counter()
        result = pipeline.predict(str(image_path))
        elapsed = time.perf_counter() - started
        inference_times.append(elapsed)
        logger.info("INFERENCE_COMPLETED round=%d/%d seconds=%.3f", index + 1, args.rounds, elapsed)

    if not result:
        logger.error("EMPTY_RESULT")
        return 3

    first = result[0]
    layout_boxes = len(first["layout_det_res"]["boxes"])
    parsed_blocks = len(first["parsing_res_list"])
    non_empty_blocks = sum(bool(block.content.strip()) for block in first["parsing_res_list"])
    inference_mean = statistics.fmean(inference_times)
    inference_stdev = statistics.pstdev(inference_times)
    program_total = time.perf_counter() - program_started

    logger.info(
        "RESULT_SUMMARY width=%s height=%s layout_boxes=%d parsed_blocks=%d non_empty_blocks=%d",
        first.get("width"),
        first.get("height"),
        layout_boxes,
        parsed_blocks,
        non_empty_blocks,
    )
    logger.info(
        "BENCHMARK_SUMMARY model_init_seconds=%.3f warmup_total_seconds=%.3f inference_total_seconds=%.3f inference_min_seconds=%.3f inference_max_seconds=%.3f inference_mean_seconds=%.3f inference_median_seconds=%.3f inference_stdev_seconds=%.3f program_total_seconds=%.3f",
        init_seconds,
        sum(warmup_times),
        sum(inference_times),
        min(inference_times),
        max(inference_times),
        inference_mean,
        statistics.median(inference_times),
        inference_stdev,
        program_total,
    )

    if not args.no_result:
        for index, block in enumerate(first["parsing_res_list"], start=1):
            logger.info(
                "OCR_BLOCK index=%d label=%s bbox=%s content=%s",
                index,
                block.label,
                block.bbox,
                block.content.replace("\r", "").replace("\n", "\\n"),
            )

    logger.info("PROGRAM_COMPLETED log=%s", log_file.resolve())
    return 0


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