import time

import cv2
import numpy as np
import pytest

from inference.onnx_engine import OnnxEngine


@pytest.fixture(scope="session")
def onnx_engine_session(fixture_onnx_model):
    return OnnxEngine(fixture_onnx_model)


@pytest.mark.performance
def test_pt_inf_01_single_image_onnx_latency(onnx_engine_session, fixture_images_dir):
    engine = onnx_engine_session
    imgs = sorted(fixture_images_dir.glob("*.jpg"))
    assert imgs
    frame = cv2.imread(str(imgs[0]))
    assert frame is not None
    model_height, model_width = engine.get_input_shape()
    n = engine.get_batch_size()
    frames = [frame] * n
    blobs = [
        cv2.dnn.blobFromImage(f, 1.0 / 255.0, (model_width, model_height), (0, 0, 0), swapRB=True, crop=False)
        for f in frames
    ]
    blob = np.vstack(blobs)
    t0 = time.perf_counter()
    engine.run(blob)
    elapsed = time.perf_counter() - t0
    assert elapsed <= 10.0