First attempt at a drift correction function for scikit-image

drift_correct

def drift_correct(image, shifts, shift_axes=[], kwargs*):
    """Align image array to correct for motion over time.
    
    Parameters
    ----------
    image : ndarray
        Image to drift correct.
    shifts : ndarray
        Amount to shift each image at each timepoint.
        Returned from skimage.feature.register_translation()
    shift_axes : list of int, optional
        Which axes of the image array are the spatial axes to shift.
    kwargs* : optional
        Keyword arguments for numpy pad() function.
    
    Returns
    -------
    output_image : ndarray
        Drift corrected image array.
    
    Examples
    --------
    ```
    >>> shifts = []
    >>> for frame in source_image:
    ...    shift, error, phasediff = feature.register_translation(frame, target)
    ...    shifts.append(shift)
    >>> shifts = np.array(shifts)
    >>> corrected_image = drift_correct(source_image, shifts, shift_axes=[1,2])
    ```
    """
    min_shifts = np.round(shifts).min(axis=0).clip(max=0).astype(np.int)
    max_shifts = np.round(shifts).max(axis=0).clip(min=0).astype(np.int)
    padding = list(zip(abs(min_shifts), max_shifts))
    no_shifts = np.zeros(shifts.shape[0])
    for axis in range(image.ndim):
        if axis not in shift_axes:
            padding.insert(axis, (0, 0))
            shifts = np.insert(shifts, axis, no_shifts, axis=1)
    # TODO: allow kwargs* to control padding mode (ie: not just 'constant')
    padded_image = np.pad(image, padding, 'constant')
    output_image = []
    # TODO: do not have hard coded values in array slicing here
    for frame, shift in zip(padded_image[...,0], shifts[:, 1:-1]):
        output_image.append(ndi.shift(frame, shift, order=0))
    output_image = np.array(output_image)
    # TODO: add option to trim padding back down with optional keyword arg
    return output_image

Some issues with the above:

• These loops are super slow, we should investigate a better way to do it faster.
• Needs kwargs* handling to pass arguments to the numpy pad function
• Could add extra functionality to trim the increased array size