在NumPy数组上实现类Matlab nlfilter的高效低内存二维滑动窗口迭代

需求说明 #

我有一个my_fun函数，可处理图像的k×k窗口（k为奇数）并返回标量。需要遍历图像每个像素，提取其周围的k×k窗口，寻找速度快、内存开销极小的迭代器，功能类似Matlab的nlfilter。

现有实现示例 #

当前通过嵌套循环提取窗口的代码如下：

import numpy as np

my_array = np.arange(25).reshape((5,5))
win_radius = (1, 1) # k = 3

for i in range(my_array.shape[0]):
    for j in range(my_array.shape[1]):
        win_rows = max(0, i - win_radius[0]), min(my_array.shape[0], i + win_radius[0] + 1)
        win_cols = max(0, j - win_radius[1]), min(my_array.shape[1], j + win_radius[1] + 1)
        my_win = my_array[win_rows[0]:win_rows[1], win_cols[0]:win_cols[1]]
        print(my_win)

输出片段：

[[0 1]
 [5 6]]
[[0 1 2]
 [5 6 7]]
[[1 2 3]
 [6 7 8]]
[[2 3 4]
 [7 8 9]]
[[3 4]
 [8 9]]
[[ 0  1]
 [ 5  6]
 [10 11]]
[[ 0  1  2]
 [ 5  6  7]
 [10 11 12]]
[[ 1  2  3]
 [ 6  7  8]
 [11 12 13]]
[[ 2  3  4]
 [ 7  8  9]
 [12 13 14]]
[[ 3  4]
 [ 8  9]
 [13 14]]
...
[[17 18 19]
 [22 23 24]]
[[18 19]
 [23 24]]

高效迭代方案 #

方案1：滑动窗口视图（内存零复制） #

利用numpy的滑动窗口视图生成完整k×k窗口，通过边界填充统一窗口尺寸，内存开销极低（仅返回原数组视图）。

import numpy as np

def full_window_iterator(arr, k):
    radius = k // 2
    # 边界填充，可替换mode为'reflect'/'edge'等适配业务需求
    padded_arr = np.pad(arr, pad_width=radius, mode='constant', constant_values=0)
    # 生成滑动窗口视图，无内存复制
    windows = np.lib.stride_tricks.sliding_window_view(padded_arr, window_shape=(k, k))
    # 逐个返回对应原像素的窗口
    for i in range(arr.shape[0]):
        for j in range(arr.shape[1]):
            yield windows[i, j]

# 测试
my_array = np.arange(25).reshape((5,5))
for win in full_window_iterator(my_array, 3):
    print(win)

方案2：轻量生成器（保留边缘窗口原始尺寸） #

如果需要保留边缘窗口的不规则大小，用生成器按需生成窗口，内存开销几乎为0，逻辑与原代码一致但更简洁。

import numpy as np

def nlfilter_style_iterator(arr, k):
    radius = k // 2
    rows, cols = arr.shape
    for i in range(rows):
        row_start = max(0, i - radius)
        row_end = min(rows, i + radius + 1)
        for j in range(cols):
            col_start = max(0, j - radius)
            col_end = min(cols, j + radius + 1)
            yield arr[row_start:row_end, col_start:col_end]

# 测试
my_array = np.arange(25).reshape((5,5))
for win in nlfilter_style_iterator(my_array, 3):
    print(win)

性能与场景适配 #

• 滑动窗口视图：适合需要统一k×k窗口的场景，numpy底层优化，速度最快。
• 生成器方案：适合保留边缘窗口原始大小的场景，内存占用极低，速度接近原生循环。
• 若my_fun支持向量化，直接对滑动窗口视图批量调用，性能会远高于逐个迭代。

内容的提问来源于stack exchange，提问作者Eric Johnson