EADST

YOLOv5: Infer the Image with ONNX

Convert to the onnx model from our well-trained pt model

python export.py --weights ./runs/train/exp/best.pt --include onnx

Infer the image with the onnx model

import os
import cv2
import numpy as np
import onnxruntime
import time

CLASSES=["b", "t"]
class YOLOV5():
    def init(self,onnxpath):
        self.onnx_session=onnxruntime.InferenceSession(onnxpath)
        self.input_name=self.get_input_name()
        self.output_name=self.get_output_name()
def get_input_name(self):
    input_name=[]
    for node in self.onnx_session.get_inputs():
        input_name.append(node.name)
    return input_name

def get_output_name(self):
    output_name=[]
    for node in self.onnx_session.get_outputs():
        output_name.append(node.name)
    return output_name

def get_input_feed(self,img_tensor):
    input_feed={}
    for name in self.input_name:
        input_feed[name]=img_tensor
    return input_feed

def inference(self,img_path):
    img=cv2.imread(img_path)
    or_img=cv2.resize(img,(640,640))
    img=or_img[:,:,::-1].transpose(2,0,1)  #BGR2RGB and HWC2CHW
    img=img.astype(dtype=np.float32)
    img/=255.0
    img=np.expand_dims(img,axis=0)
    input_feed=self.get_input_feed(img)
    pred=self.onnx_session.run(None,input_feed)[0]
    return pred,or_img

def nms(dets, thresh):
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
areas = (y2 - y1 + 1) * (x2 - x1 + 1)
scores = dets[:, 4]
keep = []
index = scores.argsort()[::-1]

while index.size > 0:
    i = index[0]
    keep.append(i)

    x11 = np.maximum(x1[i], x1[index[1:]]) 
    y11 = np.maximum(y1[i], y1[index[1:]])
    x22 = np.minimum(x2[i], x2[index[1:]])
    y22 = np.minimum(y2[i], y2[index[1:]])

    w = np.maximum(0, x22 - x11 + 1)                              
    h = np.maximum(0, y22 - y11 + 1)

    overlaps = w * h

    ious = overlaps / (areas[i] + areas[index[1:]] - overlaps)
    idx = np.where(ious <= thresh)[0]
    index = index[idx + 1]
return keep

def xywh2xyxy(x):
    # [x, y, w, h] to [x1, y1, x2, y2]
    y = np.copy(x)
    y[:, 0] = x[:, 0] - x[:, 2] / 2
    y[:, 1] = x[:, 1] - x[:, 3] / 2
    y[:, 2] = x[:, 0] + x[:, 2] / 2
    y[:, 3] = x[:, 1] + x[:, 3] / 2
    return y
def filter_box(org_box,conf_thres,iou_thres): 
    org_box=np.squeeze(org_box)
    conf = org_box[..., 4] > conf_thres
    box = org_box[conf == True]
cls_cinf = box[..., 5:]
cls = []
for i in range(len(cls_cinf)):
    cls.append(int(np.argmax(cls_cinf[i])))
all_cls = list(set(cls))

output = []
for i in range(len(all_cls)):
    curr_cls = all_cls[i]
    curr_cls_box = []
    curr_out_box = []
    for j in range(len(cls)):
        if cls[j] == curr_cls:
            box[j][5] = curr_cls
            curr_cls_box.append(box[j][:6])
    curr_cls_box = np.array(curr_cls_box)
    # curr_cls_box_old = np.copy(curr_cls_box)
    curr_cls_box = xywh2xyxy(curr_cls_box)
    curr_out_box = nms(curr_cls_box,iou_thres)
    for k in curr_out_box:
        output.append(curr_cls_box[k])
output = np.array(output)
return output

def draw(image,box_data):  
boxes=box_data[...,:4].astype(np.int32) 
scores=box_data[...,4]
classes=box_data[...,5].astype(np.int32)

for box, score, cl in zip(boxes, scores, classes):
    top, left, right, bottom = box
    print('class: {}, score: {}'.format(CLASSES[cl], score))
    print('box left,top,right,down: [{}, {}, {}, {}]'.format(top, left, right, bottom))

    text_left = left
    text_top = top         
    color = (255, 0, 0)
    if CLASSES[cl] == "b":
        text_left += 50
        text_top -= 20     
        color = (0, 0, 255)
    cv2.rectangle(image, (top, left), (right, bottom), color, 2)
    cv2.putText(image, CLASSES[cl],
                (text_top, text_left),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.6, color, 2)

if name=="main":
    onnx_path = "best.onnx"
    model = YOLOV5(onnx_path)
    img_path = "0_0.jpg"
    output, img = model.inference(img_path)
    outbox = filter_box(output,0.5,0.5)
    draw(img, outbox)
    cv2.imwrite(img_path[:-4]+"_res.jpg", img)

Reference:

YOLOV5模型转onnx并推理