yolov8 onnx推理

前言：yolov8的后处理在某些情况下会导致转模型失败，因此需要把后处理剥离出来。

代码需要做如下修改：

改完后，网络会有三个输出，如图：

最后，用python写网络的后处理：

import onnxruntime
import numpy as np
import cv2
import os
import tqdm
# 指定 ONNX 模型文件路径
onnx_model_path = 'yolov8n.onnx'
sess = onnxruntime.InferenceSession(onnx_model_path)

coco_classes = [
    'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck',
    'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench',
    'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra',
    'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
    'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
    'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup',
    'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange',
    'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
    'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse',
    'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink',
    'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier',
    'toothbrush'
]
def softmax(x):
    exp_x = np.exp(x - np.max(x))  # 避免指数溢出
    return exp_x / exp_x.sum(axis=0, keepdims=True)

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

def get_color(idx):
    idx += 3
    return (37 * idx % 255, 17 * idx % 255, 29 * idx % 255)

def onnx_infer(img_path):
    # 创建 ONNX 运行时的 Session
    

    # 构造输入数据
    img = cv2.imread(img_path)
    img = cv2.resize(img,(640,640))
    input_data = np.transpose(img,(2,0,1))[None,...]
    input_data = np.array(input_data/255,dtype=np.float32)

    # 进行推理
    output_data = sess.run(['output0','357','358'], {'images': input_data})
       
    ans_bboxs = []
    ans_score = []
    ans_cat = []
    
    for output in output_data:
        _,h,w,c = output.shape
        pred = output.reshape(-1,c)
   
   
        for i in range(h*w):
            
            centery = i//w + 0.5
            centerx = i%w + 0.5
                   
            pix = pred[i]
            ld,td,rd,bd = [],[],[],[]
            for li in range(0,16):
                ld.append(pix[li])
                td.append(pix[li+16])
                rd.append(pix[li+32])
                bd.append(pix[li+48])

            ld = softmax(ld)
            td = softmax(td)
            rd = softmax(rd)
            bd = softmax(bd)
            
            ldis,tdis,rdis,bdis=0,0,0,0
            for j in range(16):
                ldis = ldis+j*ld[j]
                tdis = tdis+j*td[j]
                rdis = rdis+j*rd[j]
                bdis = bdis+j*bd[j]

            x0 = centerx - ldis
            y0 = centery - tdis
            x1 = centerx + rdis
            y1 = centery + bdis
            
            pred_cls = sigmoid(pix[64:])
            pred_cat = np.argmax(pred_cls)
            pred_score = pred_cls[pred_cat]
            if pred_score > 0.25:
                ans_bboxs.append([x0*640/h,y0*640/h,x1*640/h,y1*640/h])
                ans_score.append(pred_score)
                ans_cat.append(pred_cat)


        
    indices = cv2.dnn.NMSBoxes(ans_bboxs, ans_score,score_threshold=0.1, nms_threshold=0.55)

    ult_bbox = np.array(ans_bboxs)[indices]
    ult_score = np.array(ans_score)[indices]
    ult_cat = np.array(ans_cat)[indices]


    for bbox,conf,cat in zip(ult_bbox,ult_score,ult_cat):
        x0,y0,x1,y1 = bbox
        cv2.rectangle(img,(int(x0),int(y0)),(int(x1),int(y1)),get_color(int(cat)),2,2)
        label = f'{coco_classes[int(cat)]}: {conf:.2f}'
        cv2.putText(img, label, (int(x0),int(y0) - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, get_color(int(cat)), 2)
    cv2.imwrite(os.path.join('onnx_infer',img_path.split('/')[-1]),img)
    
    

if __name__ == '__main__':
    data_root = 'input/data/coco/imgs/'
    imgs = os.listdir(data_root)
    for img in tqdm.tqdm(imgs):
        if img.endswith('jpg') or img.endswith('png'):
            onnx_infer(os.path.join(data_root,img))

最后结果如图所示：