前言: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))
最后结果如图所示: