【pytorch】实现1对多的IOU计算技巧

IOU计算中常出现类似情况，例如求某个真值框，相对于所有预测框的IOU，即1对多的IOU计算。

对于坐标系，x轴向右；y轴向下。

其pytorch代码为：

def calculate_iou(box_as, box_bs):
    '''
    实现box_as对box_bs的1对多IOU计算
    输入：
    box_as:size:[A,4] 其中4为x1,y1,x2,y2;即box的左上角及右下角
    box_bs:size:[B,4]
    输出：
    ious:size:[A,B,1]:其意义为，每个A对应B个IOU
    '''
    #首先进行矩阵扩张操作：
    #意义为，每个A中的重复了B次
    box_asExtend=box_as.unsqueeze(1).expand(box_as.shape[0],box_bs.shape[0],4)
    #意义为A个[B,4]
    box_bsExtend = box_bs.unsqueeze(0).expand(box_as.shape[0], box_bs.shape[0], 4)
    #经过上面的变换，1对多IOU计算已经转换为1对1IOU计算：
    #此时再使用1对1 IOU变换对进行计算：box_asExtend与box_bsExtend间的IOU值：
    box1 = box_asExtend
    box2 = box_bsExtend
    # 计算两框相并的面积b1AddB2area
    leftTop = torch.max(box1[..., 0:2], box2[..., 0:2])
    bottomRight = torch.min(box1[..., 2:4], box2[..., 2:4])
    # 如出现负值表示框不相交截取为0:输出[...,w,h]
    b1AndB2 = torch.clamp(bottomRight - leftTop, min=0)
    # 计算了相交面积，形状[...,1]
    b1AndB2Area = b1AndB2[..., 0:1] * b1AndB2[..., 1:2]

    #计算两个box各自的面积：
    # 再计算b1及b2自己的面积：形状[...,1]
    b1Area = (box1[...,2:3]-box1[...,0:1])*(box1[...,3:4]-box1[...,1:2])
    b2Area = (box2[...,2:3]-box2[...,0:1])*(box2[...,3:4]-box2[...,1:2])
    # 得出IOU的结果
    return b1AndB2Area / (b1Area + b2Area - b1AndB2Area)
#测试：
b1=torch.FloatTensor([[0,0,1,1],[1,0,2,1]])
b2=torch.FloatTensor([[0.5,0.5,1.5,1.5]])
print(calculate_iou(b1,b2))

输出结果为：

tensor([[[0.1429]],
        [[0.1429]]])