augmentations.py

from scipy.misc import imresize

import utils
import numpy as np

from scipy import ndimage
from utils import to_int


def get_random_transform_params(input_shape, rotation_range = 0., height_shift_range = 0., width_shift_range = 0.,
                                shear_range = 0., zoom_range = (1, 1), horizontal_flip = False, resize_range = None,
                                distortion_prob = 0., additive_gaussian_noise_range = None, multiplication_gaussian = 0,
                                transform_colorspace_param = None, transform_colorspace_bounds = (-1, 1)):
    """
    This closure function returns generative function that gets random instance trough parameter and
    together with closed input parameters generates random parameters for transformation matrix.

    :param distortion_prob: Probability of the downsampling and upsampling of the image
    :param resize_range: Defines uniform interval of downsampling factor
    :param input_shape: Shape of images to be transformed with matrix with this parameters
    :param rotation_range: Interval of rotation in degrees (used for in both direction)
    :param height_shift_range: Value of two-sided interval of random shift in vertical direction
    :param width_shift_range: Value of two-sided interval of random shift in horizontal direction
    :param shear_range: Value of two-sided interval of random shear in horizontal direction
    :param zoom_range: Tuple with 2 values representing range of random zoom (values > 1.0 is for zoom out)
    :param horizontal_flip: Whether do random horizontal flip image
    :return: Function that with given random instance generates random parameters for transformation matrix
    """

    def get_instance(rnd):
        U = rnd.uniform
        N = rnd.normal

        rr = rotation_range
        hs = height_shift_range
        ws = width_shift_range
        sr = shear_range
        agn = additive_gaussian_noise_range
        dp = distortion_prob
        mg = multiplication_gaussian
        tcp = transform_colorspace_param
        tcb = transform_colorspace_bounds

        return {
            'input_shape': input_shape,
            'theta': np.pi / 180 * U(-rr, rr) if rr else 0,
            'ty': U(-hs, hs) * input_shape[0] if hs else 0,
            'tx': U(-ws, ws) * input_shape[1] if ws else 0,
            'shear': U(-sr, sr) if shear_range else 0,
            'z': U(zoom_range[0], zoom_range[1]) if zoom_range != (1, 1) else 1,
            'h_flip': rnd.rand() < 0.5 if horizontal_flip else False,
            'add_noise': N(0, U(agn[0], agn[1]), input_shape) if agn is not None else None,
            'resize': U(*resize_range) if U(0, 1) < dp else None,
            'resize_smooth': U(0, 1) < 0.5,
            'mul': N(1, mg) if mg > 0 else None,
            'color_m': utils.crop_value(N(tcp[0], tcp[1], (3, 3)), tcb) if tcp is not None else None,
            'agn': agn
        }

    return get_instance


def assemble_transformation_matrix(input_shape, theta = 0, tx = 0, ty = 0, shear = 0, z = 1):
    """
    Creates transformation matrix with given parameters. That resulting matrix has origin in centre of the image

    :param input_shape: Shape of images to be transformed with matrix. Origin of transformation matrix is set
                        in the middle of image.

    :param theta: Rotation in radians
    :param tx: Translation in X axis
    :param ty: Translation in Y axis
    :param shear: Shear in horizontal direction
    :param z: Image zoom
    :return: Transformation matrix
    """

    def transform_matrix_offset_center(matrix, x, y):
        """
        Creates translation matrix from input matrix with origin in the centre of image

        :param matrix: Input matrix
        :param x: Width of the image
        :param y: Height of the image
        :return: Returns shifted input matrix with origin in [y/2, x/2]
        """
        o_x = float(x) / 2 + 0.5
        o_y = float(y) / 2 + 0.5
        offset_matrix = np.array([[1, 0, o_x], [0, 1, o_y], [0, 0, 1]])
        reset_matrix = np.array([[1, 0, -o_x], [0, 1, -o_y], [0, 0, 1]])
        t_matrix = np.dot(np.dot(offset_matrix, matrix), reset_matrix)
        return t_matrix

    rotation_matrix = np.array([[np.cos(theta), -np.sin(theta), 0],
                                [np.sin(theta), np.cos(theta), 0],
                                [0, 0, 1]])

    translation_matrix = np.array([[1, 0, ty],
                                   [0, 1, tx],
                                   [0, 0, 1]])

    shear_matrix = np.array([[1, -np.sin(shear), 0],
                             [0, np.cos(shear), 0],
                             [0, 0, 1]])

    zoom_matrix = np.array([[z, 0, 0],
                            [0, z, 0],
                            [0, 0, 1]])

    # Assembling transformation matrix
    transform_matrix = np.dot(np.dot(np.dot(rotation_matrix, translation_matrix), shear_matrix), zoom_matrix)

    # Set origin of transformation to center of the image
    h, w = input_shape[0], input_shape[1]
    transform_matrix = transform_matrix_offset_center(transform_matrix, h, w)

    return transform_matrix


def transform(v, t_matrix, h_flip = False, add_noise = None, resize = None, resize_smooth = None,
              mul = None, color_m = None):
    """
    Transform image with (inverted) transformation matrix

    :param v: Input image to be transformed
    :param t_matrix: Transformation matrix
    :param h_flip: Whether do horizontal flip
    :return: Transformed image
    """

    def apply_transform(x, transform_matrix, channel_index = 0, fill_mode = 'nearest', cval = 0.):
        x = np.rollaxis(x, channel_index, 0)
        final_affine_matrix = transform_matrix[:2, :2]
        final_offset = transform_matrix[:2, 2]
        channel_images = [ndimage.interpolation.affine_transform(x_channel, final_affine_matrix,
                                                                 final_offset, order = 2, mode = fill_mode,
                                                                 cval = cval)
                          for x_channel in x]
        x = np.stack(channel_images, axis = 0)
        x = np.rollaxis(x, 0, channel_index + 1)

        return x

    def flip_axis(x, axis):
        x = np.asarray(x).swapaxes(axis, 0)
        x = x[::-1, ...]
        x = x.swapaxes(0, axis)
        return x

    v = apply_transform(v, t_matrix, 2)

    if h_flip:
        v = flip_axis(v, 1)

    if color_m is not None or mul is not None or add_noise is not None:
        v = v.astype(np.float32)
        shape = v.shape

        if mul is not None:
            v *= mul

        if color_m is not None:
            v = np.reshape(v, [-1, 3])
            v = np.matmul(v, color_m)
            v = np.reshape(v, shape)

        if add_noise is not None:
            v += add_noise

        if resize is not None:
            interpolation = 'bilinear' if resize_smooth else 'nearest'

            v = imresize(v, (resize * np.array(shape[:2])).astype(np.uint16), interpolation)
            v = imresize(v, shape[:2], interpolation)

        v = utils.crop_value(v, [np.zeros(shape), np.ones(shape) * 255])

    return v.astype(np.uint8)


def crop_data(img, labels, new_img_size, new_label_size = None, crop_label = True):
    """
    Both images and labels will be cropped to match the given size

    :param img: Images to be cropped
    :param labels: Labels to be cropped
    :param new_img_size: New image size
    :param new_label_size: New labels size
    :return: Cropped image and labels
    """

    img_size = img.shape[-3]
    r = to_int((img_size - new_img_size) / 2)

    img = img[..., r:r + new_img_size, r:r + new_img_size, :]

    if crop_label:
        labels -= r

    if new_label_size is not None:
        labels = np.array((labels / new_img_size) * new_label_size, dtype = np.int32)

    return img, labels


def flip_body_joints(points):
    """
    Change semantic of labels after flip transformation - i.e. left leg will be now right and so on.

    :param points: Body joints to be changed
    """

    def swap(a, b):
        points[:, [a, b]] = points[:, [b, a]]

    # Leg
    swap(0, 5)
    swap(1, 4)
    swap(2, 3)

    # Arm
    swap(10, 15)
    swap(11, 14)
    swap(12, 13)


def generate_random__transformation(X, rseed = 0, t_params_f = None):
    rnd = np.random.RandomState(rseed)

    if not t_params_f:
        raise Exception('No attributes given!')

    n = X.shape[0]
    X_t = []

    t_params = t_params_f(rnd)
    h_flip = t_params.pop('h_flip')
    add_noise = t_params.pop('add_noise')
    resize = t_params.pop('resize')
    mul = t_params.pop('mul')
    agn = t_params.pop('agn')
    color_m = t_params.pop('color_m')
    resize_smooth = t_params.pop('resize_smooth')

    t_matrix = assemble_transformation_matrix(**t_params)

    for k in range(n):
        inp = np.squeeze(X[k])

        if agn is not None:
            gauss = rnd.normal(0, rnd.uniform(agn[0], agn[1]), inp.shape)
        else:
            gauss = None

        x_t = transform(inp, t_matrix, h_flip, gauss, resize, resize_smooth, mul, color_m)
        X_t.append(x_t)

    return np.array(X_t)


