在上一篇文章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_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】.文章来源:https://www.toymoban.com/news/detail-463552.html
output[xi] = x[i]
这个就是我们得到的经过NMS后的输出啦 文章来源地址https://www.toymoban.com/news/detail-463552.html
到了这里,关于YOLOv5图像分割中的NMS处理的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!