numpy sliding 2d window calculations

Question

I'm trying to learn a way to use numpy to efficiently solve problems that involve a sliding window in a variety of circumstances. Here is an example illustrating the type of problem I'm interested in:

I have a large 2d matrix, and I'd like to perform calculations on neighbors of every element in the matrix. For example, I may want to find the max value, excluding some special value of negibors at indexes at (x-1,y)(x+1,y+1) at each index, and put the result into another different 2d "solution" matrix.

Note that convolution2d, while useful won't work for me in this situation because I have specific operations to do at each pixel, and only want to do it on specific neighbors (of every pixel).

Also a bonus would be ensuring I don't go out of bounds.

Finally, is it possible to use any state as well? In cases where all neighbors are 0 I am hoping to assign a new integer id that I increment each time this occurs.

Here is an example:

Window:

0 0 1
1 0 0
0 0 0


Input:

0 0 0 0 0 0 0 0 0
0 0 0 0 0 0 0 0 0
0 0 0 0 0 0 0 0 0
0 0 9 9 0 0 9 0 0
0 0 0 0 0 0 0 0 0

Output:

0 0 0 0 0 0 0 0 0
0 0 0 0 0 0 0 0 0
0 0 0 0 0 0 0 0 0
0 0 1 1 0 0 2 0 0
0 0 0 0 0 0 0 0 0

Answer 1

Use np.roll() to create secondary matrices. Then perform whatever operations you need between the initial and secondary matrices. For example, to take the average of the central cell and two neighbors:

sec_a = np.roll(mtrx, -1, axis=0)
sec_b = np.roll(mtrx, -1, axis=1)

result = (mtrx + sec_a + sec_b) / 3

Additionally, roll() rolls around the edges, so no need to worry about bounds.

Answer 2

I once created this function to store sliding blocks from a 2D array into columns, so that any operation that we once thought to apply in a sliding window on a 2D array could be easily applied along the columns. Read more about it in this solution to Implement Matlab's im2col 'sliding' in python.

Now, NumPy supports most of its functions to be applied along a specified axis. So, with this tool, effectively we would be able to apply almost any operation in a sliding window in a vectorized way. Here's the formal definition of it -

def im2col(A,BLKSZ):   

    # Parameters
    M,N = A.shape
    col_extent = N - BLKSZ[1] + 1
    row_extent = M - BLKSZ[0] + 1

    # Get Starting block indices
    start_idx = np.arange(BLKSZ[0])[:,None]*N + np.arange(BLKSZ[1])

    # Get offsetted indices across the height and width of input array
    offset_idx = np.arange(row_extent)[:,None]*N + np.arange(col_extent)

    # Get all actual indices & index into input array for final output
    return np.take (A,start_idx.ravel()[:,None] + offset_idx.ravel())

Here's how we can use this tool to solve the problem at hand, assuming A as the 2D input array -

# Get 3x3 sliding blocks from A and set them as columns.
Acol = im2col(A,[3,3])

# Setup kernel mask
kernel = np.ones((3,3),dtype=bool)
kernel[2,1:] = 0

# Mask rows of Acol with kernel and perform any operation, let's say MAX
out = Acol[kernel.ravel()].max(0).reshape(A.shape[0]-2,A.shape[1]-2)

Sample run -

In [365]: A
Out[365]: 
array([[83, 84, 46,  9, 25],
       [32,  8, 31, 45, 58],
       [14,  8,  0, 51, 27],
       [31, 40,  7, 27, 71]])

In [366]: kernel = np.ones((3,3),dtype=bool)
     ...: kernel[2,1:] = 0
     ...: 

In [367]: im2col(A,[3,3])[kernel.ravel()].max(0).reshape(A.shape[0]-2,A.shape[1]-2)
Out[367]: 
array([[84, 84, 58],
       [32, 51, 58]])

Answer 3

Assuming your original 2D matrix is named A and has a size (n, m)

# extraction of 3x3 sub-matrices and storage in a new 2D matrix
B = [ [ A[i-1:i+2, j-1:j+2] for i in range(1, n-1) ] for j in range(1, m-1) ]
# conversion to a mask array
B = np.ma.array( B, mask=False )
# masking the unwanted elements of each sub-matrix
B.mask[:, :, 1, 2] = True
B.mask[:, :, 2, 2] = True

NB: the ranges of i and j at the creation of the sub-matrices have been chosen to avoid the boundaries.

Operations on a sub-matrix B[i, j] will ignore the masked elements.

Now, to perform an numpy operation (eg the max of the sub-matrix) on each sub-matrix and store the result in a 2D matrix :

C = [ [ np.max(B[i,j]) for i in range(n-2) ] for j in range(m-2) ]

Answer 4

I've used the following as a readable solution:

import numpy as np

def get_windows(arr, window_size=64, step=32):
  windows = []
  row = 0
  col = 0
  max_row, max_col = arr.shape
  while row < max_row:
    while col < max_col:
      windows.append(arr[row:row+window_size, col:col+window_size])
      col += step
    row += step
    col = 0
  return windows

a = np.random.rand(4, 4)
windows = get_windows(a, window_size=2, step=1)