在二维 numpy 数组中有效地查找一组连续值的索引

Question

我计算了图像的分割，其中每个超像素（区域）都由与图像大小相同的 2D 数组中的条目值定义。 我正在尝试获取每个区域的索引列表，以便稍后执行每个区域的操作。 这是我当前的代码：

index_list = []
for i in range(num_superpixels):
    indices = np.where(superpixels == i)
    index_list.append(indices)

以下是一个包含 3 个区域的 3x3 输入的最小示例。 在实践中，我使用从 640x480 图像获得的 500-1000 个超像素，事情变得非常缓慢。

>>> superpixels
array([[0, 0, 2],
       [0, 0, 2],
       [1, 1, 2]])

>>> index_list
      [[array([0, 0, 1, 1]), array([0, 1, 0, 1])],
       [array([2, 2]), array([0, 1])],
       [array([0, 1, 2]), array([2, 2, 2])]]

由于每个区域都是一个连续的块（在 2D 图像中，但不在内存中），因此在循环中使用np.where确实效率低下 - 在每次迭代时，它都会遍历width*height条目以找到约 500 个条目的区域。

我该如何加快速度？

Answer 1

首先，可以基于区域的直接索引设计更好的算法。 实际上，当前代码的复杂度为O(width * height * num_superpixels) ，而有可能达到O(width * height)复杂度。 这个想法是在bin[cellValue]中创建num_superpixels箱和 append 每个单元（二维数组）的位置。

请注意，使用 Python 循环实现它会太慢，但您可以使用Numba加快实现速度。 由于 Numba 不喜欢可变大小的 arrays（效率低下），因此可以应用第一次通过来计算每个 bin 中的单元格数量，然后填充单元格位置。

这是一个例子：

from numba import jit, njit, int32, int64, prange
from numba.types import UniTuple, List

@jit(List(UniTuple(int32[::1],2))(int64[:,::1], int64))
def fastCompute(superpixels, num_superpixels):
    # Count the number of elements
    binSize = np.zeros(num_superpixels, dtype=np.int32)
    for i in range(superpixels.shape[0]):
        for j in range(superpixels.shape[1]):
            binSize[superpixels[i,j]] += 1

    # Put the pixels location in the right bin
    result = [(np.empty(binSize[i], dtype=np.int32), np.empty(binSize[i], dtype=np.int32)) for i in range(num_superpixels)]
    binPos = np.zeros(num_superpixels, dtype=np.int32)
    for i in range(superpixels.shape[0]):
        for j in range(superpixels.shape[1]):
            binIdx = superpixels[i,j]
            tmp = result[binIdx]
            cellBinPos = binPos[binIdx]
            tmp[0][cellBinPos] = i
            tmp[1][cellBinPos] = j
            binPos[binIdx] += 1

    return result

在我的机器上，使用以下基于随机的配置，上面的 function 比初始代码快 120 倍。

# Generate a random input
num_superpixels = 500
superpixels = np.random.randint(np.ones((640, 480)) * num_superpixels)

fastCompute 的fastCompute的类型与初始代码类似（除了出于性能考虑使用元组和 32 位整数），但它不是最优的，因为它包含纯 Python object 类型，并且不是很紧凑。 调整 output 类型应该会产生更快的代码。