有效地为二维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回收箱，并在bin[cellValue]中附加每个单元格( 2D数组)的位置。

请注意，用Python循环实现这个过程太慢了，但是您可以使用Numba来加快实现速度。由于Numba不喜欢可变大小的数组(低效)，所以可以先使用第一遍来计算每个桶中的单元格数，然后填充单元格位置。

下面是一个示例：

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对象类型，而且并不是非常紧凑。调整输出类型应该会产生更快的代码。