使用numpy优化水平和垂直邻接的计算

Question

我有以下单元格：

cells = np.array([[1, 1, 1],
                  [1, 1, 0],
                  [1, 0, 0],
                  [1, 0, 1],
                  [1, 0, 0],
                  [1, 1, 1]])

我想计算水平和垂直邻接来得出这个结果：

# horizontal adjacency 
array([[3, 2, 1],
       [2, 1, 0],
       [1, 0, 0],
       [1, 0, 1],
       [1, 0, 0],
       [3, 2, 1]])

# vertical adjacency 
array([[6, 2, 1],
       [5, 1, 0],
       [4, 0, 0],
       [3, 0, 1],
       [2, 0, 0],
       [1, 1, 1]])

实际的解决方案是这样的：

def get_horizontal_adjacency(cells):
    adjacency_horizontal = np.zeros(cells.shape, dtype=int)
    for y in range(cells.shape[0]):
        span = 0
        for x in reversed(range(cells.shape[1])):
            if cells[y, x] > 0:
                span += 1
            else:
                span = 0
            adjacency_horizontal[y, x] = span
    return adjacency_horizontal

def get_vertical_adjacency(cells):
    adjacency_vertical = np.zeros(cells.shape, dtype=int)
    for x in range(cells.shape[1]):
        span = 0
        for y in reversed(range(cells.shape[0])):
            if cells[y, x] > 0:
                span += 1
            else:
                span = 0
            adjacency_vertical[y, x] = span
    return adjacency_vertical

算法基本上是（对于水平邻接）：

1. 循环通过行
2. 通过列向后循环
3. 如果单元格的 x、y 值不为零，则在实际跨度上加 1
4. 如果单元格的 x、y 值为为零，则将实际跨度重置为零
5. 将跨度设置为结果数组的新 x、y 值

由于我需要在所有数组元素上循环两次，这对于较大的数组（例如图像）来说很慢。

有没有办法使用矢量化或其他一些 numpy 魔法来改进算法？

总结：

joni 和 Mark Setchell 提出了很好的建议！

我创建了一个带有示例图像的small Repo 和一个带有比较的python 文件。结果令人惊讶：

• 原始方法：3.675 秒
• 使用 Numba：0.002 秒
• 使用 Cython：0.005 秒

Answer 1

我对 Numba 进行了非常快速的尝试，但并没有彻底检查它，尽管结果似乎是正确的：

#!/usr/bin/env python3

# https://stackoverflow.com/q/69854335/2836621
# magick -size 1920x1080 xc:black -fill white -draw "circle 960,540 960,1040" -fill black -draw "circle 960,540 960,800" a.png

import cv2
import numpy as np
import numba as nb

def get_horizontal_adjacency(cells):
    adjacency_horizontal = np.zeros(cells.shape, dtype=int)
    for y in range(cells.shape[0]):
        span = 0
        for x in reversed(range(cells.shape[1])):
            if cells[y, x] > 0:
                span += 1
            else:
                span = 0
            adjacency_horizontal[y, x] = span
    return adjacency_horizontal

@nb.jit('void(uint8[:,::1], int32[:,::1])',parallel=True)
def nb_get_horizontal_adjacency(cells, result):
    for y in nb.prange(cells.shape[0]):
        span = 0
        for x in range(cells.shape[1]-1,-1,-1):
            if cells[y, x] > 0:
                span += 1
            else:
                span = 0
            result[y, x] = span
    return 

# Load image
im = cv2.imread('a.png', cv2.IMREAD_GRAYSCALE)

%timeit get_horizontal_adjacency(im)

result = np.zeros((im.shape[0],im.shape[1]),dtype=np.int32)
%timeit nb_get_horizontal_adjacency(im, result)

时间安排很好，如果运行正常，速度会提高 4000 倍：

In [15]: %timeit nb_get_horizontal_adjacency(im, result)
695 µs ± 9.12 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [17]: %timeit get_horizontal_adjacency(im)
2.78 s ± 44.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

输入

输入图像以 1080p 尺寸创建，即 1920x1080，ImageMagick 使用：

magick -size 1920x1080 xc:black -fill white -draw "circle 960,540 960,1040" -fill black -draw "circle 960,540 960,800" a.png

输出（对比度调整）

Answer 2

正如 cmets 中所述，这是一个完美的示例，通过 Cython 或 Numba 重写函数更容易。由于 Mark 已经提供了 Numba 解决方案，让我提供一个 Cython 解决方案。首先，让我们在我的机器上计时他的解决方案，以便进行公平的比较：

In [5]: %timeit nb_get_horizontal_adjacency(im, result)
836 µs ± 36 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

假设图像 im 是 np.ndarray 和 dtype=np.uint8，并行化 Cython 解决方案如下所示：

In [6]: %%cython -f -a -c=-O3 -c=-march=native -c=-fopenmp --link-args=-fopenmp

from cython import boundscheck, wraparound, initializedcheck
from libc.stdint cimport uint8_t, uint32_t
from cython.parallel cimport prange
import numpy as np

@boundscheck(False)
@wraparound(False)
@initializedcheck(False)
def cy_get_horizontal_adjacency(uint8_t[:, ::1] cells):
    cdef int nrows = cells.shape[0]
    cdef int ncols = cells.shape[1]
    cdef uint32_t[:, ::1] adjacency_horizontal = np.zeros((nrows, ncols), dtype=np.uint32)
    cdef int x, y, span
    for y in prange(nrows, nogil=True, schedule="static"):
        span = 0
        for x in reversed(range(ncols)):
            if cells[y, x] > 0:
                span += 1
            else:
                span = 0
            adjacency_horizontal[y, x] = span
    return np.array(adjacency_horizontal, copy=False)

在我的机器上，这几乎快了两倍：

In [7]: %timeit cy_get_horizontal_adjacency(im)
431 µs ± 4.38 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)