Skip to main content
 首页 » 编程设计

python之对二维 numpy 数组中的每个 NXN 子数组执行计算的最快方法

2025年05月04日22del

我有一个代表灰度图像的 2D numpy 数组。我需要提取该数组中的每个 N x N 子数组,子数组之间有指定的重叠,并计算一个属性,例如均值、标准差或中位数。

下面的代码执行此任务但速度很慢,因为它使用 Python for 循环。关于如何矢量化此计算或以其他方式加速计算的任何想法?

import numpy as np 
 
img = np.random.randn(100, 100) 
N = 4 
step = 2 
 
h, w = img.shape 
out = [] 
for i in range(0, h - N, step): 
    outr = [] 
    for j in range(0, w - N, step): 
        outr.append(np.mean(img[i:i+N, j:j+N])) 
    out.append(outr) 
out = np.array(out) 

请您参考如下方法:

对于均值和标准差,有一个基于cumsum 的快速解决方案。

这是 500x200 图像、30x20 窗口和步长 5 和 3 的时间。为了比较,我使用 skimage.util.view_as_windows 和 numpy 均值和标准差。

mn + sd using cumsum     1.1531693299184553 ms 
mn using view_as_windows 3.495307120028883 ms 
sd using view_as_windows 21.855629019846674 ms 

代码:

import numpy as np 
from math import gcd 
from timeit import timeit 
 
def wsum2d(A, winsz, stepsz, canoverwriteA=False): 
    M, N = A.shape 
    m, n = winsz 
    i, j = stepsz 
    for X, x, s in ((M, m, i), (N, n, j)): 
        g = gcd(x, s) 
        if g > 1: 
            X //= g 
            x //= g 
            s //= g 
            A = A[:X*g].reshape(X, g, -1).sum(axis=1) 
        elif not canoverwriteA: 
            A = A.copy() 
        canoverwriteA = True 
        A[x:] -= A[:-x] 
        A = A.cumsum(axis=0)[x-1::s] 
        A = A.T 
    return A 
 
def w2dmnsd(A, winsz, stepsz): 
    # combine A and A*A into a complex, so overheads apply only once 
    M21 = wsum2d(A*(A+1j), winsz, stepsz, True) 
    M2, mean_ = M21.real / np.prod(winsz), M21.imag / np.prod(winsz) 
    sd = np.sqrt(M2 - mean_*mean_) 
    return mean_, sd 
 
# test 
np.random.seed(0) 
A = np.random.random((500, 200)) 
wsz = (30, 20) 
stpsz = (5, 3) 
mn, sd = w2dmnsd(A, wsz, stpsz) 
from skimage.util import view_as_windows 
Av = view_as_windows(A, wsz, stpsz) # this emits a warning on my system 
assert np.allclose(mn, np.mean(Av, axis=(2, 3))) 
assert np.allclose(sd, np.std(Av, axis=(2, 3))) 
from timeit import repeat 
 
print('mn + sd using cumsum    ', min(repeat(lambda: w2dmnsd(A, wsz, stpsz), number=100))*10, 'ms') 
print('mn using view_as_windows', min(repeat(lambda: np.mean(Av, axis=(2, 3)), number=100))*10, 'ms') 
print('sd using view_as_windows', min(repeat(lambda: np.std(Av, axis=(2, 3)), number=100))*10, 'ms')