Skip to content

Instantly share code, notes, and snippets.

@hanayashiki
Created December 28, 2019 07:39
Show Gist options
  • Save hanayashiki/4c4361d676726f7128bb7058a4b625ba to your computer and use it in GitHub Desktop.
Save hanayashiki/4c4361d676726f7128bb7058a4b625ba to your computer and use it in GitHub Desktop.
2d sliding window implementation for numpy
import numpy as np
from numpy.lib.stride_tricks import as_strided
from typing import *
def sliding_window_2d(a: np.array, wnd_size: Tuple[int, int]):
# a: an nd-array of (height, width, ...rest)
# wnd_size: a tuple of (wnd_width, wnd_height)
# returns an nd-array of (height - wnd_height + 1, width - wnd_width + 1, wnd_height, wnd_width, ...rest)
height, width, *rest = a.shape
s0, s1, *rest_strides = a.strides
wnd_width, wnd_height = wnd_size
return as_strided(a,
shape=(height - wnd_height + 1, width - wnd_width + 1, wnd_height, wnd_width, *rest),
strides=(s0, s1, s0, s1, *rest_strides))
if __name__ == '__main__':
a = np.arange(0, 3 * 4).reshape((3, 4))
print(a)
print(sliding_window_2d(a, (2, 2)))
assert (np.sum(sliding_window_2d(a, (2, 2)) == np.array([[2, 3], [6, 7]])) == 4)
a = np.arange(0, 3 * 4 * 3).reshape((3, 4, 3))
print(a)
print(sliding_window_2d(a, (2, 2)))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment