代码收藏家技术教程 2022-07-31

关于yolov5后处理细节

近日做项目遇到yolov5后处理操作，有点记不太清楚，所以想记录下来，方便后期复习。

我们都知道，yolov5模型的输入三个feature map的集合，加上batch的维度，也就是三维张量，即[batch，(p3*p3+p4*p4+p5*p5)*3，4+conf+cls_num],模型输出的为相对于调整图片的xywh，然后后面就要进入后处理阶段。

后处理采用了两次阈值筛选，然后送入了nms处理。

两次筛选都用的是conf_thres,只有nms求IOU的时候采用的iou_thresh.

首先，通过tensor的性质提取classes numbers，筛选出大于confience socre的box，得到的xc为bool，并采用assert进行阈值界限设置。

第一个for循环于图片数量的单张预测，即提取张量[(p3*p3+p4*p4+p5*p5)*3,5+cls_nums],取出大于阈值的box，如果有labels将labels加入到筛选的box后边，如果没有box，就继续下一张图片。

然后计算class confidence。

这段代码开始是将xywh转化为xyxy(相对于resize图像)，后面是mutil labels和single label的处理。

如果是mutil labels, 筛选出大于阈值的box，nonzero()得到位置索引[[i-n,j-n],[i,j],[i+n,j+n]…]，利用转置得到i,j，分别表示第i个框，第j个类别。最后x保留为6维，分别是，xyxy，conf，cls，其中的conf 为类别置信度。single label 只保留类别置信度最大的置信度就可以了。

参数 classes为需要筛选掉的类别，不得不说大神对于pytorch的熟练，一行代码全解决。