高效的类距离矩阵计算（手动度量函数）

Question

我想计算一个“距离”矩阵，类似于scipy.spatial.distance.cdist，但使用“边界框”（4 维向量）之间的并集交集 (IoU)，而不是典型的距离度量（如欧几里得距离）。

例如，假设我们有两个边界框集合，例如

import numpy as np
A_bboxes = np.array([[0, 0, 10, 10], [5, 5, 15, 15]])
array([[ 0,  0, 10, 10],
       [ 5,  5, 15, 15]])

B_bboxes = np.array([[1, 1, 11, 11], [4, 4, 13, 13], [9, 9, 13, 13]])
array([[ 1,  1, 11, 11],
       [ 4,  4, 13, 13],
       [ 9,  9, 13, 13]])

我想计算一个矩阵J，它的第 {i,j} 个元素将保存A_bboxes 的第 i 个 bbox 和 B_bboxes 的第 j 个 bbox 之间的 IoU。

给定以下用于计算两个给定 bbox 之间的 IoU 的函数：

def compute_iou(bbox_a, bbox_b):
    xA = max(bbox_a[0], bbox_b[0])
    yA = max(bbox_a[1], bbox_b[1])
    xB = min(bbox_a[2], bbox_b[2])
    yB = min(bbox_a[3], bbox_b[3])

    interArea = max(0, xB - xA + 1) * max(0, yB - yA + 1)
    boxAArea = (bbox_a[2] - bbox_a[0] + 1) * (bbox_a[3] - bbox_a[1] + 1)
    boxBArea = (bbox_b[2] - bbox_b[0] + 1) * (bbox_b[3] - bbox_b[1] + 1)

    iou = interArea / float(boxAArea + boxBArea - interArea)

    return iou

IoU 矩阵可以如下计算：

J = np.zeros((A_bboxes.shape[0], B_bboxes.shape[0])) 
for i in range(A_bboxes.shape[0]): 
   for j in range(B_bboxes.shape[0]): 
      J[i, j] = compute_iou(A_bboxes[i], B_bboxes[j])

导致：

J = array([[0.70422535, 0.28488372, 0.02816901],
           [0.25388601, 0.57857143, 0.20661157]])

现在，我也想做同样的事情，但没有使用双for循环。我知道scipy.spatial.distance.cdist 可以为用户定义的 2-arity 函数执行类似的任务，例如：

dm = cdist(XA, XB, lambda u, v: np.sqrt(((u-v)**2).sum()))

但是，我看不出如何将 IoU 的计算嵌入到 lambda 表达式中。有什么方法可以做到这一点，甚至有什么不同的方法可以避免 lambda 函数？

编辑：回答

使用 lambda 形式嵌入 IoU 的计算似乎真的很容易。解决方法如下：

J = cdist(A_bboxes, B_bboxes, lambda u, v: compute_iou(u, v)))
J = array([[0.70422535, 0.28488372, 0.02816901],
           [0.25388601, 0.57857143, 0.20661157]])

Answer 1

如果您想要一个高性能的解决方案，您可以使用 cython 或 numba。通过这两种方法，它的性能比您的 cdist 方法高出 3 个数量级。

模板函数

对于其他距离函数，例如 Minkowski 距离，我在一个月前写了一个 answer。

import numpy as np
import numba as nb
from scipy.spatial.distance import cdist

def gen_cust_dist_func(kernel,parallel=True):

    kernel_nb=nb.njit(kernel,fastmath=True)

    def cust_dot_T(A,B):
        assert B.shape[1]==A.shape[1]

        out=np.empty((A.shape[0],B.shape[0]),dtype=np.float64)
        for i in nb.prange(A.shape[0]):
            for j in range(B.shape[0]):
                out[i,j]=kernel_nb(A[i,:],B[j,:])
        return out

    if parallel==True:
        return nb.njit(cust_dot_T,fastmath=True,parallel=True)
    else:
        return nb.njit(cust_dot_T,fastmath=True,parallel=False)

自定义函数示例

def compute_iou(bbox_a, bbox_b):
    xA = max(bbox_a[0], bbox_b[0])
    yA = max(bbox_a[1], bbox_b[1])
    xB = min(bbox_a[2], bbox_b[2])
    yB = min(bbox_a[3], bbox_b[3])

    interArea = max(0, xB - xA + 1) * max(0, yB - yA + 1)
    boxAArea = (bbox_a[2] - bbox_a[0] + 1) * (bbox_a[3] - bbox_a[1] + 1)
    boxBArea = (bbox_b[2] - bbox_b[0] + 1) * (bbox_b[3] - bbox_b[1] + 1)

    iou = interArea / float(boxAArea + boxBArea - interArea)

    return iou

#generarte custom distance function
cust_dist=gen_cust_dist_func(compute_iou,parallel=True)


A_bboxes = np.array([[0, 0, 10, 10], [5, 5, 15, 15]]*100)
B_bboxes = np.array([[1, 1, 11, 11], [4, 4, 13, 13], [9, 9, 13, 13]]*1000)

%timeit cust_dist(A_bboxes,B_bboxes)
#1.74 ms ± 13.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit cdist(A_bboxes, B_bboxes, lambda u, v: compute_iou(u, v))
#3.33 s ± 11.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)