imreg_with_mode.py

"""
including the mode (e.g. "reflect", "wrap") into transform_img  


"""


import numpy as np
import imreg_dft
from scipy.misc import ascent
import scipy.ndimage.interpolation as ndii
import imreg_dft.utils as utils




def _get_pads_and_slices(shape_src, shape_dest):
    """returns correct pads, slices for transforming shape_src to shape_dest"""
    diff = tuple(s1 - s2 for s1, s2 in zip(shape_dest, shape_src))
    slices = tuple(slice((-d) // 2, -(((-d) + 1) // 2)) if d < 0 else slice(None, None) for d in diff)
    pads = tuple((d // 2, d - (d // 2)) if d > 0 else (0, 0) for d in diff)
    return pads, slices


def _to_shape(img, shape_dest, mode="constant", bgval = 0 ):
    """resizes img to shape woth the correct padding mode"""
    shape_src = img.shape
    if shape_src == shape_dest:
        return img
    pads, slices = _get_pads_and_slices(shape_src, shape_dest)
    kwargs = {}
    kwargs.update({"constant_values":bgval} if mode == "constant" else {})
    return np.pad(img, pads, mode=mode, **kwargs)[slices]


# changes transform_img

def transform_img_with_mode(img, scale=1.0, angle=0.0, tvec=(0, 0),
                  mode="constant", bgval=None, order=1):
    """
    Return translation vector to register images.
    Args:
        img (2D or 3D numpy array): What will be transformed.
            If a 3D array is passed, it is treated in a manner in which RGB
            images are supposed to be handled - i.e. assume that coordinates
            are (Y, X, channels).
        scale (float): The scale factor (scale > 1.0 means zooming in)
        angle (float): Degrees of rotation (clock-wise)
        tvec (2-tuple): Pixel translation vector, Y and X component.
        mode (string): The transformation mode (refer to e.g.
            :func:`scipy.ndimage.shift` and its kwarg ``mode``).
        bgval (float): Shade of the background (filling during transformations)
            If None is passed, :func:`imreg_dft.utils.get_borderval` with
            radius of 5 is used to get it.
        order (int): Order of approximation (when doing transformations). 1 =
            linear, 3 = cubic etc. Linear works surprisingly well.
    Returns:
        np.ndarray: The transformed img, may have another
        i.e. (bigger) shape than the source.
    """
    if img.ndim == 3:
        # A bloody painful special case of RGB images
        ret = np.empty_like(img)
        for idx in range(img.shape[2]):
            sli = (slice(None), slice(None), idx)
            ret[sli] = transform_img(img[sli], scale, angle, tvec,
                                     mode, bgval, order)
        return ret

    if bgval is None:
        bgval = utils.get_borderval(img)

    bigshape = list(np.round(np.array(img.shape) * 1.2).astype(int))
    # change here
    dest0 = _to_shape(img.copy(), bigshape, mode=mode, bgval = bgval)

    if scale != 1.0:
        dest0 = ndii.zoom(dest0, scale, order=order, mode=mode, cval=bgval)
    if angle != 0.0:
        dest0 = ndii.rotate(dest0, angle, order=order, mode=mode, cval=bgval)

    if tvec[0] != 0 or tvec[1] != 0:
        dest0 = ndii.shift(dest0, tvec, order=order, mode=mode, cval=bgval)

    # change here
    dest = _to_shape(dest0, img.shape, mode=mode,bgval = bgval)

    return dest



im = ascent()

# baseline from ndimage
out = ndii.shift(im, [100,100], mode ="reflect")

# mode is ignored here 
out1 = imreg_dft.transform_img(im, tvec=[100,100], mode="reflect")
# this works
out2 = transform_img_with_mode(im, tvec=[100,100], mode="reflect")