如何用Numpy在2D数组中匹配指定kernel模式 替代低效for循环

Numpy本身没有专门的内置模板匹配函数，但可以通过Numpy的滑动窗口视图实现完全向量化的匹配操作，无需写for循环，大数组下性能优异。你要实现的是二值kernel的精确匹配，本质是计算滑动窗口和kernel的互相关，结果等于kernel中1的总和的位置即为匹配位置。

方法1：纯Numpy实现（依赖Numpy 1.20+） #

通过sliding_window_view生成所有3x3滑动窗口，批量计算匹配结果：

import numpy as np

# 初始化输入数组和kernel
a = np.array([
    [0,0,1,1,1,0],
    [0,0,0,0,0,0],
    [1,0,0,0,0,1],
    [0,0,0,0,1,0],
    [0,0,0,1,0,0],
    [0,0,1,0,0,1]
], dtype=np.uint8)

k = np.array([
    [0,0,1],
    [0,1,0],
    [1,0,0]
], dtype=np.uint8)
kernel_sum = k.sum()

# 生成所有3x3滑动窗口
windows = np.lib.stride_tricks.sliding_window_view(a, k.shape)
# 计算每个窗口与kernel的匹配度，匹配时输出1
match_mask = (windows * k).sum(axis=(-2, -1)) == kernel_sum

# 填充为原数组大小，边缘补0
result = np.zeros_like(a)
result[1:-1, 1:-1] = match_mask.astype(np.uint8)

输出result即可得到你需要的结果。

方法2：基于Scipy的更高性能实现 #

如果允许使用Scipy，用signal.correlate2d实现更简洁，大数组下效率更高：

from scipy import signal
import numpy as np

a = np.array([
    [0,0,1,1,1,0],
    [0,0,0,0,0,0],
    [1,0,0,0,0,1],
    [0,0,0,0,1,0],
    [0,0,0,1,0,0],
    [0,0,1,0,0,1]
], dtype=np.uint8)

k = np.array([
    [0,0,1],
    [0,1,0],
    [1,0,0]
], dtype=np.uint8)
kernel_sum = k.sum()

# 互相关计算，输出与原数组尺寸一致
corr_result = signal.correlate2d(a, k, mode='same', boundary='fill', fillvalue=0)
result = (corr_result == kernel_sum).astype(np.uint8)

内容的提问来源于stack exchange，提问作者Guy Barash