python nms实现和soft-nms官方实现

#NMS实现
import numpy as np


boxes=np.array([[100,100,210,210,0.72],
        [250,250,420,420,0.8],
        [220,220,320,330,0.92],
        [100,100,210,210,0.72],
        [230,240,325,330,0.81],
        [220,230,315,340,0.9]]) 


def py_cpu_nms(dets, thresh):
    # dets:(m,5)  thresh:scaler
    
    x1 = dets[:,0]
    y1 = dets[:,1]
    x2 = dets[:,2]
    y2 = dets[:,3]
    
    areas = (y2-y1+1) * (x2-x1+1)
    scores = dets[:,4]
    keep = []
    
    index = scores.argsort()[::-1]
    
    while index.size >0:

        i = index[0]       # every time the first is the biggst, and add it directly
        keep.append(i)
        
        x11 = np.maximum(x1[i], x1[index[1:]])    #计算iou 
        y11 = np.maximum(y1[i], y1[index[1:]])
        x22 = np.minimum(x2[i], x2[index[1:]])
        y22 = np.minimum(y2[i], y2[index[1:]])
        
        w = np.maximum(0, x22-x11+1)    # the weights of overlap
        h = np.maximum(0, y22-y11+1)    # the height of overlap
       
        overlaps = w*h
        
        ious = overlaps / (areas[i]+areas[index[1:]] - overlaps)
        
        idx = np.where(ious<=thresh)[0]
        
        index = index[idx+1]   #index下标是从1开始的，而where计算的下标是从0开始的，故需要+1
        
    return keep
        

import matplotlib.pyplot as plt
def plot_bbox(dets, c='k'):
    
    x1 = dets[:,0]
    y1 = dets[:,1]
    x2 = dets[:,2]
    y2 = dets[:,3]
    
    
    plt.plot([x1,x2], [y1,y1], c)
    plt.plot([x1,x1], [y1,y2], c)
    plt.plot([x1,x2], [y2,y2], c)
    plt.plot([x2,x2], [y1,y2], c)
    plt.title("after nms")

plot_bbox(boxes,'k')  

keep = py_cpu_nms(boxes, thresh=0.7)
plot_bbox(boxes[keep], 'r')

# ----------------------------------------------------------
# Soft-NMS: Improving Object Detection With One Line of Code
# Copyright (c) University of Maryland, College Park
# Licensed under The MIT License [see LICENSE for details]
# Written by Navaneeth Bodla and Bharat Singh
# ----------------------------------------------------------

import numpy as np
cimport numpy as np

# 本文件是.pyx文件，是python的c扩展文件，要想被python调用、运行，仅仅写了源代码还是不够的，还要转成.c或者.c++的文件，并且再进一步转成.pyd文件；
# .pyd文件才是可以直接使用的文件，为了达到上述目的，就要写一个setup.py脚本，这个在nms文件夹中都有，就不专门介绍了；

# Cython是让Python脚本支持C语言扩展的编译器，Cython能够将Python+C混合编码的.pyx脚本转换为C代码，主要用于优化Python脚本性能或Python调用C函数库；
# 由于Python固有的性能差的问题，用C扩展Python成为提高Python性能常用方法，Cython算是较为常见的一种扩展方式；


# max函数
cdef inline np.float32_t max(np.float32_t a, np.float32_t b):
    return a if a >= b else b

# min函数
cdef inline np.float32_t min(np.float32_t a, np.float32_t b):
    return a if a <= b else b

