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

Question

我有一个代表灰度图像的 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)

Answer 1

对于均值和标准差，有一个基于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')