data_augmentation

from scipy.ndimage import zoom
from skimage.transform import rotate
import numpy as np

   def crop(image, label, patch_size):
        img_size = image.shape[1:]
        scale_range = [0.8, 1.2]
        scale = np.random.rand() * (
                scale_range[1] - scale_range[0]) + scale_range[0]
        crop_size = (
                np.array(patch_size).astype('float') / scale).astype('int')

        start = []
        for i in range(2):
            if crop_size[i] > img_size[i]:
                start.append(np.random.randint(int((img_size[i]-crop_size[i])/2)-1, 0))
            else:
                start.append(np.random.randint(0, img_size[i]-crop_size[i]+1))

        pad = [[0, 0]]
        for i in range(2):
            left_pad = max(0, -start[i])
            right_pad = max(0, start[i] + crop_size[i] - img_size[i])
            pad.append([left_pad, right_pad])
        crop = image[:,
               max(start[0], 0):min(start[0] + crop_size[0], img_size[0]),
               max(start[1], 0):min(start[1] + crop_size[1], img_size[1])]
        crop = np.pad(crop, pad, 'constant')
        crop_label = label[max(start[0], 0):min(start[0] + crop_size[0], img_size[0]),
                     max(start[1], 0):min(start[1] + crop_size[1], img_size[1])]
        crop_label = np.pad(crop_label, pad[1:], 'constant')

        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            crop = zoom(crop, [1, scale, scale], order=1)
            crop_label = zoom(crop_label, [scale, scale], order=1)
        newpad = patch_size[0] - crop.shape[1:][0]
        if newpad < 0:
            crop = crop[:, :newpad, :newpad]
            crop_label = crop_label[:newpad, :newpad]
        elif newpad > 0:
            pad2 = [[0, 0], [0, newpad], [0, newpad]]
            crop = np.pad(crop, pad2, 'constant')
            crop_label = np.pad(crop_label, pad2[1:], 'constant')
        return crop, crop_label
        
        
    def random_rotate(sample, target, angles_range):
        angles = np.float32(np.random.uniform(*angles_range))
        rot_sample = sample.copy()
        rot_target = target.copy()
        for index in range(3):
            sample_channel = sample[:, :, index]
            target_channel = target[:, :, index]
            sample_channel = rotate(sample_channel, angles, resize=False, preserve_range=True)
            target_channel = rotate(target_channel, angles, resize=False, preserve_range=True)
            rot_sample[:, :, index] = sample_channel
            rot_target[:, :, index] = target_channel
        return np.float32(rot_sample), np.float32(rot_target) 
    
    def flip(sample, target):
        flip_num = np.random.randint(0, 8)
        if flip_num == 1:
            sample = np.flipud(sample)
            target = np.flipud(target)
        elif flip_num == 2:
            sample = np.fliplr(sample)
            target = np.fliplr(target)
        elif flip_num == 3:
            sample = np.rot90(sample, k=1, axes=(1, 0))
            target = np.rot90(target, k=1, axes=(1, 0))
        elif flip_num == 4:
            sample = np.rot90(sample, k=3, axes=(1, 0))
            target = np.rot90(target, k=3, axes=(1, 0))
        elif flip_num == 5:
            sample = np.fliplr(sample)
            target = np.fliplr(target)
            sample = np.rot90(sample, k=1, axes=(1, 0))
            target = np.rot90(target, k=1, axes=(1, 0))
        elif flip_num == 6:
            sample = np.fliplr(sample)
            target = np.fliplr(target)
            sample = np.rot90(sample, k=3, axes=(1, 0))
            target = np.rot90(target, k=3, axes=(1, 0))
        elif flip_num == 7:
            sample = np.flipud(sample)
            target = np.flipud(target)
            sample = np.fliplr(sample)
            target = np.fliplr(target)
        return sample, target
    

def FlipH(img):
    return np.fliplr ( img )


def FlipV(img):
    return np.flipud ( img )


def Rotate(img, angle):
    return transform.rotate ( img, angle )


def Blur(img, sigma=0.7):
    is_colour = len ( img.shape ) == 3
    return rescale_intensity ( gaussian ( img, sigma=sigma, multichannel=is_colour ) )

def Noise(img, var=0.00001):
    return random_noise ( img, mode='gaussian', var=var )


class Zoom:
    def __init__(self, p1x, p1y, p2x, p2y):
        self.p1x = p1x
        self.p1y = p1y
        self.p2x = p2x
        self.p2y = p2y

    def process(self, img):
        h = len ( img )
        w = len ( img[0] )

        crop_p1x = max ( self.p1x, 0 )
        crop_p1y = max ( self.p1y, 0 )
        crop_p2x = min ( self.p2x, w )
        crop_p2y = min ( self.p2y, h )

        cropped_img = img[crop_p1y:crop_p2y, crop_p1x:crop_p2x]

        x_pad_before = -min ( 0, self.p1x )
        x_pad_after = max ( 0, self.p2x - w )
        y_pad_before = -min ( 0, self.p1y )
        y_pad_after = max ( 0, self.p2y - h )

        padding = [(y_pad_before, y_pad_after), (x_pad_before, x_pad_after)]
        is_colour = len ( img.shape ) == 3
        if is_colour:
            padding.append ( (0, 0) )  # colour images have an extra dimension

        padded_img = np.pad ( cropped_img, padding, 'constant' )
        return transform.resize ( padded_img, (h, w) )


