python 计算两个盒子之间的IoU

Posted 2021-05-08

def intersect(box_a, box_b):
    """
    compute the intersection area between two bounding boxes
    :param box_a: (tensor) bounding boxes， shape [A, 4]
    :param box_b: (tensor) bounding boxes， shape [B, 4]
    :return:（tensor）intersection area, shape [A, B]
    """
    A = box_a.size(0)
    B = box_b.size(0)

    max_xy = torch.min(box_a[:, 2:].unsqueeze(1).expand(A, B, 2), # shape [A,2] -> [A,1,2] -> [A,B,2]
                       box_b[:, 2:].unsqueeze(0).expand(A, B, 2)) # shape [B,2] -> [1,B,2] -> [A,B,2]
    min_xy = torch.max(box_a[:, :2].unsqueeze(1).expand(A, B, 2),
                       box_b[:, :2].unsqueeze(0).expand(A, B, 2))
    inter = torch.clamp((max_xy - min_xy), min=0)
    return inter[:, :, 0] * inter[:, :, 1]


def jaccard(box_a, box_b):
    """
    Compute the iou of box_a and box_b
    :param box_a: (tensor) bounding boxes， shape [A, 4]
    :param box_b: (tensor) bounding boxes， shape [B, 4]
    :return: (tensor) iou of bounding boxes shape [A, B]
    """
    inter = intersect(box_a, box_b)
    area_a = ((box_a[:, 2]-box_a[:, 0]) *
              (box_a[:, 3]-box_a[:, 1])).unsqueeze(1).expand_as(inter)  # [A,B]
    area_b = ((box_b[:, 2]-box_b[:, 0]) *
              (box_b[:, 3]-box_b[:, 1])).unsqueeze(0).expand_as(inter)  # [A,B]
    union = area_a + area_b - inter
    return inter / union  # [A,B]