python实现nms

import numpy as np
import operator
def iou(box1, box2):
    x1, y1, w1, h1, s1 = box1
    x2, y2, w2, h2, s2 = box2
    x3 = max(x1, x2)
    y3 = max(y1, y2)
    x4 = min(x1+w1, x2+w2)
    y4 = min(y1+h1, y2+h2)
    m=max(0.,x4-x3)
    n=max(0.,y4-y3)
    jiao=m*n
    bing=w1*h1+w2*h2-jiao
    return jiao/bing
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, [10, 2]).tolist()
size = np.random.randint(50, 100, [10, 2]).tolist()
score = [x for x in np.random.rand(10)]
record = []
for i in range(10):
    record.append(loc[i]+size[i]+[score[i]])
boxes = nms(record, 0.3)
print(boxes)