iou einops

iou_einops.py

def area(a, num_classes):
    class_idx = einops.rearrange(torch.arange(num_classes), 'i -> i 1 1 1')
    a = a == class_idx
    area = einops.reduce(a, 'c ... -> c', reduction='sum')
    return area


def intersection(a, b, num_classes) -> float:
    class_idx = einops.rearrange(torch.arange(num_classes), 'i -> i 1 1 1')
    a = a == class_idx
    b = b == class_idx
    intersection = torch.logical_and(a, b)
    intersection = einops.reduce(
        intersection, 'c ... -> c', reduction='sum'
    )
    return intersection


def union(a, b, num_classes: int) -> float:
    class_idx = einops.rearrange(torch.arange(num_classes), 'i -> i 1 1 1')
    a = a == class_idx
    b = b == class_idx
    union = torch.logical_or(a, b)
    union = einops.reduce(
        union, 'c ... -> c', reduction='sum'
    )
    return union