def generate_random_sequences(X, Y, sequence_size = 32, shift = 16, rseed = 0, final_size = None,
                              t_params_f = None, final_heatmap_size = None):
    rnd = np.random.RandomState(rseed)

    if not t_params_f:
        raise Exception('No attributes given!')

    if final_size is None:
        final_size = min(X.shape[2], X.shape[3])

    n = X.shape[0]
    perm = rnd.permutation(range(0, n, shift))
    perm_n = perm.shape[0]

    for idx in range(perm_n):
        b = range(perm[idx], min(perm[idx] + sequence_size, n))

        X_t = []
        Y_t = []

        t_params = t_params_f(rnd)
        h_flip = t_params.pop('h_flip')
        add_noise = t_params.pop('add_noise')
        resize = t_params.pop('resize')
        mul = t_params.pop('mul')
        agn = t_params.pop('agn')
        color_m = t_params.pop('color_m')
        resize_smooth = t_params.pop('resize_smooth')

        t_matrix = assemble_transformation_matrix(**t_params)

        for k in b:
            inp = np.squeeze(X[k])

            if agn is not None:
                gauss = rnd.normal(0, rnd.uniform(agn[0], agn[1]), inp.shape)
            else:
                gauss = None

            x_t = transform(inp, t_matrix, h_flip, gauss, resize, resize_smooth, mul, color_m)
            y_t = utils.get_affine_transform(np.squeeze(Y[k]), np.linalg.inv(t_matrix)) if Y is not None else None

            x_t, y_t = crop_data(x_t, y_t, final_size, final_heatmap_size)

            X_t.append(x_t)
            if Y is not None:
                if h_flip:
                    y_t[1, :] = (final_size if final_heatmap_size is None else final_heatmap_size) - y_t[1, :]
                    flip_body_joints(y_t)

                Y_t.append(y_t)

        if Y is not None:
            yield np.array(X_t), np.array(Y_t), idx
        else:
            yield np.array(X_t), idx


def generate_minibatches(X, Y = None, batch_size = 32, rseed = 0,
                         final_size = None, t_params_f = None, final_heatmap_size = None):
    """
    This function splits whole input batch of images into minibatches of given size. All images in batch are
    transformed using affine transformations in order to prevent over-fitting during training.

    :param X: Batch of input images to be divided. It has to be 4D tensor [batch, channel, height, width]
    :param Y: Labels of input images (joint positions on heatmap). 3D tensor [batch, image dimension, joint].
              E.g. joint with index 4 present in 10th image (i.e. index 9) that is in position [50, 80] is in
              indexes: Y[9, :, 4] == [50, 80]
    :param batch_size: Size of each mini-batch
    :param rseed: Random seed
    :param t_params_f: Function that generates parameters for transformation matrix (see get_random_transform_params)
    :param final_size: Transformed images are cropped to match the given size
    :param final_heatmap_size: Size of heatmaps
    :return: Sequence of randomly ordered and transformed mini-batches
    """

    rnd = np.random.RandomState(rseed)

    if not t_params_f:
        raise Exception('No attributes given!')

    if final_size is None:
        final_size = min(X.shape[2], X.shape[3])

    n = X.shape[0]
    perm = rnd.permutation(n)

    for idx in range(0, n, batch_size):
        b = perm[idx:min(idx + batch_size, n)]

        X_t = []
        Y_t = []

        for k in b:
            t_params = t_params_f(rnd)
            h_flip = t_params.pop('h_flip')
            add_noise = t_params.pop('add_noise')
            resize = t_params.pop('resize')
            mul = t_params.pop('mul')
            color_m = t_params.pop('color_m')
            agn = t_params.pop('agn')
            resize_smooth = t_params.pop('resize_smooth')

            t_matrix = assemble_transformation_matrix(**t_params)

            x_t = transform(np.squeeze(X[k]), t_matrix, h_flip, add_noise, resize, resize_smooth, mul, color_m)
            y_t = utils.get_affine_transform(np.squeeze(Y[k]), np.linalg.inv(t_matrix)) if Y is not None else None

            x_t, y_t = crop_data(x_t, y_t, final_size, final_heatmap_size)

            X_t.append(x_t)
            if Y is not None:
                if h_flip:
                    y_t[1, :] = (final_size if final_heatmap_size is None else final_heatmap_size) - y_t[1, :]
                    flip_body_joints(y_t)

                Y_t.append(y_t)

        if Y is not None:
            yield np.array(X_t), np.array(Y_t), b
        else:
            yield np.array(X_t), b