目标检测中NMS(non maximum suppression)

一、原理

---

参考网址

NMS即non maximum suppression即非极大抑制，顾名思义就是抑制不是极大值的元素，搜索局部的极大值。在最近几年常见的物体检测算法（包括rcnn、sppnet、fast-rcnn、faster-rcnn等）中，最终都会从一张图片中找出很多个可能是物体的矩形框，然后为每个矩形框为做类别分类概率。

Soft NMS法：对于得分小于阈值的边框，不再直接舍弃，而是降低其得分。Soft NMS是对NMS的优化算法，它在不增加额外参数的情况下且只需要对NMS算法进行简单的改动就能提高AP。该Soft-NMS算法在标准数据集PASCAL VOC2007（较R-FCN和Faster-RCNN提升1.7%）和MS-COCO（较R-FCN提升1.3%，较Faster-RCNN提升1.1%）上均有提升。对于大多数数据集而言，作用比较小，提升效果非常不明显，它起作用的地方是大量密集的同类重叠场景，大量密集的不同类重叠场景其实也没什么作用。

二、实现步骤

---

1、NMS算法实现

---

参考网址

假设我们已经有了预测的框，每个预测框对应的类别，每个预测框对应的类别得分。本文中使用的案例是n个[x1,y1,x2,y2,confident, 0,0,0,1,0,0] --> n个[左上坐标，右下坐标，置信度，种类的one-hot编码]。

（2）对于每个类别而言，取该类中计算得分最大的框与其余预测框之间的IoU。

（3）根据设定的阈值，剔除掉该类中IOU大于阈值的预测框。

（4）对于每个类别而言，在(3)的基础上循环(1)-(3)步骤，在该类中已经排好序的情况下，直接下一个最大值作为该类的得分最大的框与其余预测框之间的IoU；直到剩余有用的预测框个数为0。

2、softNMS算法实现

---

论文地址：https://arxiv.org/pdf/1704.04503v2.pdf

（1）对于每个类别，按照类别得分的从大到小顺序排列

（2）对于每个类别而言，取该类中计算得分最大的框与其余预测框之间的IoU。

（3）根据计算的iou，使用高斯惩罚函数公式

Si 表示置信度，当iou越大时，该预测框的置信度就会施加更重的惩罚值，从而降低置信度。

（4）对于每个类别而言，在(3)对置信度施加惩罚后，再剔除掉该类中置信度小于阈值的预测框。

（5）重复(1) - (4)步骤，直至剩余的预测框个数为0。

三、代码实现

---

1、Pytorch版

---

import numpy as np
import cv2
import matplotlib.pyplot as plt
import torch

def bbox_iou(box1, box2, x1y1x2y2=True):
    # 坐标分离
    if not x1y1x2y2:
        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
    else:
        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]
    
    # 计算交集的两坐标（即交集的左上、右下坐标）
    inter_rect_x1 = torch.max(b1_x1, b2_x1)
    inter_rect_y1 = torch.max(b1_y1, b2_y1)
    inter_rect_x2 = torch.min(b1_x2, b2_x2)
    inter_rect_y2 = torch.min(b1_y2, b2_y2)
    # 计算交集面积
    inter_area = torch.clamp(inter_rect_x2 - inter_rect_x1, min=0) *\
                torch.clamp(inter_rect_y2 - inter_rect_y1, min=0)
    # 分别计算两个目标框各自的面积
    b1_area = (b1_x2 - b1_x1) * (b1_y2 - b1_y1)
    b2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1)
    # 交集/并集
    iou = inter_area / torch.clamp(b1_area + b2_area - inter_area, min= 1e-6)
    return iou


