import onnxruntime as ort
import torch
import numpy as np
import cv2
from tqdm import tqdm
import time

weight = r"C:\Users\user\PycharmProjects\yolov5_v6.0\runs\train\exp10\weights\yolov5s_640x640_opt.onnx"

input_img = cv2.resize(cv2.imread(r"D:\Deep_Learning\yolov5\whitelines\train\images\707.png"), (640, 640))
input_img = cv2.cvtColor(input_img, cv2.COLOR_BGR2RGB).transpose(2, 0, 1)[np.newaxis, :].astype(np.float32)
input_img /= 255.0

ort_session = ort.InferenceSession(weight)
t1 = time.time()
for _ in tqdm(range(1000)):
    pred = ort_session.run(
        None,
        {"images": input_img}
    )
print(time.time() - t1)