目标检测中IOU的计算

IOU计算一般以四角顶点坐标计算，当目标框以中心点定位时，需转换为顶点表示的形式，代码如下

import torch


def iou(box1, box2, is_vertex=True):
    # 如果检测框坐标点以中心点表示，则将其变换为以顶点表示
    if not is_vertex:
        # [x, y, w, h] -> [x1, y1, x2, y2]
        box1 = [box1[0] - box1[2], box1[1] + box1[3], box1[0] + box1[2], box1[1] - box1[3]]
        box2 = [box2[0] - box2[2], box2[1] + box2[3], box2[0] + box2[2], box2[1] - box2[3]]

    # 计算交集左上角坐标和右下角坐标
    left_top = [max(box1[0], box2[0]), min(box1[1], box2[1])]
    right_bottom = [min(box1[2], box2[2]), max(box1[3], box2[3])]

    # 计算交集面积
    inter = abs(left_top[0] - right_bottom[0]) * abs(left_top[1] - right_bottom[1])

    # 计算并集面积
    union = abs(box1[2] - box1[0]) * abs(box1[3] - box1[1]) + abs(box2[2] - box2[0]) * abs(box2[3] - box2[1]) - inter

    return inter / union


# 两相同目标框，一个以顶点坐标表示，另一个以中心点坐标表示
boxA1 = torch.tensor([1, 6, 5, 2])
boxA2 = torch.tensor([3, 4, 7, 0])

boxB1 = torch.tensor([3, 4, 2, 2])
boxB2 = torch.tensor([5, 2, 2, 2])

resA = iou(boxA1, boxA2)
resB = iou(boxB1, boxB2, is_vertex=False)

print("以顶点表示的IOU：", resA)
print("以中心点表示的IOU：", resB)



# >> 以顶点表示的IOU： tensor(0.1429)
#    以中心点表示的IOU： tensor(0.1429)