# origin nms操作
def cpu_nms(np.ndarray[np.float32_t, ndim=2] dets, np.float thresh):
    cdef np.ndarray[np.float32_t, ndim=1] x1 = dets[:, 0]        # pred bbox top_x
    cdef np.ndarray[np.float32_t, ndim=1] y1 = dets[:, 1]        # pred bbox top_y
    cdef np.ndarray[np.float32_t, ndim=1] x2 = dets[:, 2]        # pred bbox bottom_x
    cdef np.ndarray[np.float32_t, ndim=1] y2 = dets[:, 3]        # pred bbox bottom_y
    cdef np.ndarray[np.float32_t, ndim=1] scores = dets[:, 4]    # pred bbox cls score

    cdef np.ndarray[np.float32_t, ndim=1] areas = (x2 - x1 + 1) * (y2 - y1 + 1)     # pred bbox areas
    cdef np.ndarray[np.int_t, ndim=1] order = scores.argsort()[::-1]                # 对pred bbox按score做降序排序，对应step-2

    cdef int ndets = dets.shape[0]                                                  # num of detected bbox
    cdef np.ndarray[np.int_t, ndim=1] suppressed = np.zeros((ndets), dtype=np.int)  # 相当于flag，与bbox对应，如果其已经在nms操作中被抑制(被认为与其他高score IoU过大，可剔除)，就置suppressed = 1，表示该bbox已经不纳入考虑
 
    cdef int _i, _j                               # nominal indices，和C的操作有点类似，先申明变量
    cdef int i, j                                 # sorted indices
    cdef np.float32_t ix1, iy1, ix2, iy2, iarea   # temp variables for box i's (the box currently under consideration)
    cdef np.float32_t xx1, yy1, xx2, yy2          # variables for computing overlap with box j (lower scoring box)
    cdef np.float32_t w, h
    cdef np.float32_t inter, ovr

    keep = []
    for _i in range(ndets):
        i = order[_i]            # 取当前index _i的score bbox，对应着此轮的最高score bbox
        if suppressed[i] == 1:   # 之前NMS操作已经被干掉了，无效bbox，那就忽略吧
            continue
        keep.append(i)           # 保留之
        ix1 = x1[i]
        iy1 = y1[i]
        ix2 = x2[i]
        iy2 = y2[i]
        iarea = areas[i]         # 面积
        for _j in range(_i + 1, ndets):    # 计算index _i的score bbox，与其之后bbox的IoU，进而做NMS
            j = order[_j]
            if suppressed[j] == 1:         # 无效bbox，忽略
                continue
            xx1 = max(ix1, x1[j])          # 为计算IoU做准备
            yy1 = max(iy1, y1[j])
            xx2 = min(ix2, x2[j])
            yy2 = min(iy2, y2[j])
            w = max(0.0, xx2 - xx1 + 1)    # Iinsection的宽、高、面积
            h = max(0.0, yy2 - yy1 + 1)
            inter = w * h
            ovr = inter / (iarea + areas[j] - inter)    # IoU
            if ovr >= thresh:         # 如果当前bbox与index _i的bbox，IoU过大，就要被抑制掉了
                suppressed[j] = 1

    return keep    # 最终NMS被保留的bbox

