IOU计算，Numpy 和 Torch 版本

简介

IOU计算一直是目标检测中最重要的一个环节。虽然iou在数学上定义很简单，但是想大规模计算还是有点复杂，我自己利用numpy和torch库仔细写了一下iou算法，从而加深对iou计算的印象。

Numpy 版本

import numpy as np
def get_iou(a_boxs,gt_boxs):
    '''
    Args:
        a_boxs (N, 4): predicted boxes.
        gt_boxs (N, 4): ground truth boxes.
    Returns:
        iou (N): IoU values.
    '''
    k = gt_boxs.shape[0]
    iou = []
    # 由于一张图像的gt_boxs数量远小于anchors，故遍历gt_boxs
    for i in range(k):
        iou.append(iou_ab(a_boxs,gt_boxs[i][np.newaxis,:]))
    iou = np.array(iou)
    print(iou)

def iou_ab(a_boxs,gt_box,eps=1e-5):
    '''
    Args:
        a_boxs (N, 4): predicted boxes.
        gt_box (1, 4): ground truth boxe.
        eps: a small number to avoid 0 as denominator.
    Returns:
        iou (N): IoU value.
    '''
    top_left = np.maximum(a_boxs[...,:2],gt_box[...,:2])
    # 将小于0 的转为 0
    bottom_right = np.minimum(a_boxs[...,2:],gt_box[...,2:])
    hw = np.clip(bottom_right-top_left,a_min=0,a_max=None)
    over_lap = hw[...,0]*hw[...,1]
    area_a = np.abs(a_boxs[:,0]-a_boxs[:,2])*np.abs(a_boxs[:,1]-a_boxs[:,3])
    area_b = np.abs(gt_box[:,0]-gt_box[:,2])*np.abs(gt_box[:,1]-gt_box[:,3])
    iou_area = over_lap/(area_a+area_b-over_lap+eps)
    return iou_area

if __name__ == '__main__':
    a_boxs = np.load('bbox_a.npy')
    b_boxs = np.load('bbox_b2.npy')
    get_iou(a_boxs,b_boxs)  

numpy 改进版

def iou_abs(a_boxs,gt_boxs,eps=1e-5):
    '''
    Args:
        a_boxs (N,1,4): predicted boxes.
        gt_box (1,M,4): ground truth boxe.
        eps: a small number to avoid 0 as denominator.
    Returns:
        iou (N): IoU value.
    '''
    top_left = np.maximum(a_boxs[...,:2],gt_boxs[...,:2])
    # 将小于0 的转为 0
    bottom_right = np.minimum(a_boxs[...,2:],gt_boxs[...,2:])
    hw = np.clip(bottom_right-top_left,a_min=0,a_max=None)
    over_lap = hw[...,0]*hw[...,1]
    area_a = np.abs(a_boxs[...,0]-a_boxs[...,2])*np.abs(a_boxs[...,1]-a_boxs[...,3])
    area_b = np.abs(gt_boxs[...,0]-gt_boxs[...,2])*np.abs(gt_boxs[...,1]-gt_boxs[...,3])
    iou_area = over_lap/(area_a+area_b-over_lap+eps)
    return iou_area

if __name__ == '__main__':
    a_boxs = np.load('bbox_a.npy')
    gt_boxs = np.load('bbox_b2.npy')
    ious = iou_abs(a_boxs[:,np.newaxis,:], gt_boxs[np.newaxis,:])
    print(ious)

torch 版本

def area_of(left_top, right_bottom) -> torch.Tensor:
    """Compute the areas of rectangles given two corners.

    Args:
        left_top (N,M,2): left top corner.
        right_bottom (N,M,2): right bottom corner.

    Returns:
        area (N): return the area.
    """
    hw = torch.clamp(right_bottom - left_top, min=0.0)
    return hw[..., 0] * hw[..., 1]


def iou_of(gt_boxs, boxes1, eps=1e-5):
    """Return intersection-over-union (Jaccard index) of boxes.

    Args:
        gt_boxs (1, M, 4): ground truth boxes.
        boxes1 (N, 1, 4): predicted boxes.
        eps: a small number to avoid 0 as denominator.
    Returns:
        iou (N): IoU values.
    """
    overlap_left_top = torch.max(gt_boxs[..., :2], boxes1[..., :2])
    overlap_right_bottom = torch.min(gt_boxs[..., 2:], boxes1[..., 2:])

    overlap_area = area_of(overlap_left_top, overlap_right_bottom)
    area0 = area_of(gt_boxs[..., :2], gt_boxs[..., 2:])
    area1 = area_of(boxes1[..., :2], boxes1[..., 2:])
    return overlap_area / (area0 + area1 - overlap_area + eps)

if __name__ == '__main__':
    a_boxs = np.load('bbox_a.npy')
    gt_boxs = np.load('bbox_b2.npy')
    a_boxs = torch.from_numpy(a_boxs)
    gt_boxs = torch.from_numpy(gt_boxs)
    iou = iou_of(gt_boxs.unsqueeze(0),a_boxs.unsqueeze(1))
    print(iou)