Vectorized 2-D moving window in numpy including edges

杀马特。学长 韩版系。学妹 提交于 2019-11-30 21:26:52

You could define a function that yields a generator and use that. The window would be the floor of the shape you want divided by 2 and the trick would be just indexing the array along that window as you move along the rows and columns.

def window(arr, shape=(3, 3)):
    # Find row and column window sizes
    r_win = np.floor(shape[0] / 2).astype(int)
    c_win = np.floor(shape[1] / 2).astype(int)
    x, y = arr.shape
     for i in range(x):
         xmin = max(0, i - r_win)
         xmax = min(x, i + r_win + 1)
         for j in range(y):
             ymin = max(0, j - c_win)
             ymax = min(y, j + c_win + 1)
             yield arr[xmin:xmax, ymin:ymax]

You could use this function like so:

arr = np.array([[1,2,3,4],
               [2,3,4,5],
               [3,4,5,6],
               [4,5,6,7]])
gen = window(arr)
next(gen)
array([[1, 2],
       [2, 3]])

Going through the generator produces all of the windows in your example.

It's not vectorized, but I'm not sure there is an existing vectorized function that returns different sized arrays. As @PaulPanzer points out you could pad your array to the size you need and use a np.lib.stride_tricks.as_strided to generate a view of the slices. Something like so:

def rolling_window(a, shape):
    s = (a.shape[0] - shape[0] + 1,) + (a.shape[1] - shape[1] + 1,) + shape
    strides = a.strides + a.strides
    return np.lib.stride_tricks.as_strided(a, shape=s, strides=strides)

def window2(arr, shape=(3, 3)):
    r_extra = np.floor(shape[0] / 2).astype(int)
    c_extra = np.floor(shape[1] / 2).astype(int)
    out = np.empty((arr.shape[0] + 2 * r_extra, arr.shape[1] + 2 * c_extra))
    out[:] = np.nan
    out[r_extra:-r_extra, c_extra:-c_extra] = arr
    view = rolling_window(out, shape)
    return view

window2(arr, (3,3))
array([[[[ nan,  nan,  nan],
         [ nan,   1.,   2.],
         [ nan,   2.,   3.]],

        [[ nan,  nan,  nan],
         [  1.,   2.,   3.],
         [  2.,   3.,   4.]],

        [[ nan,  nan,  nan],
         [  2.,   3.,   4.],
         [  3.,   4.,   5.]],

        [[ nan,  nan,  nan],
         [  3.,   4.,  nan],
         [  4.,   5.,  nan]]],


       [[[ nan,   1.,   2.],
         [ nan,   2.,   3.],
         [ nan,   3.,   4.]],

        [[  1.,   2.,   3.],
         [  2.,   3.,   4.],
         [  3.,   4.,   5.]],

        [[  2.,   3.,   4.],
         [  3.,   4.,   5.],
         [  4.,   5.,   6.]],

        [[  3.,   4.,  nan],
         [  4.,   5.,  nan],
         [  5.,   6.,  nan]]],


       [[[ nan,   2.,   3.],
         [ nan,   3.,   4.],
         [ nan,   4.,   5.]],

        [[  2.,   3.,   4.],
         [  3.,   4.,   5.],
         [  4.,   5.,   6.]],

        [[  3.,   4.,   5.],
         [  4.,   5.,   6.],
         [  5.,   6.,   7.]],

        [[  4.,   5.,  nan],
         [  5.,   6.,  nan],
         [  6.,   7.,  nan]]],


       [[[ nan,   3.,   4.],
         [ nan,   4.,   5.],
         [ nan,  nan,  nan]],

        [[  3.,   4.,   5.],
         [  4.,   5.,   6.],
         [ nan,  nan,  nan]],

        [[  4.,   5.,   6.],
         [  5.,   6.,   7.],
         [ nan,  nan,  nan]],

        [[  5.,   6.,  nan],
         [  6.,   7.,  nan],
         [ nan,  nan,  nan]]]])

This version pads the edges with np.nan to avoid confusion with any other values in your array. It is about 3x faster with the given array than the window function, but I am not sure how having padded output will impact anything you want to do downstream.

This is not strictly an answer to your question since its not vectorized, but hopefully its a helpful benchmark for any other potential solutions (there's bound to be something in image processing libraries?)

Regardless, I've implemented the window as a loop which takes the average of the window with output into a new array. Inputs are an array and the window size +/- the current index. One version using straight Python and Numpy, the other compiled with numba.

def mw_mean(in_arr,out_arr,x_win,y_win):
    xn,yn = in_arr.shape
    for x in range(xn):
        xmin = max([0,x - x_win])
        xmax = min([xn, x + x_win + 1])
        for y in range(yn):
            ymin = max([0,y - y_win])
            ymax = min([yn, y + y_win + 1])

            out_arr[x,y] = in_arr[xmin:xmax, ymin:ymax].mean()
    return out_arr



@jit("i4[:,:](i4[:,:],i4[:,:],i4,i4)", nopython = True)
def mw_mean_numba(in_arr,out_arr,x_win,y_win):
    xn,yn = in_arr.shape
    for x in range(xn):
        xmin = max(0,x - x_win)
        xmax = min(xn, x + x_win + 1)
        for y in range(yn):
            ymin = max(0,y - y_win)
            ymax = min(yn, y + y_win + 1)

            out_arr[x,y] = in_arr[xmin:xmax, ymin:ymax].mean()
    return out_arr

This is tested against three different array sizes--your original test case and two larger ones (100x100 and 1000x1000):

a = np.array([[1,2,3,4], [2,3,4,5], [3,4,5,6], [4,5,6,7]])
b = np.random.randint(1,7, size = (100,100))
c = np.random.randint(1,7, size = (1000,1000))

aout,bout,cout = np.zeros_like(a),np.zeros_like(b),np.zeros_like(c)

x_win = 1
y_win = 1

Runtimes without compilation:

%timeit mw_mean(a,aout,x_win,y_win)
1000 loops, best of 3: 225 µs per loop

%timeit mw_mean(b,bout,x_win,y_win)
10 loops, best of 3: 137 ms per loop

%timeit mw_mean(c,cout,x_win,y_win)
1 loop, best of 3: 14.1 s per loop

Runtimes with compilation:

%timeit mw_mean_numba(a,aout,x_win,y_win)
1000000 loops, best of 3: 1.22 µs per loop

%timeit mw_mean_numba(b,bout,x_win,y_win)
1000 loops, best of 3: 550 µs per loop

%timeit mw_mean_numba(c,cout,x_win,y_win)
10 loops, best of 3: 55.1 ms per loop

Edit: a previous version of this was modifying the array in place, which is obviously a big no-no for a rolling window. Benchmarks remain unchanged.

标签
易学教程内所有资源均来自网络或用户发布的内容,如有违反法律规定的内容欢迎反馈
该文章没有解决你所遇到的问题?点击提问,说说你的问题,让更多的人一起探讨吧!