位置: IT常识 - 正文
YOLOv5图像分割中的NMS处理(yolov3图像识别)
编辑：rootadmin推荐整理分享YOLOv5图像分割中的NMS处理(yolov3图像识别)，希望有所帮助，仅作参考，欢迎阅读内容。
文章相关热门搜索词:yolo 图像分割,yolov5语义分割,yolov5 图像大小,yolov5 图像大小,yolov4图像分割,yolov3图像分割,yolov4图像分割,yolo 图像分割,内容如对您有帮助，希望把文章链接给更多的朋友！
在上一篇文章YOLOv5图像分割--SegmentationModel类代码详解有讲到图像经过YOLOv5网络后得到的输出形式，主要是调用了BaseModel类下的forward得到的输出，输出的shape为【batch,25200,117】，这里的25200相当于总的anchors数量【以640*640的输入为例，共有anchors=80*80*3+40*40*3+20*20*3】，117为5[x,y,w,h,conf]+80个类+32【mask的数量】。
那么得到上面这张图的输出后又需要哪些处理呢？又是怎么处理的呢？本篇文章就是来刨析这个问题。
可以从下面的代码看到在进行model后会得到pred和proto。前者就是上面得到图的形式，后者的shape为【batch,32,160,160】，这里的32是mask的数量，160*160是针对80*80这个特征层的上采样得到的。然后是送入NMS进行处理得到新的pred输出。
pred, proto = model(im, augment=augment, visualize=visualize)[:2] # im[batch, 3, 640, 640] # pred:[batch,25200,117], proto:[batch,32,160,160] # NMS with dt[2]: pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det, nm=32)目录
NMS中发生了什么？
通过conf_thres进行筛选
xywh2xyxy
取得对应anchor中conf最大的类
NMS处理
NMS中发生了什么？
那么来看一下NMS中发生了什么
可以看到代码中传入的参数为：pred:model的输出，conf_thres：置信度阈值，iou_thres：iou阈值，classes:需要过滤的类，nm为mask的数量。
然后进入内部，prediction的shepe为【batch,25200,117】，因此通过下面的代码得到：
bs:batch size 这里得到是1
nc:117-32-5=80【coco类的数量】
xc:由于prediction在117这个维度的第5维度【这里的4维度】是指conf，所以可以通过conf_thres得到大于阈值的目标mask xc，那么此时xc的shape为【1，25200】。通过这个步骤就可以将所有anchors中大于conf的目标筛选出来【形式为False or True】。
通过conf_thres进行筛选 bs = prediction.shape[0] # batch size nc = prediction.shape[2] - nm - 5 # number of classes 85-5 xc = prediction[..., 4] > conf_thres # candidates 85=5+80=(x,y,w,h,conf) 第4维度为conf然后新建一个output全零的张量，最后一个通道的shape为6+32=38。6指的【x1,y1,x2,y1,conf,class】。
output = [torch.zeros((0, 6 + nm), device=prediction.device)] * bs
下面的xi是图像的index，对应该index的图像。
x = x[xc[xi]]：由于prediction当前batch为1 ，所以对应xi此时为0，xc是上面通过conf_thres得到shape为【1，25200】的mask，所以xc[xi]就是取出该batch中所有的anchors【这些anchors内已经内是经过conf筛选的目标】；x的shape为【25200，117】[忘记这个shape的含义可以看我最前面的图]，那么x[xc[ix]]就可以表示为通过xc中的25200个含有False or True 的anchors，筛选出conf大于阈值的目标【这个目标维度为117，含有box信息，conf,80个类，32个mask】
在我这里得到新的x的shape为：【52，117】。这就表示了在25200个anchors中有52个anchors内有目标，每个anchor又有117个维度来记录该anchor内目标的boxes信息，conf以及类信息等。
for xi, x in enumerate(prediction): # image index, image inference # Apply constraints # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # width-height x = x[xc[xi]] # confidence xc是所有anchor内的置信度，shape[1,25200]。表示为：获取当前图像中所有anchor中有预测结果的[这个是通过conf筛选过的]下面的这一行代码表示为计算conf，x[:,5:]
# Compute conf x[:, 5:] *= x[:, 4:5] # conf = obj_conf * cls_confxywh2xyxy下面的代码中box是将上面x中的box信息由center_x,center_y,wh转为x1,y1,x2,y2的形式【为后面计算iou做准备】，这里的mi=85【表示mask start index】，因为此时的xshape为【52，117=5+85+32】所以从85开始提取所有的mask。
# Box/Mask box = xywh2xyxy(x[:, :4]) # center_x, center_y, width, height) to (x1, y1, x2, y2) mask = x[:, mi:] # zero columns if no masks取得对应anchor中conf最大的类在下面的代码中x[:,5:mi]=x[:,5:85]表示为取所有anchors中所有类的conf。【0~4是box，5:85是class】，然后取出所有anchors中预测的类中conf最大的conf以及索引j【x[:,5:85]的shape为[52,80]，利用max(1,keepdim=True)在80所在的这个维度上取max】，而这里得到最大index j就是每个anchors预测得到的类的index[这样不就知道预测的种类了嘛]。
可视化看一下，这里只展示一部分。前面的一个列是所有anchor得到的最大conf值，后面的一列是对应的类别，比如第行中0.24885表示当前这个anchor预测为类别0【也就是persion类，置信度为0.24885】：
最后再用cat函数进行拼接，即将box【xyxy】，conf[所有anchor得到最大的conf], j[最大conf对应的类，mask]
conf, j = x[:, 5:mi].max(1, keepdim=True) x = torch.cat((box, conf, j.float(), mask), 1)[conf.view(-1) > conf_thres]得到新的x如下图：
下面这行代码是对上面得到x在conf这个维度上进行排序
x = x[x[:, 4].argsort(descending=True)] # sort by confidenceNMS处理agnostic参数 True表示多个类一起计算nms，False表示按照不同的类分别进行计算nms。
这里的nms是调用的torchvision下的nms。需要传入boxes，这里的boxes是加了c的偏移量【为什么加这个偏移量这里我没明白，有知道的可以留言说一下】。这里的boxes为【x1,y1,x2,y2】形式。score是已经排序好的，iou_thres是iou阈值。
# Batched NMS c = x[:, 5:6] * (0 if agnostic else max_wh) # classes boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores i = torchvision.ops.nms(boxes, scores, iou_thres) # NMS此时得到的i为：tensor([ 0, 2, 6, 16, 30], device='cuda:0')
这里的x是shape为【49，38】，49表示为有目标的49个anchors，38就是前面含有boxes,conf,index,mask的信息。i就是筛选得到的anchors索引。索引得到最终5个anchors以及信息，shape为【5，38】.
output[xi] = x[i]这个就是我们得到的经过NMS后的输出啦
上一篇：轻量级网络MobileNeXt--改进MobileNet v2的逆残差结构(轻量级网络设计)
