mirror of
https://github.com/azaion/ai-training.git
synced 2026-04-22 21:56:36 +00:00
Oleksandr Bezdieniezhnykh
34 lines
919 B
Python
34 lines
919 B
Python
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
|