import cv2
import time
from pp_onnx.onnx_paddleocr import ONNXPaddleOcr, draw_ocr

# 优化参数配置，提高速度
model = ONNXPaddleOcr(
    use_angle_cls=True,
    use_gpu=False,
    providers=['AzureExecutionProvider'],
    provider_options=[{'device_id': 0}],
    det_limit_type='max',
)

def sav2Img(org_img, result, name="./result_img/3.jpg"):
    from PIL import Image
    result = result[0]
    image = org_img[:, :, ::-1]
    boxes = [line[0] for line in result]
    txts = [line[1][0] for line in result]
    scores = [line[1][1] for line in result]
    im_show = draw_ocr(image, boxes, txts, scores)
    im_show = Image.fromarray(im_show)
    im_show.save(name)


# 执行OCR推理
img = cv2.imread('./test_img/3.jpg')
if img is None:
    print(f"❌ 未找到图像文件")
else:
    # 图像预处理 - 缩小图像尺寸加速处理
    h, w = img.shape[:2]
    max_size = 1080
    if max(h, w) > max_size:
        scale = max_size / max(h, w)
        img = cv2.resize(img, (int(w * scale), int(h * scale)))

    s = time.time()
    result = model.ocr(img)
    e = time.time()
    print(f"total time: {e - s:.3f} 秒")
    print("result:", result)
    for box in result[0]:
        print(box)
    sav2Img(img, result)