yolov5 onnx 推理代码

Posted 2021-07-04 洪流之源

onnx_yolov5_inference.py

import os
import cv2
import time
import onnxruntime as ort
import numpy as np
from common import pad_to_square
from common import save_tensor
from common import decode
from common import non_max_suppression
from common import draw_result

anchors = [[10,13], [16,30], [33,23],
           [30,61], [62,45], [59,119],
           [116,90],[156,198], [373,326]]

# 模型输入尺寸
model_input_h = 640
model_input_w = 640
# 目标置信度阈值
obj_conf_thresh = 0.3
# nms iou阈值
nms_iou_thresh = 0.5
# 类别数量
classes_num = 1
# 类别标签
class_label = (\'o\')

if __name__ == \'__main__\':
    image_name = \'B5D95V.jpg\'
    test_image_path = \'./data/test_images/{}\'.format(image_name)
    onnx_model_path = \'./models/onnx/best_20210531_sim.onnx\'
    tensor_output_dir = \'./data/output_tensor\'

    # 读取图片并转换到RGB图像格式,
    # 注意opencv图像维度顺序为：
    # shape[0]=h, sh