import time
import os
from paddle import core
from paddleocr import PaddleOCRVL

IMAGE_PATH = "images/名片02.jpg"
WARMUP_ROUNDS = 0
BENCHMARK_ROUNDS = 1

# ── 线程配置 ──
# 可通过环境变量 PADDLE_THREADS 覆盖，否则使用逻辑核心数
DEFAULT_THREADS = int(os.environ.get("PADDLE_THREADS", os.cpu_count() or 4))
core.set_num_threads(DEFAULT_THREADS)
print(f"[Threads] oneDNN compiled, using {DEFAULT_THREADS} threads (CPU cores: {os.cpu_count()})")

# ── 模型初始化计时 ──
print("=" * 60)
print("初始化模型...")
t0 = time.perf_counter()
pipeline = PaddleOCRVL(pipeline_version="v1.6")
t_init = time.perf_counter() - t0
print(f"[OK] 模型初始化耗时: {t_init:.2f}s")
print("=" * 60)

# ── 推理 Benchmark ──
print(f"\n开始 OCR 识别: {IMAGE_PATH}")
print(f"预热 {WARMUP_ROUNDS} 轮 + 正式测试 {BENCHMARK_ROUNDS} 轮\n")

# 预热
for i in range(WARMUP_ROUNDS):
    print(f"  预热 {i + 1}/{WARMUP_ROUNDS}...")
    _ = pipeline.predict(IMAGE_PATH)

# 正式计时
times = []
for i in range(BENCHMARK_ROUNDS):
    print(f"  推理 {i + 1}/{BENCHMARK_ROUNDS}...", end=" ", flush=True)
    t0 = time.perf_counter()
    result = pipeline.predict(IMAGE_PATH)
    elapsed = time.perf_counter() - t0
    times.append(elapsed)
    print(f"{elapsed:.2f}s")

print("\n" + "=" * 60)
print("[Benchmark]")
print(f"  图片尺寸: {result[0]['width']} x {result[0]['height']}")
print(f"  检测文本块: {len(result[0]['layout_det_res']['boxes'])} 个")
print(f"  识别文本块: {len(result[0]['parsing_res_list'])} 个")
print(f"  推理次数: {BENCHMARK_ROUNDS}")
print(f"  最快: {min(times):.2f}s")
print(f"  最慢: {max(times):.2f}s")
print(f"  平均: {sum(times) / len(times):.2f}s")
if len(times) > 1:
    print(f"  标准差: {(sum((t - sum(times) / len(times)) ** 2 for t in times) / len(times)) ** 0.5:.2f}s")
print("=" * 60)

# ── 输出识别结果 ──
print("\n[识别结果]\n")
for item in result:
    for block in item["parsing_res_list"]:
        print(f"  [{block.label}] ({block.bbox})")
        print(f"  {block.content}\n")