def nms(boxes, nms_thres=0.5):
    result = []
    boxes = torch.Tensor(boxes)
    #------------------------------------------#
    #   获得预测结果中包含的所有种类
    #------------------------------------------#
    unique_labels = boxes[:, -1].cpu().unique()
    
    for c in unique_labels:
        #------------------------------------------#
        #   获得某一类得分筛选后全部的预测结果
        #------------------------------------------#
        detections_class = boxes[boxes[:, -1] == c]
        
        # 按照存在物体的置信度排序
        _, conf_sort_index = torch.sort(detections_class[:, 4], descending=True)
        detections_class = detections_class[conf_sort_index]
        # 进行非极大抑制
        max_detections = []
        while detections_class.size(0):
            # 取出这一类置信度最高的，一步一步往下判断，判断重合程度是否大于nms_thres，如果是则去除掉
            max_detections.append(detections_class[0].unsqueeze(0))
            if len(detections_class) == 1:
                break
            ious = bbox_iou(max_detections[-1], detections_class[1:])
            detections_class = detections_class[1:][ious < nms_thres]
        # 堆叠
        max_detections = torch.cat(max_detections).data
        result.append(max_detections)
    result = torch.cat(result).data

    # ===================================================================================== #
    #                                          案例绘图                                      #
    # ===================================================================================== #
    plt.figure()
    colors = [(255,0,0),(0,255,0),(0,0,255)]
    img = np.zeros((616, 616, 3))
    plt.subplot(121)
    for i in result:
        cv2.rectangle(img, (int(i[0]), int(i[1])), (int(i[2]), int(i[3])), colors[int(i[5])], 2)
    plt.imshow(img)
    plt.title('nms')

    plt.subplot(122)
    img_ = np.zeros((616, 616, 3))
    for i in boxes:
        cv2.rectangle(img_, (int(i[0]), int(i[1])), (int(i[2]), int(i[3])), colors[int(i[5])], 2)
    plt.imshow(img_)
    plt.title('original')
    plt.show()
    # ===================================================================================== #
    return result # 返回nms后的结果
   

def soft_nms(boxes,conf_thres=0.5,sigma=0.5):
    result = []
    boxes = torch.Tensor(boxes)
    #------------------------------------------#
    #   获得预测结果中包含的所有种类
    #------------------------------------------#
    unique_labels = boxes[:, -1].cpu().unique()
    
    for c in unique_labels:
        #------------------------------------------#
        #   获得某一类得分筛选后全部的预测结果
        #------------------------------------------#
        detections_class = boxes[boxes[:, -1] == c]
        
        # 按照存在物体的置信度排序
        _, conf_sort_index = torch.sort(detections_class[:, 4], descending=True)
        detections_class = detections_class[conf_sort_index]
        # 进行非极大抑制
        max_detections = []
        while detections_class.size(0):
            # 取出这一类置信度最高的，一步一步往下判断，根据iou设置一个对置信度的惩罚因子，去除置信度小于conf_thres的框
            max_detections.append(detections_class[0].unsqueeze(0))
            if len(detections_class) == 1:
                break
            ious                    = bbox_iou(max_detections[-1], detections_class[1:])
            detections_class[1:, 4] = torch.exp(-(ious * ious) / sigma) * detections_class[1:, 4]
            detections_class        = detections_class[1:]
            detections_class        = detections_class[detections_class[:, 4] >= conf_thres]
            arg_sort                = torch.argsort(detections_class[:, 4], descending = True)
            detections_class        = detections_class[arg_sort]

        # 堆叠
        max_detections = torch.cat(max_detections).data
        result.append(max_detections)
    result = torch.cat(result).data

    # ===================================================================================== #
    #                                          案例绘图                                      #
    # ===================================================================================== #
    plt.figure()
    colors = [(255,0,0),(0,255,0),(0,0,255)]
    img = np.zeros((616, 616, 3))
    plt.subplot(121)
    for i in result:
        cv2.rectangle(img, (int(i[0]), int(i[1])), (int(i[2]), int(i[3])), colors[int(i[5])], 2)
    plt.imshow(img)
    plt.title('nms')

    plt.subplot(122)
    img_ = np.zeros((616, 616, 3))
    for i in boxes:
        cv2.rectangle(img_, (int(i[0]), int(i[1])), (int(i[2]), int(i[3])), colors[int(i[5])], 2)
    plt.imshow(img_)
    plt.title('original')
    plt.show()
    # ===================================================================================== #
    return result # 返回nms后的结果

