有没有比使用 np.where 更快的方法来迭代一个非常大的 2D numpy 数组？

Posted 2023-02-19

【中文标题】有没有比使用 np.where 更快的方法来迭代一个非常大的 2D numpy 数组？

【英文标题】：Is there a faster method for iterating over a very big 2D numpy array than using np.where?

【发布时间】：2022-01-17 01:43:37

【问题描述】：

我有一个巨大的 2D numpy 数组，其中填充了整数值。我通过 gdal.GetRasterBand() 从 .tif-image 中收集它们。 图像的像素值代表唯一的集群识别号。所以一个簇内的所有像素都具有相同的值。 在我的脚本中，我想检查集群的像素是否超过特定阈值。如果簇大小大于阈值，我想保留簇并给它们一个像素值 1。如果簇的像素小于阈值，则该簇的所有像素都应为 0。

到目前为止，我的代码可以正常工作，但速度非常慢。而且因为我想改变阈值，所以它需要永远。 我将衷心感谢您的帮助。谢谢。

# Import GeoTIFF via GDAL and convert to NumpyArray
data = gdal.Open(image)
raster = data.GetRasterBand(1)
raster = raster.ReadAsArray()

# Different thresholds for iteration
thresh = [0,10,25,50,100,1000,2000]

for threshold in thresh:
        clusteredRaster = np.array(raster.copy(), dtype = int)

        for clump in np.unique(clusteredRaster): # Unique ids of the clusters in image

            if clusteredRaster[np.where(clusteredRaster == clump)].size >= threshold: 
                clusteredRaster[np.where(clusteredRaster == clump)] = int(1)

            else:
                clusteredRaster[np.where(clusteredRaster == clump)] = int(0)
'''

[ClusterImage][1]

In the image you can see the cluster image. Each color stands vor a specific clusternumber. I want to delete the small ones (under a specific size) and just keep the big ones.

  [1]: https://i.stack.imgur.com/miEKg.png

【问题讨论】：

可以将np.unique(clusteredRaster) 移出threshold 循环吗？

很遗憾不是，因为我使用不同的图像，并且每张图像的唯一值都不同

【参考方案1】：

可以进行许多修改以提高性能，

clusteredRaster = np.array(raster.copy(), dtype = int)

可以替换为

clusteredRaster = raster.astype(int)

这本质上是一个复制和转换操作符，所以这个操作更快。

现在

clusteredRaster[np.where(clusteredRaster == clump)] = int(1)

你不需要调用 np.where，这样会更快

clusteredRaster[clusteredRaster == clump] = int(1)

这部分也做了

clusteredRaster[np.where(clusteredRaster == clump)].size

您还可以按以下方式两次删除 clusteredRaster==clump 的评估：

for clump in np.unique(clusteredRaster): # Unique ids of the clusters in image
    indicies = clusteredRaster==clump

    if clusteredRaster[indicies].size >= threshold: 
        clusteredRaster[indicies] = int(1)

    else:
        clusteredRaster[indicies] = int(0)

我认为您的函数现在的运行速度将提高一倍，但是如果您想运行得更快，您必须使用较小的数据类型，如 np.uint8 而不是普通 int，前提是您的图像以 RGB 编码并且可以用 8 表示位整数（或者如果 8 位太低，可能是 np.uint16 ？） 这是从 python 端获得的最快速度。

有更快的方法，例如使用带有 openmp 的 C 模块在多个内核上多线程处理您的工作，这可以通过 numba 或 cython 之类的方法轻松完成，而无需担心编写 C 代码，但是有很多如果您想获得有史以来最快的性能，请阅读要做的事情，例如使用哪个线程后端（TBB 与 openmp）以及一些操作系统和设备相关的功能。

【讨论】：

感谢您的评论！我的迭代时间从 216 秒下降到 46 秒，这已经很有帮助了！不幸的是，我至少需要 uint32，因为我的独特价值会很高......但是我会尝试读入 numba 或 cython。非常感谢

【参考方案2】：

除了changes suggested by Ahmed Mohamed AEK，您还可以在 for 循环之外计算唯一值、索引和计数。另外，您不需要每次都复制raster - 您可以创建一个np.uint8s 的数组。

这给出了与您的原始实现相同的结果：

data = gdal.Open(image)
raster = data.GetRasterBand(1).ReadAsArray()

# Different thresholds for iteration
thresh = [0, 10, 25, 50, 100, 1000, 2000]

# determine the unique clumps and their frequencies outside of the for loops
clumps, counts = np.unique(raster, return_counts=True)

# only determine the indices once, rather than for each threshold
indices = np.asarray([raster==clump for clump in clumps])

for threshold in thresh:
    
    clustered_raster = np.zeros_like(raster, dtype=np.uint8)

    for clump_indices, clump_counts in zip(indices, counts):
    
        clustered_raster[clump_indices] = clump_counts >= threshold

【讨论】：

非常感谢您的回答。根据您的想法，我设法完全避免了一个新的 for 循环，这大大减少了时间！

【参考方案3】：

根据您的有用答案，我得到了一个简单的解决方案！ 这个想法是找到每个阈值的唯一值和集群大小并立即填写正确的值，从而避免循环。 它将迭代时间从最初的每次迭代 142 秒减少到 0.52 秒，并重现了相同的结果。

data = gdal.Open(image)
raster = data.GetRasterBand(1).ReadAsArray()

thresh = [0, 10, 25, 50, 100, 1000, 2000]   
for threshold in thresh:
    # Create new 0-raster with same dimensions as input raster
    clusteredRaster = np.zeros(raster.shape, dtype = uint8)
    
    # Get unique cluster IDs and count the size of the occurence
    clumps, counts = np.unique(raster, return_counts=True)

    # Get only the clumps which are bigger than the threshold
    biggerClumps = clumps[counts >= threshold]

    # fill in ones for the relevant cluster IDs
    clusteredRaster[np.isin(raster,biggerClumps)] = 1