使用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提出了很好的建议！

我创建了一个小型Repo，其中包含一个示例图像和一个带有比较的python文件。结果令人惊讶：

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

- Lukas Weber

---

你是否正在尝试使用OpenCV的distanceTransform？ - Quang Hoang

我想用Python实现https://www.evryway.com/largest-interior/。 - Lukas Weber

在我看来，加速函数最简单的方法就是将它们都重写为Cython。这几乎是相同的代码。 - joni

Answer 1

我用Numba试图快速尝试了一下，但没有仔细检查过，虽然结果看起来差不多正确：

#!/usr/bin/env python3

# https://dev59.com/dcLra4cB1Zd3GeqPOacA
# 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

输出（对比度调整后）

- Mark Setchell

---

for x in range(cells.shape[1]-1,0,-1): needs to be range(cells.shape[1]-1,-1,-1). if not, the last x will be 1 instead of 0 - Lukas Weber

@LukasWeber 感谢您的更正 - 我已经更新了答案。有趣的是，那正是我没有仔细检查的部分！ - Mark Setchell

你能解释一下为什么nb.prange()只在第一个for循环中使用吗？ - Lukas Weber

Numba可以并行化（即分布式地）处理使用prange编写的任何循环。 如果您有4个CPU核心，则它将在每个核心上执行1/4的循环。 如果您在内部循环中再次执行prange，则实际上并没有16个核心可用，因此没必要进一步分割循环。 我可能是错误的，也没有计时，但通常我会使最外层的循环并行化，并且认为在尝试进一步分割内部循环时已经没有太大帮助了。 如果我错了，请有人告诉我！ - Mark Setchell

我应该在多波段图像中使用uint8[:, ::1, :]还是uint8[:, :, ::1]？ - Lukas Weber

Answer 2

如评论中所述，这是一个完美的例子，可以通过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.uint8的np.ndarray，并行化的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)

- joni

---

很酷 - 干得好！ - Mark Setchell

你好，joni！这是我第一次尝试使用Cython。通过我的编译脚本，你的方法比Numba慢了一点（请参见编辑后的问题）。你能否看一下我是否正确编译了所有内容？https://github.com/lukasalexanderweber/SO-Improve-Performance-of-Numpy-Iterations/tree/main/cy - Lukas Weber

你应该设置正确的编译器标志。在Windows上，应该在*.pyx文件的顶部设置# distutils: extra_compile_args=/Ox /arch:AVX2 /openmp。然而，根据我的经验，与gcc/clang相比，MSVC编译器在SIMD自动向量化方面表现不佳（这就是gcc/clang编译器的“-march=native”标志的目的）。因此，如果Numba解决方案在Windows上更快，我也不会感到惊讶。 - joni