计算交并比

边界框的面积交并比

import numpy as np

def compute_iou(box_a, box_b):
    box_a = np.array(box_a) if isinstance(box_a, list) else box_a
    box_b = np.array(box_b) if isinstance(box_b, list) else box_b
    box_a = box_a[None, :] if box_a.ndim == 1 else box_a
    box_b = box_b[None, :] if box_b.ndim == 1 else box_b
    assert box_a.ndim == 2 and box_a.shape[1] == 4
    assert box_b.ndim == 2 and box_b.shape[1] == 4

    inter_lt = np.maximum(box_a[:, None, 0:2], box_b[:, 0:2])
    inter_rb = np.minimum(box_a[:, None, 2:4], box_b[:, 2:4])
    inter_wh = np.maximum(inter_rb - inter_lt + 1, 0)
    inter = inter_wh[:, :, 0] * inter_wh[:, :, 1]
    area_a = (box_a[:, 2] - box_a[:, 0] + 1) * (box_a[:, 3] - box_a[:, 1] + 1)
    area_b = (box_b[:, 2] - box_b[:, 0] + 1) * (box_b[:, 3] - box_b[:, 1] + 1)
    iou = inter / (area_a[:, None] + area_b - inter)
    return iou

if __name__ == '__main__':
    box_a = [[0, 0, 1, 1], [0, 0, 2, 2], [0, 0, 3, 3]]
    box_b = [[0, 0, 2, 2], [0, 0, 3, 3], [0, 0, 4, 4], [0, 0, 5, 5]]
    iou = compute_iou(box_a, box_b)
    print(iou)