python - 一种优雅/更快的方法来找到图像上线条的端点？

Question

我一直致力于通过将数组操作的 for 循环替换为适当的 NumPy 函数来提高代码速度。

该函数的目的是获取一条线的端点，这是 255 个像素点中唯一具有一个相邻像素的点。

有没有办法让我可以从 np.where 获得两个点，有条件或我不熟悉的一些 NumPy 函数可以完成这项工作？

def get_end_points(图像):

x1=-1
y1=-1
x2=-1
y2=-1
    for i in range(image.shape[0]):
        for j in range(image.shape[1]):
            if image[i][j]==255 and neighbours_sum(i,j,image) == 255:
                if x1==-1:
                    x1 = j
                    y1 = i
                else:
                    x2=j
                    y2=i
return x1,y1,x2,y2

Answer 1

这是一个使用卷积的解决方案:

import numpy as np
import scipy.signal

def find_endpoints(img):
    # Kernel to sum the neighbours
    kernel = [[1, 1, 1],
              [1, 0, 1],
              [1, 1, 1]]
    # 2D convolution (cast image to int32 to avoid overflow)
    img_conv = scipy.signal.convolve2d(img.astype(np.int32), kernel, mode='same')
    # Pick points where pixel is 255 and neighbours sum 255
    endpoints = np.stack(np.where((img == 255) & (img_conv == 255)), axis=1)
    return endpoints

# Test
img = np.zeros((1000, 1000), dtype=np.uint8)
# Draw a line from (200, 130) to (800, 370)
for i in range(200, 801):
    j = round(i * 0.4 + 50)
    img[i, j] = 255
print(find_endpoints(img))
# [[200 130]
#  [800 370]]

编辑:

您也可以考虑使用 Numba 来实现此目的。该代码几乎是您已经拥有的代码，因此可能不是特别“优雅”，但速度要快得多。例如，像这样:

import numpy as np
import numba as nb

@nb.njit
def find_endpoints_nb(img):
    endpoints = []
    # Iterate through every row and column
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            # Check current pixel is white
            if img[i, j] != 255:
                continue
            # Sum neighbours
            s = 0
            for ii in range(max(i - 1, 0), min(i + 2, img.shape[0])):
                for jj in range(max(j - 1, 0), min(j + 2, img.shape[1])):
                    s += img[ii, jj]
            # Sum including self pixel for simplicity, check for two white pixels
            if s == 255 * 2:
                endpoints.append((i, j))
                if len(endpoints) >= 2:
                    break
        if len(endpoints) >= 2:
            break
    return np.array(endpoints)

print(find_endpoints_nb(img))
# [[200 130]
#  [800 370]]

这在我的计算机上运行速度相对较快:

%timeit find_endpoints(img)
# 34.4 ms ± 64.4 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit find_endpoints_nb(img)
# 552 µs ± 4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

此外，它应该使用更少的内存。上面的代码假设只有两个端点。如果添加并行化，您可能可以使其速度更快(尽管您必须进行一些更改，因为您无法从并行线程修改列表端点)。