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的数量】。

YOLOv5图像分割中的NMS处理_第1张图片

 

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

可以从下面的代码看到在进行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_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】

YOLOv5图像分割中的NMS处理_第2张图片

最后再用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如下图:YOLOv5图像分割中的NMS处理_第3张图片

下面这行代码是对上面得到x在conf这个维度上进行排序 

x = x[x[:, 4].argsort(descending=True)]  # sort by confidence

NMS处理 

agnostic参数 True表示多个类一起计算nmsFalse表示按照不同的类分别进行计算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后的输出啦 

 

 

 

 

 

 

 

 

 

你可能感兴趣的:(图像分割,深度学习,计算机视觉,人工智能)