如何使用 numpy 和 pytorch 快速计算 IOU

前言

在目标检测中用交并比（Interection-over-unio，简称 IOU）来衡量两个边界框之间的重叠程度，下面就使用 numpy 和 pytorch 两种框架的矢量计算方式来快速计算各种情况下的 IOU。

一对一

先来计算最简单的单框对单框的交并比。假设两个预测框 \(b_0=(x_{min0},\ y_{min0},\ x_{max0},\ y_{max0})\) 和 \(b_1=(x_{min1},\ y_{min1},\ x_{max1},\ y_{max1})\) ，固定 \(b_0\) 的位置不变，移动 \(b_1\)，他们之间会有四种重叠情况，如下图所示，此时交并比计算公式为 \(IOU=C/(A+B-C)\)，就是交集面积除以并集面积。虽然图中有四种重叠情况，但是计算的时候可以合并为一种 \(C=w_c*h_c\)：

• 交集 \(C\) 的宽度 \(w_c=x_2-x_1\)，其中 \(x_2=\min\{x_{max0},\ x_{max1}\}\)，\(x_1=\max\{x_{min0},\ x_{min1} \}\)；

• 交集 \(C\) 的高度 \(h_c=y_2-y_1\)，其中 \(y_2=\min\{y_{max0},\ y_{max1}\}\)，\(y_1=\max\{y_{min0},\ y_{min1} \}\)；

numpy

def iou(box0: np.ndarray, box1: np.ndarray):
    """ 计算一对一交并比

    Parameters
    ----------
    box0, box1: `~np.ndarray` of shape `(4, )`
        边界框
    """
    xy_max = np.minimum(box0[2:], box1[2:])
    xy_min = np.maximum(box0[:2], box1[:2])

    # 计算交集
    inter = np.clip(xy_max-xy_min, a_min=0, a_max=np.inf)
    inter = inter[0]*inter[1]

    # 计算并集
    area_0 = (box0[2]-box0[0])*(box0[3]-box0[1])
    area_1 = (box1[2]-box1[0])*(box1[3]-box1[1])
    union = area_0 + area_1- inter

    return inter/union

pytorch

def iou(box0: torch.Tensor, box1: torch.Tensor):
    """ 计算一对一交并比

    Parameters
    ----------
    box0, box1: Tensor of shape `(4, )`
        边界框
    """
    xy_max = torch.min(box0[2:], box1[2:])
    xy_min = torch.max(box0[:2], box1[:2])

    # 计算交集
    inter = torch.clamp(xy_max-xy_min, min=0)
    inter = inter[0]*inter[1]

    # 计算并集
    area_0 = (box0[2]-box0[0])*(box0[3]-box0[1])
    area_1 = (box1[2]-box1[0])*(box1[3]-box1[1])
    union = area_0 + area_1- inter

    return inter/union

一对多

一对多的代码和一对一几乎是一模一样的，就是有一个参数多了一个维度而已。

numpy

def iou(box: np.ndarray, boxes: np.ndarray):
    """ 计算一个边界框和多个边界框的交并比

    Parameters
    ----------
    box: `~np.ndarray` of shape `(4, )`
        边界框

    boxes: `~np.ndarray` of shape `(n, 4)`
        其他边界框

    Returns
    -------
    iou: `~np.ndarray` of shape `(n, )`
        交并比
    """
    # 计算交集
    xy_max = np.minimum(boxes[:, 2:], box[2:])
    xy_min = np.maximum(boxes[:, :2], box[:2])
    inter = np.clip(xy_max-xy_min, a_min=0, a_max=np.inf)
    inter = inter[:, 0]*inter[:, 1]

    # 计算面积
    area_boxes = (boxes[:, 2]-boxes[:, 0])*(boxes[:, 3]-boxes[:, 1])
    area_box = (box[2]-box[0])*(box[3]-box[1])

    return inter/(area_box+area_boxes-inter)

pytorch

def iou(box: torch.Tensor, boxes: torch.Tensor):
    """ 计算一个边界框和多个边界框的交并比

    Parameters
    ----------
    box: Tensor of shape `(4, )`
        一个边界框

    boxes: Tensor of shape `(n, 4)`
        多个边界框

    Returns
    -------
    iou: Tensor of shape `(n, )`
        交并比
    """
    # 计算交集
    xy_max = torch.min(boxes[:, 2:], box[2:])
    xy_min = torch.max(boxes[:, :2], box[:2])
    inter = torch.clamp(xy_max-xy_min, min=0)
    inter = inter[:, 0]*inter[:, 1]

    # 计算并集
    area_boxes = (boxes[:, 2]-boxes[:, 0])*(boxes[:, 3]-boxes[:, 1])
    area_box = (box[2]-box[0])*(box[3]-box[1])

    return inter/(area_box+area_boxes-inter)

