divide image into overlapping tiles

tile_image.py

from skimage.util import pad
from collections import OrderedDict


# FIXME: probably a little neater to return a 4d array of shape (n_tile_rows, n_tile_col, H, W)
def tile_image(im, tile_size, overlap=(0,0)):
    """
    Divide an image or video into spatially overlapping tiles.
    
    Args:
        im (array): a 2- 3- or 4-dimensional image/video array. The array
            will be zero padded as needed, then tiled along the first two
            dimensions.
        tile_size (int or 2-tuple of ints)
        overlap (int or 2-tuple of ints): consecutive tiles will overlap by
            this many pixels.
            
    Returns:
        dict: a mapping from (row_ind, col_ind) to image tiles. Each image tile
            is exactly `tile_size` in the first two dimensions. In general, the
            tiles in the last row/column will be zero padded.
    """
    # Validate/prep input args
    if isinstance(tile_size, int):
        tile_size = (tile_size, tile_size)
    if isinstance(overlap, int):
        overlap = (overlap, overlap)
    tile_size = np.asarray(tile_size, dtype=int)
    im_shape = np.asarray(im.shape[:2], dtype=int)
    overlap = np.asarray(overlap, dtype=int)
    
    # Compute some derived parameters
    stride = tile_size - overlap
    num_tiles = np.ceil((im_shape + overlap - 1).astype(float) / tile_size).astype(int)
    padding = num_tiles * tile_size - im_shape - overlap
    
    # Pad image to fit an even number of overlapping tiles
    if padding.max() > 0:
        if im.ndim == 2:
            padding = (0, padding[0]), (0, padding[1])
        elif im.ndim == 3:
            padding = (0, padding[0]), (0, padding[1]), (0,0)
        else:
            raise ValueError('input image should be 2 or 3 dimensions')
        padded_im = pad(im, padding, mode='constant')
        padded_im[:im_shape[0], :im_shape[1], ...] = im
    else:
        padded_im = im
    
    # Compute the tiles, put in dict
    trailing_dims = im.shape[2:]
    tile_arr_shape = tuple(np.hstack([num_tiles, tile_size, trailing_dims]))
    tile_arr = np.zeros(tile_arr_shape, dtype=np.float32)
    for tile_row in range(num_tiles[0]):
        row_sta = tile_row * stride[0]
        row_end = row_sta + tile_size[0]
        for tile_col in range(num_tiles[1]):
            col_sta = tile_col * stride[1]
            col_end = col_sta + tile_size[1]
            tile = padded_im[row_sta:row_end, col_sta:col_end]
            tile_arr[tile_row, tile_col] = tile
    return tile_arr


tile_size = (1024, 1024)
overlap = (256, 256)

tiles = tile_image(mosaic_im, tile_size, overlap)

plt.figure(figsize=(16,16))
for r in range(tiles.shape[0]):
    for c in range(tiles.shape[1]):
        num = r * tiles.shape[1] + c
        tile = normalize(tiles[r, c, :, :, 3:])
        plt.subplot(tiles.shape[0], tiles.shape[1], num+1)
        plt.imshow(tile)
        plt.title(np.isnan(tile).sum())