import cv2
import numpy as np

try:
    import scipy
    # scipy.ndimage cannot be accessed until explicitly imported
    from scipy import ndimage
except ImportError:
    scipy = None


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


def random_rotation(x, rg, row_axis=0, col_axis=1, channel_axis=2,
                    fill_mode='nearest', cval=0., interpolation_order=1):
    """Performs a random rotation of a Numpy image tensor.
    # Arguments
        x: Input tensor. Must be 3D.
        rg: Rotation range, in degrees.
        row_axis: Index of axis for rows in the input tensor.
        col_axis: Index of axis for columns in the input tensor.
        channel_axis: Index of axis for channels in the input tensor.
        fill_mode: Points outside the boundaries of the input
            are filled according to the given mode
            (one of `{'constant', 'nearest', 'reflect', 'wrap'}`).
        cval: Value used for points outside the boundaries
            of the input if `mode='constant'`.
        interpolation_order int: order of spline interpolation.
            see `ndimage.interpolation.affine_transform`
    # Returns
        Rotated Numpy image tensor.
    """
    theta = np.random.uniform(-rg, rg)
    x = apply_affine_transform(x, theta=theta, channel_axis=channel_axis,
                               fill_mode=fill_mode, cval=cval,
                               order=interpolation_order)
    return x


def random_shift(x, wrg, hrg, row_axis=0, col_axis=1, channel_axis=2,
                 fill_mode='nearest', cval=0., interpolation_order=1):
    """Performs a random spatial shift of a Numpy image tensor.
    # Arguments
        x: Input tensor. Must be 3D.
        wrg: Width shift range, as a float fraction of the width.
        hrg: Height shift range, as a float fraction of the height.
        row_axis: Index of axis for rows in the input tensor.
        col_axis: Index of axis for columns in the input tensor.
        channel_axis: Index of axis for channels in the input tensor.
        fill_mode: Points outside the boundaries of the input
            are filled according to the given mode
            (one of `{'constant', 'nearest', 'reflect', 'wrap'}`).
        cval: Value used for points outside the boundaries
            of the input if `mode='constant'`.
        interpolation_order int: order of spline interpolation.
            see `ndimage.interpolation.affine_transform`
    # Returns
        Shifted Numpy image tensor.
    """
    h, w = x.shape[row_axis], x.shape[col_axis]
    tx = np.random.uniform(-hrg, hrg) * h
    ty = np.random.uniform(-wrg, wrg) * w
    x = apply_affine_transform(x, tx=tx, ty=ty, channel_axis=channel_axis,
                               fill_mode=fill_mode, cval=cval,
                               order=interpolation_order)
    return x


def random_shear(x, intensity, row_axis=0, col_axis=1, channel_axis=2,
                 fill_mode='nearest', cval=0., interpolation_order=1):
    """Performs a random spatial shear of a Numpy image tensor.
    # Arguments
        x: Input tensor. Must be 3D.
        intensity: Transformation intensity in degrees.
        row_axis: Index of axis for rows in the input tensor.
        col_axis: Index of axis for columns in the input tensor.
        channel_axis: Index of axis for channels in the input tensor.
        fill_mode: Points outside the boundaries of the input
            are filled according to the given mode
            (one of `{'constant', 'nearest', 'reflect', 'wrap'}`).
        cval: Value used for points outside the boundaries
            of the input if `mode='constant'`.
        interpolation_order int: order of spline interpolation.
            see `ndimage.interpolation.affine_transform`
    # Returns
        Sheared Numpy image tensor.
    """
    shear = np.random.uniform(-intensity, intensity)
    x = apply_affine_transform(x, shear=shear, channel_axis=channel_axis,
                               fill_mode=fill_mode, cval=cval,
                               order=interpolation_order)
    return x


