如何使用 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)