guoju0808/ry-ocr/tests/bench.py

"""基准测试：模型预热后，统计 PNG / PDF 单次识别耗时"""
import sys
import time
from pathlib import Path

try:
    sys.stdout.reconfigure(encoding="utf-8")
except Exception:
    pass

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

from app.core import warmup
from app.services.recognize_service import recognize_file


def bench(label: str, file_path: Path, runs: int = 5):
    content = file_path.read_bytes()
    print(f"\n=== {label}: {file_path.name} ({len(content)} bytes) ===")

    # 预热（不计耗时） - 现在 warmup 会真跑一次 dummy 图预测
    print("预热中（不计入）...", flush=True)
    warmup()
    print("预热完成", flush=True)

    times = []
    for i in range(runs):
        t0 = time.time()
        result = recognize_file(file_path.name, content)
        elapsed = (time.time() - t0) * 1000
        times.append(elapsed)
        print(f"  第 {i+1} 次: {elapsed:7.1f} ms  success={result.success}")

    times.sort()
    print(f"\n  min:    {times[0]:7.1f} ms")
    print(f"  median: {times[len(times)//2]:7.1f} ms")
    print(f"  max:    {times[-1]:7.1f} ms")
    print(f"  avg:    {sum(times)/len(times):7.1f} ms")


if __name__ == "__main__":
    bench("PNG", ROOT / "fapiao_1.png", runs=5)
    bench("PDF", ROOT / "fapiao.pdf", runs=5)