YOLOv5图像分割中的NMS处理

在上一篇文章YOLOv5图像分割--SegmentationModel类代码详解有讲到图像经过YOLOv5网络后得到的输出形式，主要是调用了BaseModel类下的forward得到的输出，输出的shape为【batch,25200,117】，这里的25200相当于总的anchors数量【以640640的输入为例，共有anchors=80803+40403+2020*3】，117为5[x,y,w,h,conf]+80个类+32【mask的数量】。

那么得到上面这张图的输出后又需要哪些处理呢？又是怎么处理的呢？本篇文章就是来刨析这个问题。

可以从下面的代码看到在进行model后会得到pred和proto。前者就是上面得到图的形式，后者的shape为【batch,32,160,160】，这里的32是mask的数量，160160是针对8080这个特征层的上采样得到的。然后是送入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中发生了什么？

那么来看一下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_conf

xywh2xyxy

下面的代码中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 confidence

NMS处理

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后的输出啦