def random_zoom(x, zoom_range, row_axis=0, col_axis=1, channel_axis=2,
                fill_mode='nearest', cval=0., interpolation_order=1):
    """Performs a random spatial zoom of a Numpy image tensor.
    # Arguments
        x: Input tensor. Must be 3D.
        zoom_range: Tuple of floats; zoom range for width and height.
        row_axis: Index of axis for rows in the input tensor.
        col_axis: Index of axis for columns in the input tensor.
        channel_axis: Index of axis for channels in the input tensor.
        fill_mode: Points outside the boundaries of the input
            are filled according to the given mode
            (one of `{'constant', 'nearest', 'reflect', 'wrap'}`).
        cval: Value used for points outside the boundaries
            of the input if `mode='constant'`.
        interpolation_order int: order of spline interpolation.
            see `ndimage.interpolation.affine_transform`
    # Returns
        Zoomed Numpy image tensor.
    # Raises
        ValueError: if `zoom_range` isn't a tuple.
    """
    if len(zoom_range) != 2:
        raise ValueError('`zoom_range` should be a tuple or list of two'
                         ' floats. Received: %s' % (zoom_range,))

    if zoom_range[0] == 1 and zoom_range[1] == 1:
        zx, zy = 1, 1
    else:
        zx, zy = np.random.uniform(zoom_range[0], zoom_range[1], 2)
    x = apply_affine_transform(x, zx=zx, zy=zy, channel_axis=channel_axis,
                               fill_mode=fill_mode, cval=cval,
                               order=interpolation_order)
    return x


def random_channel_shift(x, intensity, channel_axis=0):
    x = np.rollaxis(x, channel_axis, 0)
    min_x, max_x = np.min(x), np.max(x)
    channel_images = [np.clip(x_channel + np.random.uniform(-intensity, intensity), min_x, max_x)
                      for x_channel in x]
    x = np.stack(channel_images, axis=0)
    x = np.rollaxis(x, 0, channel_axis + 1)
    return x


# For curving soybean pods. L.C.Uzal
def random_curves_transform(x, strength=0.1, range=(0.,255.)):
    low, high = range
    delta = (high - low) * strength / 2.
    xp = np.random.uniform(low=low + delta, high=high - delta)
    yp = np.random.uniform(low=xp-delta, high=xp+delta)
    xp = np.asarray([low, xp, high])
    yp = np.asarray([low, yp, high])
    return np.interp(x,xp,yp)


def transform_matrix_offset_center(matrix, x, y):
    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]])
    transform_matrix = np.dot(np.dot(offset_matrix, matrix), reset_matrix)
    return transform_matrix


def apply_affine_transform(x, theta=0, tx=0, ty=0, shear=0, zx=1, zy=1,
                           row_axis=0, col_axis=1, channel_axis=2,
                           fill_mode='nearest', cval=0., order=1):
    """Applies an affine transformation specified by the parameters given.
    # Arguments
        x: 2D numpy array, single image.
        theta: Rotation angle in degrees.
        tx: Width shift.
        ty: Heigh shift.
        shear: Shear angle in degrees.
        zx: Zoom in x direction.
        zy: Zoom in y direction
        row_axis: Index of axis for rows in the input image.
        col_axis: Index of axis for columns in the input image.
        channel_axis: Index of axis for channels in the input image.
        fill_mode: Points outside the boundaries of the input
            are filled according to the given mode
            (one of `{'constant', 'nearest', 'reflect', 'wrap'}`).
        cval: Value used for points outside the boundaries
            of the input if `mode='constant'`.
        order int: order of interpolation
    # Returns
        The transformed version of the input.
    """
    if scipy is None:
        raise ImportError('Image transformations require SciPy. '
                          'Install SciPy.')
    transform_matrix = None
    if theta != 0:
        theta = np.deg2rad(theta)
        rotation_matrix = np.array([[np.cos(theta), -np.sin(theta), 0],
                                    [np.sin(theta), np.cos(theta), 0],
                                    [0, 0, 1]])
        transform_matrix = rotation_matrix

    if tx != 0 or ty != 0:
        shift_matrix = np.array([[1, 0, tx],
                                 [0, 1, ty],
                                 [0, 0, 1]])
        if transform_matrix is None:
            transform_matrix = shift_matrix
        else:
            transform_matrix = np.dot(transform_matrix, shift_matrix)

    if shear != 0:
        shear = np.deg2rad(shear)
        shear_matrix = np.array([[1, -np.sin(shear), 0],
                                 [0, np.cos(shear), 0],
                                 [0, 0, 1]])
        if transform_matrix is None:
            transform_matrix = shear_matrix
        else:
            transform_matrix = np.dot(transform_matrix, shear_matrix)

    if zx != 1 or zy != 1:
        zoom_matrix = np.array([[zx, 0, 0],
                                [0, zy, 0],
                                [0, 0, 1]])
        if transform_matrix is None:
            transform_matrix = zoom_matrix
        else:
            transform_matrix = np.dot(transform_matrix, zoom_matrix)

    if transform_matrix is not None:
        h, w = x.shape[row_axis], x.shape[col_axis]
        transform_matrix = transform_matrix_offset_center(
            transform_matrix, h, w)
        x = np.rollaxis(x, channel_axis, 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=order,
            mode=fill_mode,
            cval=cval) for x_channel in x]
        x = np.stack(channel_images, axis=0)
        x = np.rollaxis(x, 0, channel_axis + 1)
    return x


