在 2D 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) 的复杂度。这个想法是创建num_superpixels bins 并在bin[cellValue] 中附加每个单元格（二维数组）的位置。

请注意，使用 Python 循环实现它会太慢，但您可以使用 Numba 加快实现速度。由于 Numba 不喜欢可变大小的数组（效率低下），因此可以应用第一遍来计算每个 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

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

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

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