目标检测中的LOU(交并比)和NMS(非极大值抑制)代码实现

1、LOU, 两个box框的交集比上并集，示意图如下所示：

 

 代码如下所示：

#假设box1的维度为[N,4] box2的维度为[M,4]
def Lou(box1, box2):
    N = box1.size(0)
    M = box2.size(0)
    #计算两个边框的交集，首先获取相交的左上和右下的坐标
    lt = torch.max(
          box1[:,:2].unsqueeze(1).expend(N, M, 2),
          box2[:,:2].unsqueeze(0).expend(N, M, 2)
    )
    rb = torch.min(
          box1[:,2:].unsqueeze(1).expend(N, M, 2),
          box2[:,2:].unsqueeze(0).expend(N, M, 2)
    )
    wh = rb - lt #(N,M,2)
    wh[wh< 0] = 0
    inter = wh[:,:,0] * wh[:,:,1] #交集面积（N, M）
    
    #计算并集
    #box1面积
    area1 = (box1[:,2] - box1[:, 0]) * (box1[:,3] - box1[:,1]) #面积（N）
    area2 = (box2[:,2] - box2[:, 0]) * (box2[:,3] - box2[:,1]) #面积（M）
    area1 = area1.unsqueeze(1).expend(N,M)
    area2 = area2.unsqueeze(0).expend(N,M) 
    
    lou = inter/ (area1+area2-inter)
    return lou
    

2.NMS(非极大值抑制)

过程如下：

1 将识别模块产生的box按照预测的值（score）进行降序排列

2 然后选择score最大值的框，并将其置为当前box，保存idex，然后依次后面的box计算与当前box的IOU值，若大于阈值，则抑制，扔掉，

3 完成一轮遍历后，继续选择下一个非抑制的box作为当前box，重复2步骤，直至识别模块预测的序列为空。

4 返回没有被抑制的index即符合条件的box

代码如下所示：

def NMS(detect, score,threshold=0.5):
    #detect(N, 4)表示检测框
    #score表示预测值
   #依次把预测框的四个值取出来，以便后面的计算
    x1 = detect[:, 0]
    x2 = detect[:, 2]
    y1 = detect[:, 1]
    y2 = detect[:, 3]
    area = (x2-x1) * (y2-y1)  #(N,)

    #用于存放保存到框
    keep = []
    
    index = score.argsort()[::-1]

    while index.size > 0:
        i = index[0]
        keep.append(i)
        x11 = np.maximum(x1[i], x1[index[1:]])
        y11 = np.maximum(y1[i], y1[index[1:]])
        x22 = np.minimun(x2[i], x2[index[1:]])
        y22 = np.minimum(y2[i], y2[index[1:]])

        w = np.maximum(0, x22-x11+1)
        h = np.maximum(0, y22-y11+1)
        
        overlaps = w*h
        lous = overlaps/(area[i] + area[index[1:]] - overlaps)

        idx = np.where(lous <= threshold)[0]    
        index = index[idx+1]
return keep