python编程之NMS

非极大值抑制（NMS）
NMS的算法步骤如下：

# INPUT：所有预测出的bounding box (bbx)信息（坐标和置信度confidence），　IOU阈值（大于该阈值的bbx将被移除）
for object in all objects:
    (1) 获取当前目标类别下所有bbx的信息
    (2) 将bbx按照confidence从高到低排序,并记录当前confidence最大的bbx
    (3) 计算最大confidence对应的bbx与剩下所有的bbx的IOU,移除所有大于IOU阈值的bbx
    (4) 对剩下的bbx，循环执行(2)和(3)直到所有的bbx均满足要求（即不能再移除bbx）
注意的是，NMS是对所有的类别分别执行的。举个栗子，假设最后预测出的矩形框有2类（分别为cup, pen）,在NMS之前，每个类别可能都会有不只一个bbx被预测出来，这个时候我们需要对这两个类别分别执行一次NMS过程。 
我们用python编写NMS代码，假设对于一张图片，所有的bbx信息已经保存在一个字典中，保存形式如下：
 

def compute_iou(x1,y1,x2,y2,xx1,yy1,xx2,yy2):
    s1=(x2-x1+1)*(y2-y1+1)
    s2=(xx2-xx1+1)*(yy2-yy1+1)
    xxx1=max(x1,xx1)
    yyy1=max(y1,yy1)
    xxx2=min(x2,xx2)
    yyy2=min(y2,yy2)
    ins= max(0,xxx2-xxx1)*max(0,yyy2-yyy1)
    return ins/(s1+s2-ins+1e-6)

def NMS(dic,threshold):
    for class_name,bbox in dic.items():
        bbox.sort(key=lambda x:x[-1],reverse=1)
        res = []
        while bbox:
            cur = bbox.pop(0)
            res.append(cur)
            for i in bbox:
                iou = compute_iou(cur[0],cur[1],cur[2],cur[3],i[0],i[1],i[2],i[3])
                if iou>threshold:
                    bbox.remove(i)
        return res

print(NMS({'cup': [[59, 120, 137, 368, 0.124648176], [221, 89, 369, 367, 0.35818103], [54, 154, 148, 382, 0.13638769]]},0.5))