手撕NMS和IOU代码(python)

手撕NMS代码(python)

import numpy as np
def nms(dets,score):
    x1 = dets[:,0]
    y1 = dets[:,1]
    x2 = dets[:,2]
    y2 = dets[:,3]
    area = (x2-x1+1)*(y2-y1+1)
    scores = dets[:,4]
    idx = scores.argsort()[::-1]
    keep = []
    while idx.size > 0:
        i = idx[0]
        keep.append(i)
        xx1 = np.maximum(x1[i],x1[idx[1:]])
        xx2 = np.minimum(x2[i],x2[idx[1:]])
        yy1 = np.maximum(y1[i],y1[idx[1:]])
        yy2 = np.minimum(y2[i],y2[idx[1:]])
        w = np.maximum(xx2-xx1,0)
        h = np.maximum(yy2-yy1,0)
        overlap = w*h
        iou = overlap/(area[i]+area[idx[1:]]-overlap)
        idx_i = np.where(iou<=score_thr)[0]
        idx = idx[idx_i+1]
    return keep

手撕IOU代码

import numpy as np

def ComputeIOU(boxA, boxB):
    ## 计算相交框的坐标
    x1 = np.max([boxA[0], boxB[0]])
    x2 = np.min([boxA[2], boxB[2]])
    y1 = np.max([boxA[1], boxB[1]])
    y2 = np.min([boxA[3], boxB[3]])
    
    ## 计算交区域，并区域，及IOU
    interArea = np.max([x2-x1+1, 0])*np.max([y2-y1+1,0])	##一定要和0比较大小，如果是负数就说明压根不相交
    unionArea = (boxA[2]-boxA[0]+1)*(boxA[3]-boxA[1]+1) + (boxB[2]-boxB[0]+1)*(boxB[3]-boxB[1]+1)-interArea
    iou = interArea/unionArea
    return iou

boxA = [1,1,3,3]
boxB = [2,2,4,4]
IOU = ComputeIOU(boxA, boxB)