YOLOv5改进之十九:非极大值抑制NMS算法改进DIoU NMS

​前 言:作为当前先进的深度学习目标检测算法YOLOv5,已经集合了大量的trick,但是还是有提高和改进的空间,针对具体应用场景下的检测难点,可以不同的改进方法。此后的系列文章,将重点对YOLOv5的如何改进进行详细的介绍,目的是为了给那些搞科研的同学需要创新点或者搞工程项目的朋友需要达到更好的效果提供自己的微薄帮助和参考。

解决问题:YOLOv5默认采用NMS算法,主要是通过IoU来筛选出候选框。NMS主要就是通过迭代的形式,不断的以最大得分的框去与其他框做IoU操作,并过滤那些IoU较大(即交集较大)的框。NMS缺点:1、NMS算法中的最大问题就是它将相邻检测框的分数均强制归零(即将重叠部分大于重叠阈值Nt的检测框移除)。在这种情况下,如果一个真实物体在重叠区域出现,则将导致对该物体的检测失败并降低了算法的平均检测率。2、NMS的阈值也不太容易确定,设置过小会出现误删,设置过高又容易增大误检。之前已经分享了用soft-nms软性非极大值抑制的方法优化原YOLOv5算法,有兴趣的朋友可以回顾一下。本文对nms进一步改进为DIoU NMS,在IoU损失中引入了一个惩罚项,直接最小化两个box的归一化距离。

YOLOv5改进之八:非极大值抑制NMS算法改进Soft-nms_人工智能算法研究院的博客-CSDN博客_yolov5 非极大值抑制https://blog.csdn.net/m0_70388905/article/details/125448230?spm=1001.2014.3001.5501

原理:

论文:[1911.08287] Distance-IoU Loss: Faster and Better Learning for Bounding Box Regression (arxiv.org)https://arxiv.org/abs/1911.08287

DIoU-NMS出现于Distance-IoU一文,研究者认为若相邻框的中心点越靠近当前最大得分框 的中心点,则其更有可能是冗余框。也就是说,当IoU相同时,如图所示。

YOLOv5改进之十九:非极大值抑制NMS算法改进DIoU NMS_第1张图片

当相邻框的中心点越靠近当前最大得分框的中心点,则可认为其更有可能是冗余框。第一种相比于第三种更不太可能是冗余框。因此,研究者使用所提出的DIoU替代IoU作为NMS的评判准则,公式如下:

方 法:

第一步:修改general.py,重新定义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 indices
            of the elements that have been kept
            by 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,CIOU
        iou = 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)

第二步:将general.py中将NMS改为DIOU NMS。

结 果:本人在多个数据集上做了大量实验,针对有遮挡目标的数据集检测效果提升明显。

预告一下:下一篇内容分享CIOU NMS。有兴趣的朋友可以关注一下我,有问题可以留言或者私聊我哦

PS:NMS的改进的方法不仅仅是适用改进YOLOv5,也可以改进其他的YOLO网络,比如YOLOv4、v3等。

最后,希望能互粉一下,做个朋友,一起学习交流。

你可能感兴趣的:(YOLO算法改进系列,深度学习,人工智能,计算机视觉,pytorch)