def random_transform(x, rotation_range=0,
                     width_shift_range=0.,
                     height_shift_range=0.,
                     shear_range=0.,
                     zoom_range=0.,
                     channel_shift_range=0.,
                     horizontal_flip=False,
                     vertical_flip=False,
                     random_curves_strength=0.):
    # Generate params
    if rotation_range:
        theta = np.random.uniform(-rotation_range, rotation_range)
    else:
        theta = 0
    
    h, w = x.shape[0], x.shape[1]
    if height_shift_range:
        tx = np.random.uniform(-height_shift_range, height_shift_range) * h
    else:
        tx = 0
    if width_shift_range:
        ty = np.random.uniform(-width_shift_range, width_shift_range) * w
    else:
        ty = 0

    if shear_range:
        shear = np.random.uniform(-shear_range, shear_range)
    else:
        shear = 0

    if np.isscalar(zoom_range):
        zoom_range = [1 - zoom_range, 1 + zoom_range]
    elif len(zoom_range) == 2:
        zoom_range = [zoom_range[0], zoom_range[1]]
    else:
        raise ValueError('`zoom_range` should be a float or '
                         'a tuple or list of two floats. '
                         'Received arg: ', zoom_range)
    if zoom_range[0] == 1 and zoom_range[1] == 1:
        zx, zy = 1, 1
    else:
        zx, zy = np.random.uniform(zoom_range[0], zoom_range[1], 2)

    # Apply transforms
    x = apply_affine_transform(x,
                               theta,
                               tx, ty,
                               shear,
                               zx, zy)

    if channel_shift_range != 0:
        x = random_channel_shift(x, channel_shift_range, 2)

    if horizontal_flip:
        if np.random.random() < 0.5:
            x = flip_axis(x, 1)

    if vertical_flip:
        if np.random.random() < 0.5:
            x = flip_axis(x, 0)

    if random_curves_strength > 0.:
        x = random_curves_transform(x, random_curves_strength)

    return x


if __name__ == "__main__":
    import argparse
    from PIL import Image
    parser = argparse.ArgumentParser()
    parser.add_argument('--image', default='', type=str)
    parser.add_argument('--imageOut', default='result.png', type=str)
    args = parser.parse_args()

    im = np.array(Image.open(args.image))
    img = random_transform(im, rotation_range=10, shear_range=.5, zoom_range=.2, channel_shift_range=10., horizontal_flip=True)
    Image.fromarray(np.uint8(img)).save(args.imageOut)

