我可以对这个 python 代码进行矢量化吗？

Question

我编写了这个 python 代码来获取标签的邻居（一组共享一些共同属性的像素）。标签的邻居被定义为位于边界另一侧的其他标签（相邻标签共享一个边界）。所以，我写的代码可以工作，但速度极慢：

# segments: It is a 2-dimensional numpy array (an image really)
# where segments[x, y] = label_index. So each entry defines the
# label associated with a pixel.

# i: The label whose neighbours we want.

def get_boundaries(segments, i):
    neighbors = []
    for y in range(1, segments.shape[1]):
        for x in range(1, segments.shape[0]):
            # Check if current index has the label we want 
            if segments[x-1, y] == i:
                # Check if neighbour in the x direction has
                # a different label
                if segments[x-1, y] != segments[x, y]:
                    neighbors.append(segments[x,y])

            # Check if neighbour in the y direction has
            # a different label
            if segments[x, y-1] == i:
                if segments[x, y-1] != segments[x, y]:
                    neighbors.append(segments[x, y])

    return np.unique(np.asarray(neighbors))

你可以想象，我在这里可能完全误用了 python。我想知道是否有办法优化此代码以使其更具 Python 风格。

Answer 1

给你：

def get_boundaries2(segments, i):
    x, y = np.where(segments == i) # where i is
    right = x + 1
    rightMask = right < segments.shape[0] # keep in bounds
    down = y + 1
    downMask = down < segments.shape[1]
    rightNeighbors = segments[right[rightMask], y[rightMask]]
    downNeighbors = segments[x[downMask], down[downMask]]
    neighbors = np.union1d(rightNeighbors, downNeighbors)
    return neighbors

如您所见，根本没有 Python 循环；我还尝试最小化副本（第一次尝试制作了带有 NAN 边框的 segments 的副本，但后来我设计了“保持边界”检查）。

请注意，我没有从这里的“邻居”中过滤掉i 本身；如果需要，您可以在最后轻松添加。一些时间安排：

输入 2000x3000：原始需要 13 秒，我需要 370 毫秒（35 倍加速）。

输入 1000x300：原始需要 643 毫秒，我需要 17.5 毫秒（36 倍加速）。

Answer 2

你需要用 numpy 的隐式循环替换你的 for 循环。

我对你的代码了解不多，无法直接转换，但我可以举个例子。

假设你有一个由 100000 个随机整数组成的数组，你需要得到一个每个元素除以其邻居的数组。

import random, numpy as np 
a = np.fromiter((random.randint(1, 100) for i in range(100000)), int)

一种方法是：

[a[i] / a[i+1] for i in range(len(a)-1)]

或者这个更快：

a / np.roll(a, -1)

时间：

initcode = 'import random, numpy as np; a = np.fromiter((random.randint(1, 100) for i in range(100000)), int)'
timeit.timeit('[a[i] / a[i+1] for i in range(len(a)-1)]', initcode, number=100)
5.822079309000401
timeit.timeit('(a / np.roll(a, -1))', initcode, number=100)
0.1392055350006558