有沒有辦法加快循環numpy.where？

Question

想象一下，您有一個分割圖，其中每個對象都由一個唯一索引標識，例如看起來類似於：

對於每個對象，我想保存它覆蓋的像素，但到目前為止我只能想出標准for循環。 不幸的是，對於包含數千個單獨對象的較大圖像，這會變得非常緩慢——至少對於我的真實數據而言。 我能以某種方式加快速度嗎？

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from skimage.draw import random_shapes


# please ignore that this does not always produce 20 objects each with a
# unique color. it is simply a quick way to produce data that is similar to
# my problem that can also be visualized.
segmap, labels = random_shapes(
    (100, 100), 20, min_size=6, max_size=20, multichannel=False,
    intensity_range=(0, 20), num_trials=100,
)
segmap = np.ma.masked_where(segmap == 255, segmap)

object_idxs = np.unique(segmap)[:-1]
objects = np.empty(object_idxs.size, dtype=[('idx', 'i4'), ('pixels', 'O')])

# important bit here:

# this I can vectorize
objects['idx'] = object_idxs
# but this I cannot. and it takes forever.
for i in range(object_idxs.size):
    objects[i]['pixels'] = np.where(segmap == i)

# just plotting here
fig, ax = plt.subplots(constrained_layout=True)
image = ax.imshow(
    segmap, cmap='tab20', norm=mpl.colors.Normalize(vmin=0, vmax=20)
)
fig.colorbar(image)
fig.show()

Answer 1

在循環中使用np.where在算法上效率不高，因為時間復雜度為O(snm)其中s = object_idxs.size和n, m = segmap.shape 。 此操作可以在O(nm)中完成。

使用 Numpy 的一種解決方案是首先選擇所有對象像素位置，然后根據它們在segmap中的關聯對象對它們進行排序，最后根據對象數量對其進行拆分。 這是代碼：

background = np.max(segmap)
mask = segmap != background
objects = segmap[mask]
uniqueObjects, counts = np.unique(objects, return_counts=True)
ordering = np.argsort(objects)
i, j = np.where(mask)
indices = np.vstack([i[ordering], j[ordering]])
indicesPerObject = np.split(indices, counts.cumsum()[:-1], axis=1)

objects = np.empty(uniqueObjects.size, dtype=[('idx', 'i4'), ('pixels', 'O')])
objects['idx'] = uniqueObjects
for i in range(uniqueObjects.size):
    # Use `tuple(...)` to get the exact same type as the initial code here
    objects[i]['pixels'] = tuple(indicesPerObject[i])
# In case the conversion to tuple is not required, the loop can also be accelerated:
# objects['pixels'] = indicesPerObject

Answer 2

聽起來您想查看任何對象的位置。 因此，如果我們從一個矩陣開始（也就是說，所有形狀都在一個數組中，其中空格為零，對象 1 由 1 組成，對象 2 由 2 組成，等等），那么您可以創建一個遮罩，顯示哪些像素（或值在矩陣中）是非零的，如下所示：

my_array != 0