if __name__ == '__main__':
    boxes = np.array([[100, 110, 210, 210, 0.71, 0.7, 0.3, 0.5],
                    [250, 250, 420, 420, 0.8, 0.1, 0.8, 0.6],
                    [220, 200, 320, 330, 0.92, 0.2, 0.5, 1.0],
                    [120, 100, 210, 210, 0.72, 0.8, 0.2, 0.3],
                    [230, 240, 325, 330, 0.81, 0.1, 0.9, 0.2],
                    [220, 230, 315, 340, 0.91, 0.2, 0.7, 0.6]])
    # 转换成[x1,y1,x2,y2,confident,种类下标]
    boxes = np.hstack((boxes[...,:5],np.expand_dims(np.argmax(boxes[:,5:],axis=-1),axis=-1)))
    nms(boxes,nms_thres=0.5)
    soft_nms(boxes,0.5,0.1)    #boxes,threshold,λ(gauss函数参数）

效果图：

2、Numpy版

---

import numpy as np
import cv2
import matplotlib.pyplot as plt

'''
    numpy版
'''
import numpy as np
def bbox_iou(box1, box2, x1y1x2y2=True):
    if not x1y1x2y2:
        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
    else:
        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]
    
    inter_rect_x1 = np.maximum(b1_x1, b2_x1)
    inter_rect_y1 = np.maximum(b1_y1, b2_y1)
    inter_rect_x2 = np.minimum(b1_x2, b2_x2)
    inter_rect_y2 = np.minimum(b1_y2, b2_y2)
    
    inter_area = np.maximum(inter_rect_x2 - inter_rect_x1, 0) *\
                np.maximum(inter_rect_y2 - inter_rect_y1, 0)
    b1_area = (b1_x2 - b1_x1) * (b1_y2 - b1_y1)
    b2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1)

    iou = inter_area / np.maximum(b1_area + b2_area - inter_area, 1e-6)
    return iou

def nms(boxes, nms_thres=0.5):
    result = []
    #------------------------------------------#
    #   获得预测结果中包含的所有种类
    #------------------------------------------#
    unique_labels = np.unique(boxes[:, -1])
    for c in unique_labels:
        #------------------------------------------#
        #   获得某一类得分筛选后全部的预测结果
        #------------------------------------------#
        detections_class = boxes[boxes[:, -1] == c]
        # 按照存在物体的置信度排序
        conf_sort_index = np.argsort(detections_class[:, 4])[::-1]
        detections_class = detections_class[conf_sort_index]
        # 进行非极大抑制
        max_detections = []
        while len(detections_class) != 0:
            # 取出这一类置信度最高的，一步一步往下判断，判断重合程度是否大于nms_thres，如果是则去除掉
            max_detections.append(np.expand_dims(detections_class[0],axis=0))
            if len(detections_class) == 1:
                break
            ious = bbox_iou(max_detections[-1], detections_class[1:])
            detections_class = detections_class[1:][ious < nms_thres]
        # 堆叠
        max_detections = np.concatenate(max_detections)
        result.append(max_detections)
    result = np.concatenate(result)

    # ===================================================================================== #
    #                                          案例绘图                                       #
    # ===================================================================================== #
    plt.figure()
    colors = [(255,0,0),(0,255,0),(0,0,255)]
    img = np.zeros((616, 616, 3))
    plt.subplot(121)
    for i in result:
        cv2.rectangle(img, (int(i[0]), int(i[1])), (int(i[2]), int(i[3])), colors[int(i[5])], 2)
    plt.imshow(img)
    plt.title('nms')

    plt.subplot(122)
    img_ = np.zeros((616, 616, 3))
    for i in boxes:
        cv2.rectangle(img_, (int(i[0]), int(i[1])), (int(i[2]), int(i[3])), colors[int(i[5])], 2)
    plt.imshow(img_)
    plt.title('original')
    plt.show()
    # ========================================================================================== #
    return result # 返回nms后的结果
   

