Numpy 有效地找到标记图像区域的质心

我有一个分割图像作为一个二维矩阵的唯一标签1。。。K例如：

img = 
    [1 1 2 2 2 2 2 3 3]
    [1 1 1 2 2 2 2 3 3]
    [1 1 2 2 2 2 3 3 3]
    [1 4 4 4 2 2 2 2 3]
    [4 4 4 5 5 5 2 3 3]
    [4 4 4 5 5 6 6 6 6]
    [4 4 5 5 5 5 6 6 6]

我试图确定区域质心。也就是说，每个标签，质量中心的X，Y坐标是多少？例如，标签1的质心为（1.25，0.625）。只需将行数（

（0+0+1+1+1+2+2+3）/8=1.25）和列数（
（0+0+0+0+1+1+1+2）/8=0.625）求平均值即可

我知道如何做到这一点的唯一方法是使用从1到k（或者在我的示例中，从1到6）的for循环，找到每个标签的点的索引，并通过索引图像的网格来平均它们的坐标

然而，我希望这样做的方式优化的GPU计算。因此，使用for循环并不理想（对于几百个标签，在一个漂亮的GPU上每个图像大约需要1秒）。我正在使用PyTorch，但实际上任何numpy解决方案都应该足够了

这个任务有没有一个GPU高效的解决方案？
这个计算需要累加，我不知道在GPU上有多高效。这是psuedo代码中的顺序算法：

int n[k] = 0
int sx[k] = 0
int sy[k] = 0
loop over y:
   loop over x:
      i = img[x,y]
      ++n[i]
      sx[i] += x
      sy[i] += y
for i = 1 to k
    sx[i] /= n[i]
    sy[i] /= n[i]

当然，
（sx[i]，sy[i]）
是物体的质心

它在CPU上的速度非常快，因此不值得为此而将数据发送到GPU，除非它已经存在。
考虑使用或重用它们（基于numpy/scipy）

下面是一个演示：

import numpy as np
from skimage import measure
from time import perf_counter as pc

img = np.array([[1, 1, 2, 2, 2, 2, 2, 3, 3],
                [1, 1, 1, 2, 2, 2, 2, 3, 3],
                [1, 1, 2, 2, 2, 2, 3, 3, 3],
                [1, 4, 4, 4, 2, 2, 2, 2, 3],
                [4, 4, 4, 5, 5, 5, 2, 3, 3],
                [4, 4, 4, 5, 5, 6, 6, 6, 6],
                [4, 4, 5, 5, 5, 5, 6, 6, 6]])

# assuming already labels of 1, 2, ... n
times = [pc()]
props = measure.regionprops(img)
times.append(pc())
for i in range(np.unique(img).shape[0]):
    print(props[i].centroid)
    times.append(pc())

print(np.diff(times))

输出：

(1.25, 0.625)
(1.5555555555555556, 4.4444444444444446)
(1.8999999999999999, 7.4000000000000004)
(4.3636363636363633, 1.1818181818181819)
(5.1111111111111107, 3.6666666666666665)
(5.4285714285714288, 6.7142857142857144)
[  9.05569615e-05   8.51235438e-04   2.48126075e-04   2.59294767e-04
   2.42692657e-04   2.00734598e-04   2.34542530e-04]

一种想法是使用输入数组中的数字作为存储单元来累积每个区域的行和列索引，从而得到矢量化的解决方案，如下所示-

m,n = a.shape
r,c = np.mgrid[:m,:n]
count = np.bincount(a.ravel())
centroid_row = np.bincount(a.ravel(),r.ravel())/count
centroid_col = np.bincount(a.ravel(),c.ravel())/count

样本运行-

In [77]: a
Out[77]: 
array([[1, 1, 2, 2, 2, 2, 2, 3, 3],
       [1, 1, 1, 2, 2, 2, 2, 3, 3],
       [1, 1, 2, 2, 2, 2, 3, 3, 3],
       [1, 4, 4, 4, 2, 2, 2, 2, 3],
       [4, 4, 4, 5, 5, 5, 2, 3, 3],
       [4, 4, 4, 5, 5, 6, 6, 6, 6],
       [4, 4, 5, 5, 5, 5, 6, 6, 6]])

In [78]: np.c_[centroid_row, centroid_col]
Out[78]: 
array([[  nan,   nan], 
       [ 1.25,  0.62], # centroid for region-1
       [ 1.56,  4.44], # centroid for region-2
       [ 1.9 ,  7.4 ], # centroid for region-3 and so on.
       [ 4.36,  1.18],
       [ 5.11,  3.67],
       [ 5.43,  6.71]])

是的，这也是我能想到的最好的解决办法。数据已经在GPU上了，因此我想在那里做它的原因：-/。我来看看CPU是怎么运行的。这当然和它得到的一样有效——每个像素接触一次，但我想知道是否有一种方法可以在GPU上并行化……如果上面算法的简单GPU实现太慢，那么可以做的一件事是为每个图像列创建数组
n
，
sx
和
sy
，然后把它们加在一起。这可能会减少等待对数组值进行原子更新的内核数量。我不确定，以原子方式更新全局数组中的单个值是否有效？你会在GPU上遇到错误的共享问题吗？