python版本的非极大值抑制(nms)的简单实现

因为在深度学习中的目标检测中会检测出多个目标框，后期需要通过非极大值抑制去除得分低并且iou大于阈值的目标框。

因此，在此我们实现了一个简单的nms的python程序，以此作为记录。

nms代码：

# --*-- coding:utf8 --*--
import operator
import numpy as np

def iou(box1, box2):
    x1, y1, w1, h1, s1 = box1
    x2, y2, w2, h2, s2 = box2
    xl = max(x1, x2)
    yl = max(y1, y2)
    xr = min(x1+w1, x2+w2)
    yr = min(y1+h1, y2+h2)
    area = (xr-xl)*(yr-yl)

    return float(area)/float(w1*h1+w2*h2-area)

def nms(boxes, thresh):
    boxes.sort(key=operator.itemgetter(4), reverse=True)
    c = 1
    while c<len(boxes):
        box = boxes[c-1]
        i = c
        while i<len(boxes):
            if iou(box, boxes[i]) >= thresh:
                del boxes[i]
                i -= 1
            i += 1
        c += 1
    return boxes

loc = np.random.randint(1, 50, [50, 2]).tolist()
size = np.random.randint(50, 100, [50, 2]).tolist()
score = [(x+0.4)/1.4 for x in np.random.rand(50)]

record = []

for i in range(50):
    record.append(loc[i]+size[i]+[score[i]])

boxes = nms(record, 0.3)

print(boxes)