Pytorch机器学习八—— YOLOV5中NMS非极大值抑制与DIOU-NMS等改进
目录
Pytorch机器学习八—— YOLOV5中NMS非极大值抑制与DIOU-NMS等改进
前言
一、NMS非极大值抑制算法
二、Hard-NMS非极大值代码
三、DIOU-NMS
四、soft-NMS
前言
在目标检测的预测阶段时会输出许多候选的anchor box其中有很多是明显重叠的预测边界框都围绕着同一个目标这时候我就可以使用NMS来合并同一目标的类似边界框或者说是保留这些边界框中最好的一个。
如果对IOU等知识不了解的可以看我上篇博客Pytorch机器学习五——目标检测中的损失函数l2IOUGIOUDIOU, CIOU
一、NMS非极大值抑制算法
我们先看一下NMS的直观理解左图为两个ground truth的bbox右图为我自己模拟网络输出的预测框。
而下图则是我使用Pytorch官方提供的NMS实现的非极大值抑制可以看到经过NMS后预测框保留了效果最好的去除了冗余的预测框。
下面来讲讲NMS算法的流程其实也是十分简单的
一.从所有候选框中选取置信度最高的预测边界框B1作为基准然后将所有与B1的IOU超过预定阈值的其他边界框移除。
这时所有边界框中B1为置信度最高的边界框且没有和其太过相似的边界框——非极大值置信度的边界框被抑制了
二.从所有候选框中选取置信度第二高的边界框B2作为一个基准将所有与B2的IOU超过预定阈值的其他边界框移除。
三.重复上述操作直到所有预测框都被当做基准——这时候没有一对边界框过于相似
二、Hard-NMS非极大值代码
在YOLOV5的源码当中作者是直接调用了Pytorch官方的NMS的API
在general.py中的non_max_suppression函数中
"""其中boxes为Nx4的tensorN为框的数量4则为x1 y1 x2 y2socres为N维的tensor表示每个框的置信度iou_thres则为上面算法中的IOU阈值返回值为一个去除了过于相似框后的根据置信度降序排列的列表我们就可以根据此列表输出预测框"""i torchvision.ops.nms(boxes, scores, iou_thres) # NMS
为了便于后续其他NMS的改进这里我们也自己写一个NMS算法这里借鉴了沐神的代码b站链接大家可以直接在YOLOV5中把上面的torchvision.ops.nms更改为下面的NMS函数
def NMS(boxes, scores, iou_thres, GIoUFalse, DIoUFalse, CIoUFalse): """ :param boxes: (Tensor[N, 4])): are expected to be in (x1, y1, x2, y2) :param scores: (Tensor[N]): scores for each one of the boxes :param iou_thres: discards all overlapping boxes with IoU > iou_threshold :return:keep (Tensor): int64 tensor with the indices of the elements that have been kept by NMS, sorted in decreasing order of scores """ # 按conf从大到小排序 B torch.argsort(scores, dim-1, descendingTrue) keep [] while B.numel() > 0: # 取出置信度最高的 index B[0] keep.append(index) if B.numel() 1: break # 计算iou,根据需求可选择GIOU,DIOU,CIOU iou bbox_iou(boxes[index, :], boxes[B[1:], :], GIoUGIoU, DIoUDIoU, CIoUCIoU) # 找到符合阈值的下标 inds torch.nonzero(iou < iou_thres).reshape(-1) B B[inds 1] return torch.tensor(keep)
这里的计算IOU的函数——bbox_iou则是直接引用了YOLOV5中的代码其简洁的集成了对与GIOU,DIOU,CIOU的计算。
def bbox_iou(box1, box2, x1y1x2y2True, GIoUFalse, DIoUFalse, CIoUFalse, eps1e-9): # Returns the IoU of box1 to box2. box1 is 4, box2 is nx4 box2 box2.T # Get the coordinates of bounding boxes if x1y1x2y2: # x1, y1, x2, y2 box1 b1_x1, b1_y1, b1_x2, b1_y2 box1[0], box1[1], box1[2], box1[3] b2_x1, b2_y1, b2_x2, b2_y2 box2[0], box2[1], box2[2], box2[3] else: # transform from xywh to xyxy b1_x1, b1_x2 box1[0] - box1[2] / 2, box1[0] box1[2] / 2 b1_y1, b1_y2 box1[1] - box1[3] / 2, box1[1] box1[3] / 2 b2_x1, b2_x2 box2[0] - box2[2] / 2, box2[0] box2[2] / 2 b2_y1, b2_y2 box2[1] - box2[3] / 2, box2[1] box2[3] / 2 # Intersection area inter (torch.min(b1_x2, b2_x2) - torch.max(b1_x1, b2_x1)).clamp(0) * \ (torch.min(b1_y2, b2_y2) - torch.max(b1_y1, b2_y1)).clamp(0) # Union Area w1, h1 b1_x2 - b1_x1, b1_y2 - b1_y1 eps w2, h2 b2_x2 - b2_x1, b2_y2 - b2_y1 eps union w1 * h1 w2 * h2 - inter eps iou inter / union if GIoU or DIoU or CIoU: cw torch.max(b1_x2, b2_x2) - torch.min(b1_x1, b2_x1) # convex (smallest enclosing box) width ch torch.max(b1_y2, b2_y2) - torch.min(b1_y1, b2_y1) # convex height if CIoU or DIoU: # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1 c2 cw ** 2 ch ** 2 eps # convex diagonal squared rho2 ((b2_x1 b2_x2 - b1_x1 - b1_x2) ** 2 (b2_y1 b2_y2 - b1_y1 - b1_y2) ** 2) / 4 # center distance squared if DIoU: return iou - rho2 / c2 # DIoU elif CIoU: # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47 v (4 / math.pi ** 2) * torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2) with torch.no_grad(): alpha v / ((1 eps) - iou v) return iou - (rho2 / c2 v * alpha) # CIoU else: # GIoU https://arxiv.org/pdf/1902.09630.pdf c_area cw * ch eps # convex area return iou - (c_area - union) / c_area # GIoU else: return iou # IoU
三、DIOU-NMS
其实DIOU-NMS就是把我上面说的NMS算法中的IOU阈值改为DIOU将NMS代码中的DIOU设置为True即可。
根据DIOU的论文如果只是单纯的使用NMS即是使用IOU作为阈值去筛掉其他预测框时当两个物体过于接近时很有可能另外一个物体的预测框就被滤除了。
就像下图中的摩托。使用DIOU-NMS可以一定程度上提升对于靠近的物体的检测。
四、soft-NMS
网上还有一种soft-NMS的算法其思想就是传统的NMS如果只通过IOU值就将其他的框直接去掉有可能会不妥于是就引入了soft-NMS。
具体流程就是我们把NMS算法中去除其他边界框改成修改其他边界框的置信度。
以下引一个博主的图
其中的f函数现在都是使用的高斯函数
si即为置信度M为当前最大置信度的边界框bi为其他边界框
网上对此的效果看法也是褒贬不一我自己也没有试过但从直觉来说我个人觉得效果不会有很大的提升如果感兴趣的可以自己试一试。
来源：lzzzzzzm