多对多

假设边界框矩阵 boxes0 的维度为 (A, 4)，boxes1 的维度为 (B, 4)，要计算这二者间的交并比，一种直观的想法是开个循环，一次从 boxes0 中拿出一个边界框 box，然后和 boxes1 里面的所有边界框计算交并比。但是在边界框很多的情况下，这种计算是很低效的。有没有什么办法可以代替循环呢？

来仔细思考一下：开循环的原因是我们不能像之前那样对两个边界框矩阵进行切片，然后计算 \([x_2, y_2]\) 以及 \([x_1, y_1]\)。所以我们只要想个办法能像之前那样切片就好了。 这时候就需要使用广播机制，将 boxes0 广播为 (A, B, 4) 维度的矩阵，boxes1 也广播为 (A, B, 4) 维度的矩阵。广播之后的 boxes0 任意沿着某行切下来的矩阵 boxes0[:, i, :]都是相同的（就是广播前的 boxes0），广播之后的 boxes1 沿着某一个通道切下来的矩阵 boxes1[i, :, :] 也是相同的。

numpy

def iou(boxes0: np.ndarray, boxes1: np.ndarray):
    """ 计算多个边界框和多个边界框的交并比

    Parameters
    ----------
    boxes0: `~np.ndarray` of shape `(A, 4)`
        边界框

    boxes1: `~np.ndarray` of shape `(B, 4)`
        边界框

    Returns
    -------
    iou: `~np.ndarray` of shape `(A, B)`
        交并比
    """
    A = boxes0.shape[0]
    B = boxes1.shape[0]

    xy_max = np.minimum(boxes0[:, np.newaxis, 2:].repeat(B, axis=1),
                        np.broadcast_to(boxes1[:, 2:], (A, B, 2)))
    xy_min = np.maximum(boxes0[:, np.newaxis, :2].repeat(B, axis=1),
                        np.broadcast_to(boxes1[:, :2], (A, B, 2)))

    # 计算交集面积
    inter = np.clip(xy_max-xy_min, a_min=0, a_max=np.inf)
    inter = inter[:, :, 0]*inter[:, :, 1]

    # 计算每个矩阵的面积
    area_0 = ((boxes0[:, 2]-boxes0[:, 0])*(
        boxes0[:, 3] - boxes0[:, 1]))[:, np.newaxis].repeat(B, axis=1)
    area_1 = ((boxes1[:, 2] - boxes1[:, 0])*(
        boxes1[:, 3] - boxes1[:, 1]))[np.newaxis, :].repeat(A, axis=0)

    return inter/(area_0+area_1-inter)

pytorch

def iou(boxes0: Tensor, boxes1: Tensor):
    """ 计算多个边界框和多个边界框的交并比

    Parameters
    ----------
    boxes0: Tensor of shape `(A, 4)`
        边界框

    boxes1: Tensor of shape  `(B, 4)`
        边界框

    Returns
    -------
    iou: Tensor of shape `(A, B)`
        交并比
    """
    A = boxes0.size(0)
    B = boxes1.size(0)

    # 将先验框和边界框真值的 xmax、ymax 以及 xmin、ymin进行广播使得维度一致，(A, B, 2)
    # 再计算 xmax 和 ymin 较小者、xmin 和 ymin 较大者，W=xmax较小-xmin较大，H=ymax较小-ymin较大
    xy_max = torch.min(boxes0[:, 2:].unsqueeze(1).expand(A, B, 2),
                       boxes1[:, 2:].broadcast_to(A, B, 2))
    xy_min = torch.max(boxes0[:, :2].unsqueeze(1).expand(A, B, 2),
                       boxes1[:, :2].broadcast_to(A, B, 2))

    # 计算交集面积
    inter = (xy_max-xy_min).clamp(min=0)
    inter = inter[:, :, 0]*inter[:, :, 1]

    # 计算每个矩形的面积
    area_boxes0 = ((boxes0[:, 2]-boxes0[:, 0]) *
                  (boxes0[:, 3]-boxes0[:, 1])).unsqueeze(1).expand(A, B)
    area_boxes1 = ((boxes1[:, 2]-boxes1[:, 0]) *
                 (boxes1[:, 3]-boxes1[:, 1])).broadcast_to(A, B)

    return inter/(area_boxes0+area_boxes1-inter)