[((1282, 1308), 'numpy.random.uniform', 'np.random.uniform', (['(-rg)', 'rg'], {}), '(-rg, rg)\n', (1299, 1308), True, 'import numpy as np\n'), ((3905, 3945), 'numpy.random.uniform', 'np.random.uniform', (['(-intensity)', 'intensity'], {}), '(-intensity, intensity)\n', (3922, 3945), True, 'import numpy as np\n'), ((5802, 5833), 'numpy.rollaxis', 'np.rollaxis', (['x', 'channel_axis', '(0)'], {}), '(x, channel_axis, 0)\n', (5813, 5833), True, 'import numpy as np\n'), ((6022, 6054), 'numpy.stack', 'np.stack', (['channel_images'], {'axis': '(0)'}), '(channel_images, axis=0)\n', (6030, 6054), True, 'import numpy as np\n'), ((6063, 6098), 'numpy.rollaxis', 'np.rollaxis', (['x', '(0)', '(channel_axis + 1)'], {}), '(x, 0, channel_axis + 1)\n', (6074, 6098), True, 'import numpy as np\n'), ((6286, 6339), 'numpy.random.uniform', 'np.random.uniform', ([], {'low': '(low + delta)', 'high': '(high - delta)'}), '(low=low + delta, high=high - delta)\n', (6303, 6339), True, 'import numpy as np\n'), ((6349, 6399), 'numpy.random.uniform', 'np.random.uniform', ([], {'low': '(xp - delta)', 'high': '(xp + delta)'}), '(low=xp - delta, high=xp + delta)\n', (6366, 6399), True, 'import numpy as np\n'), ((6405, 6432), 'numpy.asarray', 'np.asarray', (['[low, xp, high]'], {}), '([low, xp, high])\n', (6415, 6432), True, 'import numpy as np\n'), ((6442, 6469), 'numpy.asarray', 'np.asarray', (['[low, yp, high]'], {}), '([low, yp, high])\n', (6452, 6469), True, 'import numpy as np\n'), ((6481, 6501), 'numpy.interp', 'np.interp', (['x', 'xp', 'yp'], {}), '(x, xp, yp)\n', (6490, 6501), True, 'import numpy as np\n'), ((6630, 6677), 'numpy.array', 'np.array', (['[[1, 0, o_x], [0, 1, o_y], [0, 0, 1]]'], {}), '([[1, 0, o_x], [0, 1, o_y], [0, 0, 1]])\n', (6638, 6677), True, 'import numpy as np\n'), ((6697, 6746), 'numpy.array', 'np.array', (['[[1, 0, -o_x], [0, 1, -o_y], [0, 0, 1]]'], {}), '([[1, 0, -o_x], [0, 1, -o_y], [0, 0, 1]])\n', (6705, 6746), True, 'import numpy as np\n'), ((11067, 11090), 'numpy.isscalar', 'np.isscalar', (['zoom_range'], {}), '(zoom_range)\n', (11078, 11090), True, 'import numpy as np\n'), ((12259, 12284), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {}), '()\n', (12282, 12284), False, 'import argparse\n'), ((2630, 2658), 'numpy.random.uniform', 'np.random.uniform', (['(-hrg)', 'hrg'], {}), '(-hrg, hrg)\n', (2647, 2658), True, 'import numpy as np\n'), ((2672, 2700), 'numpy.random.uniform', 'np.random.uniform', (['(-wrg)', 'wrg'], {}), '(-wrg, wrg)\n', (2689, 2700), True, 'import numpy as np\n'), ((5476, 5526), 'numpy.random.uniform', 'np.random.uniform', (['zoom_range[0]', 'zoom_range[1]', '(2)'], {}), '(zoom_range[0], zoom_range[1], 2)\n', (5493, 5526), True, 'import numpy as np\n'), ((5853, 5862), 'numpy.min', 'np.min', (['x'], {}), '(x)\n', (5859, 5862), True, 'import numpy as np\n'), ((5864, 5873), 'numpy.max', 'np.max', (['x'], {}), '(x)\n', (5870, 5873), True, 'import numpy as np\n'), ((6777, 6806), 'numpy.dot', 'np.dot', (['offset_matrix', 'matrix'], {}), '(offset_matrix, matrix)\n', (6783, 6806), True, 'import numpy as np\n'), ((8171, 8188), 'numpy.deg2rad', 'np.deg2rad', (['theta'], {}), '(theta)\n', (8181, 8188), True, 'import numpy as np\n'), ((8474, 8519), 'numpy.array', 'np.array', (['[[1, 0, tx], [0, 1, ty], [0, 0, 1]]'], {}), '([[1, 0, tx], [0, 1, ty], [0, 0, 1]])\n', (8482, 8519), True, 'import numpy as np\n'), ((8787, 8804), 'numpy.deg2rad', 'np.deg2rad', (['shear'], {}), '(shear)\n', (8797, 8804), True, 'import numpy as np\n'), ((9178, 9223), 'numpy.array', 'np.array', (['[[zx, 0, 0], [0, zy, 0], [0, 0, 1]]'], {}), '([[zx, 0, 0], [0, zy, 0], [0, 0, 1]])\n', (9186, 9223), True, 'import numpy as np\n'), ((9648, 9679), 'numpy.rollaxis', 'np.rollaxis', (['x', 'channel_axis', '(0)'], {}), '(x, channel_axis, 0)\n', (9659, 9679), True, 'import numpy as np\n'), ((10039, 10071), 'numpy.stack', 'np.stack', (['channel_images'], {'axis': '(0)'}), '(channel_images, axis=0)\n', (10047, 10071), True, 'import numpy as np\n'), ((10084, 10119), 'numpy.rollaxis', 'np.rollaxis', (['x', '(0)', '(channel_axis + 1)'], {}), '(x, 0, channel_axis + 1)\n', (10095, 10119), True, 'import numpy as np\n'), ((10578, 10628), 'numpy.random.uniform', 'np.random.uniform', (['(-rotation_range)', 'rotation_range'], {}), '(-rotation_range, rotation_range)\n', (10595, 10628), True, 'import numpy as np\n'), ((10986, 11030), 'numpy.random.uniform', 'np.random.uniform', (['(-shear_range)', 'shear_range'], {}), '(-shear_range, shear_range)\n', (11003, 11030), True, 'import numpy as np\n'), ((11514, 11564), 'numpy.random.uniform', 'np.random.uniform', (['zoom_range[0]', 'zoom_range[1]', '(2)'], {}), '(zoom_range[0], zoom_range[1], 2)\n', (11531, 11564), True, 'import numpy as np\n'), ((12462, 12484), 'PIL.Image.open', 'Image.open', (['args.image'], {}), '(args.image)\n', (12472, 12484), False, 'from PIL import Image\n'), ((219, 232), 'numpy.asarray', 'np.asarray', (['x'], {}), '(x)\n', (229, 232), True, 'import numpy as np\n'), ((8712, 8750), 'numpy.dot', 'np.dot', (['transform_matrix', 'shift_matrix'], {}), '(transform_matrix, shift_matrix)\n', (8718, 8750), True, 'import numpy as np\n'), ((9089, 9127), 'numpy.dot', 'np.dot', (['transform_matrix', 'shear_matrix'], {}), '(transform_matrix, shear_matrix)\n', (9095, 9127), True, 'import numpy as np\n'), ((9413, 9450), 'numpy.dot', 'np.dot', (['transform_matrix', 'zoom_matrix'], {}), '(transform_matrix, zoom_matrix)\n', (9419, 9450), True, 'import numpy as np\n'), ((9809, 9937), 'scipy.ndimage.interpolation.affine_transform', 'ndimage.interpolation.affine_transform', (['x_channel', 'final_affine_matrix', 'final_offset'], {'order': 'order', 'mode': 'fill_mode', 'cval': 'cval'}), '(x_channel, final_affine_matrix,\n    final_offset, order=order, mode=fill_mode, cval=cval)\n', (9847, 9937), False, 'from scipy import ndimage\n'), ((10736, 10794), 'numpy.random.uniform', 'np.random.uniform', (['(-height_shift_range)', 'height_shift_range'], {}), '(-height_shift_range, height_shift_range)\n', (10753, 10794), True, 'import numpy as np\n'), ((10863, 10919), 'numpy.random.uniform', 'np.random.uniform', (['(-width_shift_range)', 'width_shift_range'], {}), '(-width_shift_range, width_shift_range)\n', (10880, 10919), True, 'import numpy as np\n'), ((11907, 11925), 'numpy.random.random', 'np.random.random', ([], {}), '()\n', (11923, 11925), True, 'import numpy as np\n'), ((11999, 12017), 'numpy.random.random', 'np.random.random', ([], {}), '()\n', (12015, 12017), True, 'import numpy as np\n'), ((5916, 5956), 'numpy.random.uniform', 'np.random.uniform', (['(-intensity)', 'intensity'], {}), '(-intensity, intensity)\n', (5933, 5956), True, 'import numpy as np\n'), ((12634, 12647), 'numpy.uint8', 'np.uint8', (['img'], {}), '(img)\n', (12642, 12647), True, 'import numpy as np\n'), ((8226, 8239), 'numpy.cos', 'np.cos', (['theta'], {}), '(theta)\n', (8232, 8239), True, 'import numpy as np\n'), ((8298, 8311), 'numpy.sin', 'np.sin', (['theta'], {}), '(theta)\n', (8304, 8311), True, 'import numpy as np\n'), ((8313, 8326), 'numpy.cos', 'np.cos', (['theta'], {}), '(theta)\n', (8319, 8326), True, 'import numpy as np\n'), ((8899, 8912), 'numpy.cos', 'np.cos', (['shear'], {}), '(shear)\n', (8905, 8912), True, 'import numpy as np\n'), ((8242, 8255), 'numpy.sin', 'np.sin', (['theta'], {}), '(theta)\n', (8248, 8255), True, 'import numpy as np\n'), ((8843, 8856), 'numpy.sin', 'np.sin', (['shear'], {}), '(shear)\n', (8849, 8856), True, 'import numpy as np\n')]