Extracting patches from arrays

sliding_window.py

"""
Code for extracting sliding windows from arrays. 
Particularly useful when you want to take patches from images prior to performing
some operation on them, like convolution or computing the mean. 
"""
import numpy as np
from itertools import zip_longest
from PIL import Image
from numpy.lib.stride_tricks import as_strided


def sliding_window_lastaxis(a, ws, ss):
    """Create a sliding window along the last axis of the array.

    Parameters
    ----------
    a: numpy.ndarray
        The array from which to extract the window
    ws: int
        The size of the window along the dimension of the last axis
    ss: int
        The stride size; how much to move in between windows
        For examples, if `ss = 1`, then the windows are offset by a single
        pixel, whereas if `ss = ws`, the windows are non-overlapping, with a
        new window "starting" immediately after the previous one.
    """
    if ws < 1:
        raise ValueError("`ws` must be at least 1.")
    if ws > a.shape[-1]:
       raise ValueError("`ws` must be shorter than the last axis.")

    ns = 1 + (a.shape[-1]-ws)//ss
    shape = a.shape[:-1] + (ns, ws)
    strides = a.strides[:-1] + (a.strides[-1]*ss, a.strides[-1])
    return as_strided(a, shape=shape, strides=strides)


def sliding_window(a, window, offset, dense=False):
    """
    Create windows/patches into an array.

    Parameters
    ----------
    a : numpy.ndarray
        The array from which to extract windows.
    window: int or iterable
        The window size.
    offset: int or iterable
        The distance to move between windows.
    dense: bool
        If `offset` has fewer elements than `window` and `dense` is set to True,
        the patches returned will be offset by one in between each other when
        `offset` is not specified.
        Otherwise, if `dense` is False (the default), they will be offset by the
        size of the window in that dimension (non-overlapping).
    """
    if not hasattr(window, '__iter__'):
        window = [window]
    if not hasattr(offset, '__iter__'):
        offset = [offset]

    for i, (win, off) in enumerate(zip_longest(window, offset)):
        if off is None:
            if dense:
                off = 1
            else:
                off = win
        a = a.swapaxes(i, -1)
        a = sliding_window_lastaxis(a, win, off)
        a = a.swapaxes(-2, i)
    return a

def glue_patches2d(a, shape=None):
    """Glue patches together, assuming array shape is (x, y, ..., wx, wy).
    Mainly useful for visualizing what the sliding windows look like.
    """
    if shape is None:
        shape = a.shape[:-2]
    dx, dy = a.shape[-2:]
    ret = np.empty((dx*shape[0], dy*shape[1]), dtype=a.dtype)
    for i in np.ndindex(*shape):
        x = i[0]*dx
        y = i[1]*dy
        ret[x:x+dx, y:y+dy] = a[i]

    return ret