Pytorch机器学习（八）—— YOLOV5中NMS非极大值抑制与DIOU-NMS等改进

前言

一、NMS非极大值抑制算法

二、Hard-NMS非极大值代码

三、DIOU-NMS

四、soft-NMS

前言

在目标检测的预测阶段时，会输出许多候选的anchor box，其中有很多是明显重叠的预测边界框都围绕着同一个目标，这时候我就可以使用NMS来合并同一目标的类似边界框，或者说是保留这些边界框中最好的一个。

如果对IOU等知识不了解的可以看我上篇博客Pytorch机器学习（五）——目标检测中的损失函数（l2，IOU，GIOU，DIOU, 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的tensor，N为框的数量，4则为x1 y1 x2 y2
socres为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, GIoU=False, DIoU=False, CIoU=False):""":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 indicesof the elements that have been keptby NMS, sorted in decreasing order of scores"""# 按conf从大到小排序B = torch.argsort(scores, dim=-1, descending=True)keep = []while B.numel() > 0:# 取出置信度最高的index = B[0]keep.append(index)if B.numel() == 1: break# 计算iou,根据需求可选择GIOU,DIOU,CIOUiou = bbox_iou(boxes[index, :], boxes[B[1:], :], GIoU=GIoU, DIoU=DIoU, CIoU=CIoU)# 找到符合阈值的下标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, x1y1x2y2=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-9):# Returns the IoU of box1 to box2. box1 is 4, box2 is nx4box2 = box2.T# Get the coordinates of bounding boxesif x1y1x2y2:  # x1, y1, x2, y2 = box1b1_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 xyxyb1_x1, b1_x2 = box1[0] - box1[2] / 2, box1[0] + box1[2] / 2b1_y1, b1_y2 = box1[1] - box1[3] / 2, box1[1] + box1[3] / 2b2_x1, b2_x2 = box2[0] - box2[2] / 2, box2[0] + box2[2] / 2b2_y1, b2_y2 = box2[1] - box2[3] / 2, box2[1] + box2[3] / 2# Intersection areainter = (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 Areaw1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + epsw2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + epsunion = w1 * h1 + w2 * h2 - inter + epsiou = inter / unionif GIoU or DIoU or CIoU:cw = torch.max(b1_x2, b2_x2) - torch.min(b1_x1, b2_x1)  # convex (smallest enclosing box) widthch = torch.max(b1_y2, b2_y2) - torch.min(b1_y1, b2_y1)  # convex heightif CIoU or DIoU:  # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1c2 = cw ** 2 + ch ** 2 + eps  # convex diagonal squaredrho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 +(b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4  # center distance squaredif DIoU:return iou - rho2 / c2  # DIoUelif CIoU:  # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47v = (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)  # CIoUelse:  # GIoU https://arxiv.org/pdf/1902.09630.pdfc_area = cw * ch + eps  # convex areareturn iou - (c_area - union) / c_area  # GIoUelse:return iou  # IoU