Compute IOU and Non Maximum Suppression in numpy

注意bbox的格式是左下角的点和右上角的点

# bbox: [x1,y1,x2,y2] -> [bottom_left,upper_right]
def CalIOU(bbox1,bbox2):
    xx1 = max(bbox1[0],bbox2[0])
    yy1 = max(bbox1[1],bbox2[1])
    xx2 = min(bbox1[2],bbox2[2])
    yy2 = min(bbox1[3],bbox2[3])
    
    w = max(0,xx2-xx1)
    h = max(0,yy2-yy1)
    inter = w*h 
    area = (bbox1[2]-bbox1[0]) * (bbox1[3]-bbox1[1]) + (bbox2[2]-bbox2[0]) * (bbox2[3]-bbox2[1])
    iou = inter / (area - inter)
    return iou 

# bbox1 = [0,0,2,2]
# bbox2 = [1,1,3,3]
# iou = CalIOU(bbox1,bbox2)
# iou = CalIOU(bbox2,bbox1)
# print(iou) 

# bbox1 = [0,1,2,3]
# bbox2 = [1,0,3,2]
# iou = CalIOU(bbox1,bbox2)
# iou = CalIOU(bbox2,bbox1)
# print(iou) 

bbox1 = [0,0,2,2]
bbox2 = [3,3,5,5]
iou = CalIOU(bbox1,bbox2)
# iou = CalIOU(bbox2,bbox1)
print(iou) 

import numpy as np 
# dets: [x1,y1,x2,y2,score] -> [bottom_left,upper_right]
def nms(dets,thresh):
    x1 = dets[:,0]
    y1 = dets[:,1]
    x2 = dets[:,2]
    y2 = dets[:,3]
    scores = dets[:,4]
    
    areas = (x2-x1+1) * (y2-y1+1)
    order = scores.argsort()[::-1]
    
    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i],x1[order[1:]])
        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 = np.maximum(0.0,xx2-xx1+1)
        h = np.maximum(0,0,yy2-yy1+1)
        
        inter = w * h 
        iou = inter / (areas[i] + areas[order[1:]] - inter)
        
        inds = np.where(iou <= thresh)[0]
        order = order[inds+1]
    return keep

出处：https://github.com/rbgirshick/fast-rcnn/blob/master/lib/utils/nms.py