NMS和Soft-NMS的代码实现，数据可视化

这部分代码主要参考这位博主的讲解，原文链接：https://blog.csdn.net/AliceH1226/article/details/123429849

为什么要使用NMS

在目标检测中，模型输出的预测框数量往往远大于实际的真实框数量，并且是堆叠在一起的，非极大值抑制(Non-Maximum Suppression,NMS)可以从中筛选出符合要求的预测目标。因此，NMS通常往往作为目标检测中的后处理算法。

NMS算法流程

1> 根据检测框的置信度得分进行降序排序，选取分数最高的检测框A，
2> 分别计算检测框A与相邻检测框的重叠度IOU，对大于阈值的检测框得分直接置零，也就是扔掉；并标记将检测框A保留下来
3> 重复这个过程，找到所有被保留下来的检测框

Soft-NMS算法流程

1> 根据检测框的置信度得分进行降序排序，选取分数最高的检测框A，
2> 分别计算检测框A与相邻检测框的重叠度IOU，对大于阈值的检测框设置一个惩罚函数，降低这些检测框的置信度得分
3> 重复这个过程，找到所有被保留下来的检测框

 

import numpy as np
import numpy as np
import cv2

# 定义一个nms函数
def nms(boxes, threshold):
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]
    scores = boxes[:, 4]
    areas = (x2 - x1 + 1) * (y2 - y1 + 1)

    keep = []
    idxs = scores.argsort()[::-1]  # 从大到下排序,argsort函数返回从小到大的索引,[::-1]反序变成从大到小
    while (idxs.size > 0):
        i = idxs[0]
        keep.append(i)
        x11 = np.maximum(x1[i], x1[idxs[1:]])
        y11 = np.maximum(y1[i], y1[idxs[1:]])
        x22 = np.minimum(x2[i], y2[idxs[1:]])
        y22 = np.minimum(y2[i], y2[idxs[1:]])

        w = np.maximum(0, x22 - x11 + 1)
        h = np.maximum(0, y22 - y11 + 1)

        overlaps = w * h
        ious = overlaps / (areas[i] + areas[idxs[1:]] - overlaps)
        idxs2 = np.where(ious < threshold)[0]  # np.where函数
        idxs = idxs[idxs2 + 1]  # 注意这个+1

    return keep

# 定义一个soft_nms函数
def soft_nms(dets, thresh=0.3, sigma=0.5): # score大于thresh的才能存留下来，当设定的thresh过低，存留下来的框就很多,所以要根据实际情况调参
    '''
    input：
        dets: dets是(n,5)的ndarray，第0维度的每个元素代码一个框：[x1, y1, x2, y2, score] 
        thresh: float
        sigma: flaot
    output：
        index
    '''

    x1 = dets[:, 0] # dets:(n,5)  x1:(n,)  dets是ndarray, x1是ndarray
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    scores = dets[:, 4] # scores是ndarray


    # 每一个候选框的面积
    areas = (x2 - x1 + 1) * (y2 - y1 + 1) # areas:(n,)

    # order是按照score降序排序的
    order = scores.argsort()[::-1] # order:(n,) 降序下标 order是ndarray


    keep = []
    while order.size > 0:
        i = order[0] # i 是当下分数最高的框的下标
        # print(i)
        keep.append(i)
        # 计算当前概率最大矩形框与其他矩形框的相交框的坐标，会用到numpy的broadcast机制，得到的是向量

        # 当order只有一个值的时候，order[1]会报错说index out of range，而order[1:]会是[]，不报错，[]也可以作为x1的索引，x1[[]]为[]
   
        xx1 = np.maximum(x1[i], x1[order[1:]]) # xx1:(n-1,)的ndarray x1[i]:numpy_64浮点数一个，x1[order[1:]]是个ndarray，可以是空的ndarray，如果是空ndarray那么xx1为空ndarray，如果非空，那么x1[order[1:]]有多少个元素，xx1就是有多少个元素的ndarray。x1[]是不是ndarray看中括号内的是不是ndarray，看中括号内的是不是ndarray看中括号内的order[]的中括号内有没有冒号，有冒号的是ndarray，没有的是一个数。
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])
        

        # 计算相交框的面积,注意矩形框不相交时w或h算出来会是负数，用0代替
        w = np.maximum(0.0, xx2 - xx1 + 1) # xx2-xx1是(n-1,)的ndarray，w是(n-1,)的ndarray, n会逐渐减小至1
        # 当xx2和xx1是空的，那w是空的
        h = np.maximum(0.0, yy2 - yy1 + 1)

        inter = w * h # inter是(n,)的ndarray
        # 当w和h是空的，inter是空的

        # 计算重叠度IOU：重叠面积/（面积1+面积2-重叠面积）
        eps = np.finfo(areas.dtype).eps # 除法考虑分母为0的情况，np.finfo(dtype).eps，np.finfo(dtype)是个类，它封装了机器极限浮点类型的数，比如eps，episilon的缩写，表示小正数。
        ovr = inter / np.maximum(eps, areas[i] + areas[order[1:]] - inter) # n-1   #一旦（面积1+面积2-重叠面积）为0，就用eps进行替换
        # 当inter为空，areas[i]无论inter空不空都是有值的，那么ovr也为空

        # 更新分数
        weight = np.exp(-ovr*ovr/sigma)
        scores[order[1:]] *= weight

        # 更新order
        score_order = scores[order[1:]].argsort()[::-1] + 1
        order = order[score_order]

        keep_ids = np.where(scores[order]>thresh)[0]

        order = order[keep_ids]

    return keep

# 数据可视化
def plot(img,boxes,title):
    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = 0.5
    thickness = 1
    
    for box in boxes:
        x1, y1, x2, y2, score = int(box[0]), int(box[1]), int(box[2]), int(box[3]), box[4] # x:col y:row
        (w, h), baseline = cv2.getTextSize(str(score), font, font_scale, thickness)
        cv2.rectangle(img, (x1, y1), (x2, y2), (0,255,0), 2)
        cv2.rectangle(img, (x1, y1 - (2 * baseline + 5)), (x1 + w, y1), (0, 255, 0), -1)
        cv2.putText(img, str(score), (x1, y1), font, font_scale, (0, 0, 0), thickness)
    cv2.imshow('demo', img)
    cv2.imwrite("{}.jpg".format(title), img)
    cv2.waitKey(0)


if __name__=="__main__":
    # 读入图片，录入原始人框（[x1, y1, x2, y2, score]）
    image = cv2.imread('person.jpg')

    boxes = np.array([[149,149,282,546, 0.9874],
                     [116,103,303,565, 0.9356],
                     [59,133,474,620, 0.8898]])

    
    # 将预测框绘制在图像上
    image_for_nms_box = image.copy()
    plot(image_for_nms_box,boxes,"Image")
    
    # 1.使用nms对预测框进行筛选
    keep_1 = nms(boxes,threshold=0.3)
    nms_boxes = boxes[keep_1]
    
    # 将筛选过后的预测框绘制在图像上
    image_for_nms_box = image.copy()
    plot(image_for_nms_box,nms_boxes,"NMS")

    # 2.使用soft_nms对框进行筛选
    keep_2 = soft_nms(boxes)
    soft_nms_boxes = boxes[keep_2]
    
    # 将筛选过后的预测框绘制在图像上
    image_for_nms_box = image.copy()
    plot(image_for_nms_box,soft_nms_boxes,"Soft_NMS")

    cv2.destroyAllWindows()