验证onnx格式的YOLO

刀么克瑟拉莫

于 2023-07-04 15:45:54 发布

阅读量263

点赞数

分类专栏： deeplearning 文章标签： onnx

本文链接：https://blog.youkuaiyun.com/random_repick/article/details/131536419

版权

该代码段展示了一个使用ONNXRuntime进行目标检测模型推理的过程。它涉及读取图像，调整大小，通过ONNX模型运行推理，获取边界框坐标，应用非极大值抑制(NMS)算法，并用OpenCV在图像上绘制检测到的框。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

相关参数改成自己的
注意输出的shape，onnx有时候会合并输出，目前不知道什么原因

import onnxruntime   
import cv2, torch
import numpy as np

anchors=[[22, 55,  38, 44,  32, 88,  70, 56], [22, 22,  17, 35,  37, 24,  27, 34]]
ans=len(anchors[0])//2
cls_num=1
conf_th=0.25
nms_th=0.45
model_w=
model_h=

def get_boxes(output, anchors):
	h=output.size(2)
	w=output.size(3)
	output=output.view(ans,int(cls_num+5),h,w).permute(0,2,3,1).contiguous()
	# conf
	conf = torch.sigmoid(output[..., 4])
	cl = torch.sigmoid(output[..., 5:])
	#conf = output[..., 4]
	#cl = output[..., 5:]
	clv, cli = torch.max(cl, -1)
	#print(conf[conf>0.5])
	conf = conf * clv
	#print(conf[conf>0.15])
	mask = conf > conf_th
	conf = conf[mask].unsqueeze(-1)
	cli = cli[mask].unsqueeze(-1)
	# grid
	FloatTensor = torch.cuda.FloatTensor if conf.is_cuda else torch.FloatTensor
	grid_h, grid_w = torch.meshgrid(torch.arange(h), torch.arange(w))
	grid_h = grid_h.repeat(ans,1,1).type(FloatTensor)
	grid_w = grid_w.repeat(ans,1,1).type(FloatTensor