def soft_nms(boxes,conf_thres=0.5,sigma=0.5):
    result = []
    #------------------------------------------#
    #   获得预测结果中包含的所有种类
    #------------------------------------------#
    unique_labels = np.unique(boxes[:, -1])
    
    for c in unique_labels:
        #------------------------------------------#
        #   获得某一类得分筛选后全部的预测结果
        #------------------------------------------#
        detections_class = boxes[boxes[:, -1] == c]
        
        # 按照存在物体的置信度排序
        conf_sort_index = np.argsort(detections_class[:, 4])[::-1]
        detections_class = detections_class[conf_sort_index]
        # 进行非极大抑制
        max_detections = []
        while len(detections_class) != 0:
            # 取出这一类置信度最高的，一步一步往下判断，根据iou设置一个对置信度的惩罚因子，去除置信度小于conf_thres的框
            max_detections.append(np.expand_dims(detections_class[0],axis=0))
            if len(detections_class) == 1:
                break
            ious                    = bbox_iou(max_detections[-1], detections_class[1:])
            detections_class[1:, 4] = np.exp(-(ious * ious) / sigma) * detections_class[1:, 4]
            detections_class        = detections_class[1:]
            detections_class        = detections_class[detections_class[:, 4] >= conf_thres]
            arg_sort                = np.argsort(detections_class[:, 4])[::-1]
            detections_class        = detections_class[arg_sort]

        # 堆叠
        max_detections = np.concatenate(max_detections)
        result.append(max_detections)
    result = np.concatenate(result)

    # =================================================================================== #
    #                                          案例绘图                                    #
    # =================================================================================== #
    plt.figure()
    colors = [(255,0,0),(0,255,0),(0,0,255)]
    img = np.zeros((616, 616, 3))
    plt.subplot(121)
    for i in result:
        cv2.rectangle(img, (int(i[0]), int(i[1])), (int(i[2]), int(i[3])), colors[int(i[5])], 2)
    plt.imshow(img)
    plt.title('nms')

    plt.subplot(122)
    img_ = np.zeros((616, 616, 3))
    for i in boxes:
        cv2.rectangle(img_, (int(i[0]), int(i[1])), (int(i[2]), int(i[3])), colors[int(i[5])], 2)
    plt.imshow(img_)
    plt.title('original')
    plt.show()
    # ========================================================================================== #
    return result # 返回nms后的结果

if __name__ == '__main__':
    boxes = np.array([[100, 110, 210, 210, 0.71, 0.7, 0.3, 0.5],
                    [250, 250, 420, 420, 0.8, 0.1, 0.8, 0.6],
                    [220, 200, 320, 330, 0.92, 0.2, 0.5, 1.0],
                    [120, 100, 210, 210, 0.72, 0.8, 0.2, 0.3],
                    [230, 240, 325, 330, 0.81, 0.1, 0.9, 0.2],
                    [220, 230, 315, 340, 0.91, 0.2, 0.7, 0.6]])
    # 转换成[x1,y1,x2,y2,confident,种类下标]
    boxes = np.hstack((boxes[...,:5],np.expand_dims(np.argmax(boxes[:,5:],axis=-1),axis=-1)))

    nms(boxes,nms_thres=0.5)
    soft_nms(boxes,0.5,0.1)    #boxes,threshold,λ(gauss函数参数）

效果图：

三、总结

---

本文章分别由算法原理、算法实现步骤和代码展示三个部分组成。如果你觉得本章论文对你有帮助，请点个，谢谢。

四、参考

---

https://blog.csdn.net/weixin_44791964/article/details/106222846

https://blog.csdn.net/qq_25344301/article/details/120094530

https://www.bilibili.com/video/BV1aB4y1K7za/?p=1&vd_source=bb59b9d0f88f71826610fe604fbe264e

https://blog.csdn.net/lzzzzzzm/article/details/120151155?ops_request_misc=%257B%2522request%255Fid%2522%253A%2522167272576716800222869958%2522%252C%2522scm%2522%253A%252220140713.130102334.pc%255Fall.%2522%257D&request_id=167272576716800222869958&biz_id=0&utm_medium=distribute.pc_search_result.none-task-blog-2~all~first_rank_ecpm_v1~rank_v31_ecpm-4-120151155-null-null.142v68control,201v4add_ask,213v2t3_control2&utm_term=yolov5%E9%9D%9E%E6%9E%81%E5%A4%A7%E5%80%BC%E6%8A%91%E5%88%B6%E7%AE%97%E6%B3%95&spm=1018.2226.3001.4187