yolov5实例分割结果处理（YOLOv5图像分割中的NMS处理）
在上一篇文章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后的输出啦
创心域SEO版权声明：以上内容作者已申请原创保护，未经允许不得转载，侵权必究！授权事宜、对本内容有异议或投诉，敬请联系网站管理员，我们将尽快回复您，谢谢合作！