目标检测算法——YOLOv5将NMS替换为DIoU-NMS

论文题目：《Distance-IoU Loss: Faster and Better Learning for Bounding Box Regression》
论文地址：  https://arxiv.org/pdf/1807.06521.pdf

非最大抑制（NMS）主要用于基于深度学习的目标检测模型输出的后处理，从而去除冗余的检测框，获得正确的检测结果。具体意义：在目标检测的预测阶段时，会输出许多候选的anchor box，其中有很多是明显重叠的预测边界框都围绕着同一个目标，这时候就可以使用NMS来合并同一目标的类似边界框，或者说是保留这些边界框中最好的一个。

在经典的NMS中，得分最高的检测框和其它检测框逐一算出一个对应的IOU值，并将该值超过NMS threshold的框全部过滤掉。可以看出，在经典NMS算法中，IOU是唯一考量的因素。但是在实际应用场景中，当两个不同物体挨得很近时，由于IOU值比较大，往往经过NMS处理后，只剩下一个检测框，这样导致漏检的错误情况发生。基于此，DIOU-NMS就不仅仅考虑IOU，还考虑两个框中心点之间的距离。如果两个框之间IOU比较大，但是两个框的距离比较大时，可能会认为这是两个物体的框而不会被过滤掉。 

实验证明，将NMS替换为DIoU-NMS，可初步改善YOLOv5对重叠遮挡目标的识别。近期较忙，想要代码的小伙伴请私信！

DIoU-NMS代码如下：

   def diou_box_nms(self, scores, boxes, iou_thres):
        if boxes.shape[0] == 0:
            return torch.zeros(0,device=boxes.device).long()
        x1,y1,x2,y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:,3]
        areas = (x2 - x1 + 1) * (y2 - y1 + 1)
        order = torch.sort(scores, descending=True)[1] #(?,)
        keep =[]
        while order.numel() > 0:
            if order.numel() == 1:
                keep.append(order.item())
                break
            else:
                i = order[0].item()
                keep.append(i)

                xmin = torch.clamp(x1[order[1:]], min = float(x1[i]))
                ymin = torch.clamp(y1[order[1:]], min = float(y1[i]))
                xmax = torch.clamp(x2[order[1:]], max = float(x2[i]))
                ymax = torch.clamp(y2[order[1:]], max = float(y2[i]))

                inter_area = torch.clamp(xmax - xmin, min=0.0) * torch.clamp(ymax - ymin, min=0.0)

                iou = inter_area / (areas[i] + areas[order[1:]] - inter_area + 1e-16)

                # diou add center
                # inter_diag
                cxpreds = (x2[i] + x1[i]) / 2
                cypreds = (y2[i] + y1[i]) / 2

                cxbbox = (x2[order[1:]] + x1[order[1:]]) / 2
                cybbox = (y1[order[1:]] + y2[order[1:]]) / 2

                inter_diag = (cxbbox - cxpreds) ** 2 + (cybbox - cypreds) ** 2

                # outer_diag
                ox1 = torch.min(x1[order[1:]], x1[i])
                oy1 = torch.min(y1[order[1:]], y1[i])
                ox2 = torch.max(x2[order[1:]], x2[i])
                oy2 = torch.max(y2[order[1:]], y2[i])

                outer_diag = (ox1 - ox2) ** 2 + (oy1 - oy2) ** 2

                diou = iou - inter_diag / outer_diag
                diou = torch.clamp(diou, min=-1.0, max=1.0)


                # mask_ind = (iou <= iou_thres).nonzero().squeeze()
                mask_ind = (diou <= iou_thres).nonzero().squeeze()

                if mask_ind.numel() == 0:
                    break
                order = order[mask_ind + 1]
        return torch.LongTensor(keep)

由于原始的NMS中，IoU指标用于抑制多余的检测框，但由于仅考虑了重叠区域，经常会造成错误的抑制，特别是在bbox包含的情况下。因此，可以使用DIoU作为NMS的标准，不仅考虑重叠区域，还考虑了中心点距离。基于DIoU作为NMS标准，虽然多了距离这个维度去考虑问题，但和NMS面对的同样的情况是当两个不同的目标本身就靠的很近的时候还是会造成错误的抑制。

单纯的使用NMS，即是使用IOU作为阈值去筛掉其他预测框时，当两个物体过于接近时，很有可能另外一个物体的预测框就被滤除了。然而，使用DIOU-NMS可以一定程度上提升对于靠近物体的检测。