class RandomScale ( object ):
    """Randomly scales an image
    Parameters
    ----------
    scale: float or tuple(float)
        if **float**, the image is scaled by a factor drawn
        randomly from a range (1 - `scale` , 1 + `scale`). If **tuple**,
        the `scale` is drawn randomly from values specified by the
        tuple

    Returns
    -------

    numpy.ndaaray
        Scaled image in the numpy format of shape `HxWxC`

    """

    def __init__(self, scale=0.2, diff=False):
        self.scale = scale

        if type ( self.scale ) == tuple:
            assert len ( self.scale ) == 2, "Invalid range"
            assert self.scale[0] > -1, "Scale factor can't be less than -1"
            assert self.scale[1] > -1, "Scale factor can't be less than -1"
        else:
            assert self.scale > 0, "Please input a positive float"
            self.scale = (max ( -1, -self.scale ), self.scale)

        self.diff = diff

    def __call__(self, img):
        # Chose a random digit to scale by
        img_shape = img.shape
        if self.diff:
            scale_x = random.uniform ( *self.scale )
            scale_y = random.uniform ( *self.scale )
        else:
            scale_x = random.uniform ( *self.scale )
            scale_y = scale_x
        resize_scale_x = 1 + scale_x
        resize_scale_y = 1 + scale_y
        img = cv2.resize ( img, None, fx=resize_scale_x, fy=resize_scale_y )
        canvas = np.zeros ( img_shape, dtype=np.uint8 )
        y_lim = int ( min ( resize_scale_y, 1 ) * img_shape[0] )
        x_lim = int ( min ( resize_scale_x, 1 ) * img_shape[1] )
        canvas[:y_lim, :x_lim, :] = img[:y_lim, :x_lim, :]
        img = canvas
        return img


class RandomRotate ( object ):
    """Randomly rotates an image

        Parameters
        ----------
        angle: float or tuple(float)
            if **float**, the image is rotated by a factor drawn
            randomly from a range (-`angle`, `angle`). If **tuple**,
            the `angle` is drawn randomly from values specified by the
            tuple

        Returns
        -------

        numpy.ndaaray
            Rotated image in the numpy format of shape `HxWxC`

        """

    def __init__(self, angle=10):
        self.angle = angle

    def __call__(self, img):
        angles = np.float32 ( np.random.uniform ( *self.angle ) )
        rot_sample = img.copy ( )
        for index in range ( 3 ):
            sample_channel = img[:, :, index]
            sample_channel = rotate ( sample_channel, angles, resize=False, preserve_range=True )
            rot_sample[:, :, index] = sample_channel
        return np.float32 ( rot_sample )


class RandomHSV ( object ):
    """HSV Transform to vary hue saturation and brightness

    Hue has a range of 0-179
    Saturation and Brightness have a range of 0-255.
    Chose the amount you want to change thhe above quantities accordingly.




    Parameters
    ----------
    hue : None or int or tuple (int)
        If None, the hue of the image is left unchanged. If int,
        a random int is uniformly sampled from (-hue, hue) and added to the
        hue of the image. If tuple, the int is sampled from the range
        specified by the tuple.

    saturation : None or int or tuple(int)
        If None, the saturation of the image is left unchanged. If int,
        a random int is uniformly sampled from (-saturation, saturation)
        and added to the hue of the image. If tuple, the int is sampled
        from the range  specified by the tuple.

    brightness : None or int or tuple(int)
        If None, the brightness of the image is left unchanged. If int,
        a random int is uniformly sampled from (-brightness, brightness)
        and added to the hue of the image. If tuple, the int is sampled
        from the range  specified by the tuple.

    Returns
    -------

    numpy.ndaaray
        Transformed image in the numpy format of shape `HxWxC`

    numpy.ndarray
        Resized bounding box co-ordinates of the format `n x 4` where n is
        number of bounding boxes and 4 represents `x1,y1,x2,y2` of the box

    """

    def __init__(self, hue=None, saturation=None, brightness=None):
        if hue:
            self.hue = hue
        else:
            self.hue = 0

        if saturation:
            self.saturation = saturation
        else:
            self.saturation = 0

        if brightness:
            self.brightness = brightness
        else:
            self.brightness = 0

        if type ( self.hue ) != tuple:
            self.hue = (-self.hue, self.hue)

        if type ( self.saturation ) != tuple:
            self.saturation = (-self.saturation, self.saturation)

        if type ( brightness ) != tuple:
            self.brightness = (-self.brightness, self.brightness)

    def __call__(self, img):

        hue = random.randint ( *self.hue )
        saturation = random.randint ( *self.saturation )
        brightness = random.randint ( *self.brightness )

        img = img.astype ( int )

        a = np.array ( [hue, saturation, brightness] ).astype ( int )
        img += np.reshape ( a, (1, 1, 3) )

        img = np.clip ( img, 0, 255 )
        img[:, :, 0] = np.clip ( img[:, :, 0], 0, 179 )

        img = img.astype ( np.uint8 )

        return img