# soft_nms操作，这里假设boxes是无序(未按score做降序)的，所以每轮soft_nms迭代都需要类似冒泡排序操作，选择当前top-1 bbox做NMS
# Nt：计算IoU的阈值，IoU > Nt，对应bbox的score权重就要降低
# threshold：降权后通过threshold进一步剔除低权重bbox
def cpu_soft_nms(np.ndarray[float, ndim=2] boxes, float sigma=0.5, float Nt=0.3, float threshold=0.001, unsigned int method=0):
    cdef unsigned int N = boxes.shape[0]    # num of detected bbox
    cdef float iw, ih, box_area
    cdef float ua
    cdef int pos = 0
    cdef float maxscore = 0
    cdef int maxpos = 0
    cdef float x1,x2,y1,y2,tx1,tx2,ty1,ty2,ts,area,weight,ov

    for i in range(N):
        maxscore = boxes[i, 4]    # 获取当前index下的bbox
        maxpos = i

        tx1 = boxes[i,0]
        ty1 = boxes[i,1]
        tx2 = boxes[i,2]
        ty2 = boxes[i,3]
        ts = boxes[i,4]

        pos = i + 1      # 下面操作就很常规了，找到当前index i之后所有bboxes中，score最大的bbox，并将之赋值给maxscore、maxpos
        while pos < N:
            if maxscore < boxes[pos, 4]:
                maxscore = boxes[pos, 4]
                maxpos = pos
            pos = pos + 1

        # 下面操作更简单，想想我们最开始学C语言，a、b两变量如何交换
	    # add max box as a detection 
        boxes[i,0] = boxes[maxpos,0]    # maxpos内的信息，放到index i处，也是当前需要处理的bbox
        boxes[i,1] = boxes[maxpos,1]
        boxes[i,2] = boxes[maxpos,2]
        boxes[i,3] = boxes[maxpos,3]
        boxes[i,4] = boxes[maxpos,4]

	    # swap ith box with position of max box
        boxes[maxpos,0] = tx1           # 别忘了tx1中可是保存了boxes[i,0]备份的
        boxes[maxpos,1] = ty1
        boxes[maxpos,2] = tx2
        boxes[maxpos,3] = ty2
        boxes[maxpos,4] = ts

        tx1 = boxes[i,0]   # 此时tx1就保存的maxpos位置的bbox信息了
        ty1 = boxes[i,1]
        tx2 = boxes[i,2]
        ty2 = boxes[i,3]
        ts = boxes[i,4]

        pos = i + 1
	    # NMS iterations, note that N changes if detection boxes fall below threshold，N值是动态变化的
        while pos < N:     # 向后做NMS比较
            x1 = boxes[pos, 0]   # 当前位置的bbox
            y1 = boxes[pos, 1]
            x2 = boxes[pos, 2]
            y2 = boxes[pos, 3]
            s = boxes[pos, 4]

            area = (x2 - x1 + 1) * (y2 - y1 + 1)          # pos下box的面积
            iw = (min(tx2, x2) - max(tx1, x1) + 1)        # 计算Insection的宽iw，如果iw < 0，说明没相交，可以直接忽略了
            if iw > 0:
                ih = (min(ty2, y2) - max(ty1, y1) + 1)    # 计算Insection的宽ih，如果ih < 0，说明没相交，可以直接忽略了
                if ih > 0:
                    ua = float((tx2 - tx1 + 1) * (ty2 - ty1 + 1) + area - iw * ih)   # U的面积
                    ov = iw * ih / ua                                                # iou between max box and detection box

                    if method == 1:                       # soft_nms中linear降权操作，与ov负相关
                        if ov > Nt: 
                            weight = 1 - ov
                        else:
                            weight = 1
                    elif method == 2:                     # soft_nms中gaussian降权操作
                        weight = np.exp(-(ov * ov)/sigma)
                    else:                                 # original NMS，weight = 0就直接把score置0
                        if ov > Nt: 
                            weight = 0
                        else:
                            weight = 1

                    boxes[pos, 4] = weight * boxes[pos, 4]  # 权重重新调整
		    
		            # if box score falls below threshold, discard the box by swapping with last box，update N
                    # 如果bbox调整后的权重，已经小于阈值threshold，那么这个bbox就可以忽略了，
                    # 操作方式是直接用最后一个有效的bbox替换当前pos上的bbox
                    if boxes[pos, 4] < threshold:
                        boxes[pos,0] = boxes[N-1, 0]
                        boxes[pos,1] = boxes[N-1, 1]
                        boxes[pos,2] = boxes[N-1, 2]
                        boxes[pos,3] = boxes[N-1, 3]
                        boxes[pos,4] = boxes[N-1, 4]
                        N = N - 1           # N-1位置上的bbox已经赋值到前面了，该bbox就可以忽略了；
                        pos = pos - 1       # pos位置上引入了新的有效bbox(N-1)，就需要再计算一遍了

            pos = pos + 1 # 当前pos bbox计算完毕

    # 求满足soft_nms筛选条件的所有bbox数量，并打散为list，但一个问题是：如何与bbox index对应起来？
    # 方式很简单，bbox也做了对应的调整、筛选，bbox list中top-N就对应着最高score，且soft-nms筛选通过的bbox，
    # 不过每个bbox的score也同样经过soft-nms调整了
    keep = [i for i in range(N)]

    return keep