# Script for data augmentation functions

import numpy as np
from collections import deque
from PIL import Image
import cv2
from data.config import *

def imread_cv2(image_path):
    """
    Read image_path with cv2 format (H, W, C)
    if image is '.gif' outputs is a numpy array of {0,1}
    """
    image_format = image_path[-3:]
    if image_format == 'jpg':
        image = cv2.imread(image_path)
    else:
        image = np.array(Image.open(image_path))

    return image

def resize_cv2(image, heigh=1280, width=1918):
    return cv2.resize(image, (width, heigh), cv2.INTER_LINEAR)

def image_to_tensor(image, mean=0, std=1.):
    """Transform image (input is numpy array, read in by cv2) """
    if len(image.shape) == 2:
        image = image.reshape(image.shape[0], image.shape[1], 1)
    image = image.astype(np.float32)
    image = (image-mean)/std
    image = image.transpose((2,0,1))
    tensor = torch.from_numpy(image)

    return tensor

# --- Data Augmentation functions --- #
# A lot of functions can be found here:
# https://github.com/fchollet/keras/blob/master/keras/preprocessing/image.py#L223

# transform image and label
def randomHorizontalFlip(image, mask, p=0.5):
    """Do a random horizontal flip with probability p"""
    if np.random.random() < p:
        image = np.fliplr(image)
        mask = np.fliplr(mask)
    return image, mask

def randomVerticalFlip(image, mask, p=0.5):
    """Do a random vertical flip with probability p"""
    if np.random.random() < p:
        image = np.flipud(image)
        mask = np.flipud(mask)
    return image, mask

def randomHorizontalShift(image, mask, max_shift=0.05, p=0.5):
    """Do random horizontal shift with max proportion shift and with probability p
    Elements that roll beyond the last position are re-introduced at the first."""
    max_shift_pixels = int(max_shift*image.shape[1])
    shift = np.random.choice(np.arange(-max_shift_pixels, max_shift_pixels+1))
    if np.random.random() < p:
        image = np.roll(image, shift, axis=1)
        mask = np.roll(mask, shift, axis=1)
    return image, mask

def randomVerticalShift(image, mask, max_shift=0.05, p=0.5):
    """Do random vertical shift with max proportion shift and probability p
    Elements that roll beyond the last position are re-introduced at the first."""
    max_shift_pixels = int(max_shift*image.shape[0])
    shift = np.random.choice(np.arange(-max_shift_pixels, max_shift_pixels+1))
    if np.random.random() < p:
            image = np.roll(image, shift, axis=0)
            mask = np.roll(mask, shift, axis=0)
    return image, mask

def randomInvert(image, mask, p=0.5):
    """Randomly invert image with probability p"""
    if np.random.random() < p:
        image = 255 - image
        mask = mask
    return image, mask

def randomBrightness(image, mask, p=0.75):
    """With probability p, randomly increase or decrease brightness.
    See https://stackoverflow.com/questions/37822375/python-opencv-increasing-image-brightness-without-overflowing-uint8-array"""
    if np.random.random() < p:
        max_value = np.percentile(255-image, q=25) # avoid burning out white cars, so take image-specific maximum
        value = np.random.choice(np.arange(-max_value, max_value))
        if value > 0:
            image = np.where((255 - image) < value,255,image+value).astype(np.uint8)
        else:
            image = np.where(image < -value,0,image+value).astype(np.uint8)
    return image, mask

def randomHue(image, mask, p=0.25, max_value=75):
    """With probability p, randomly increase or decrease hue.
    See https://stackoverflow.com/questions/32609098/how-to-fast-change-image-brightness-with-python-opencv"""

    if np.random.random() < p:
        value = np.random.choice(np.arange(-max_value, max_value))
        hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        hsv[:,:,0] = hsv[:,:,0] + value
        hsv = np.clip(hsv, a_min=0, a_max=255).astype(np.uint8)
        image = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)

    return image, mask

def GaussianBlur(image, mask, kernel=(1, 1),sigma=1,  p=0.5):
    """With probability p, apply Gaussian blur"""
    # TODO
    return image, mask

def randomRotate(image, mask, max_angle, p=0.5):
    """Perform random rotation with max_angle and probability p"""
    # TODO

    return(image, mask)

[((23, 11, 23, 62), 'cv2.resize', 'cv2.resize', ({(23, 22, 23, 27): 'image', (23, 29, 23, 43): '(width, heigh)', (23, 45, 23, 61): 'cv2.INTER_LINEAR'}, {}), '(image, (width, heigh), cv2.INTER_LINEAR)', False, 'import cv2\n'), ((16, 16, 16, 38), 'cv2.imread', 'cv2.imread', ({(16, 27, 16, 37): 'image_path'}, {}), '(image_path)', False, 'import cv2\n'), ((43, 7, 43, 25), 'numpy.random.random', 'np.random.random', ({}, {}), '()', True, 'import numpy as np\n'), ((44, 16, 44, 32), 'numpy.fliplr', 'np.fliplr', ({(44, 26, 44, 31): 'image'}, {}), '(image)', True, 'import numpy as np\n'), ((45, 15, 45, 30), 'numpy.fliplr', 'np.fliplr', ({(45, 25, 45, 29): 'mask'}, {}), '(mask)', True, 'import numpy as np\n'), ((50, 7, 50, 25), 'numpy.random.random', 'np.random.random', ({}, {}), '()', True, 'import numpy as np\n'), ((51, 16, 51, 32), 'numpy.flipud', 'np.flipud', ({(51, 26, 51, 31): 'image'}, {}), '(image)', True, 'import numpy as np\n'), ((52, 15, 52, 30), 'numpy.flipud', 'np.flipud', ({(52, 25, 52, 29): 'mask'}, {}), '(mask)', True, 'import numpy as np\n'), ((59, 29, 59, 77), 'numpy.arange', 'np.arange', ({(59, 39, 59, 56): '-max_shift_pixels', (59, 58, 59, 76): 'max_shift_pixels + 1'}, {}), '(-max_shift_pixels, max_shift_pixels + 1)', True, 'import numpy as np\n'), ((60, 7, 60, 25), 'numpy.random.random', 'np.random.random', ({}, {}), '()', True, 'import numpy as np\n'), ((61, 16, 61, 45), 'numpy.roll', 'np.roll', (), '', True, 'import numpy as np\n'), ((62, 15, 62, 43), 'numpy.roll', 'np.roll', (), '', True, 'import numpy as np\n'), ((69, 29, 69, 77), 'numpy.arange', 'np.arange', ({(69, 39, 69, 56): '-max_shift_pixels', (69, 58, 69, 76): 'max_shift_pixels + 1'}, {}), '(-max_shift_pixels, max_shift_pixels + 1)', True, 'import numpy as np\n'), ((70, 7, 70, 25), 'numpy.random.random', 'np.random.random', ({}, {}), '()', True, 'import numpy as np\n'), ((71, 20, 71, 49), 'numpy.roll', 'np.roll', (), '', True, 'import numpy as np\n'), ((72, 19, 72, 47), 'numpy.roll', 'np.roll', (), '', True, 'import numpy as np\n'), ((77, 7, 77, 25), 'numpy.random.random', 'np.random.random', ({}, {}), '()', True, 'import numpy as np\n'), ((85, 7, 85, 25), 'numpy.random.random', 'np.random.random', ({}, {}), '()', True, 'import numpy as np\n'), ((86, 20, 86, 50), 'numpy.percentile', 'np.percentile', (), '', True, 'import numpy as np\n'), ((98, 7, 98, 25), 'numpy.random.random', 'np.random.random', ({}, {}), '()', True, 'import numpy as np\n'), ((100, 14, 100, 52), 'cv2.cvtColor', 'cv2.cvtColor', ({(100, 27, 100, 32): 'image', (100, 34, 100, 51): 'cv2.COLOR_BGR2HSV'}, {}), '(image, cv2.COLOR_BGR2HSV)', False, 'import cv2\n'), ((103, 16, 103, 52), 'cv2.cvtColor', 'cv2.cvtColor', ({(103, 29, 103, 32): 'hsv', (103, 34, 103, 51): 'cv2.COLOR_HSV2BGR'}, {}), '(hsv, cv2.COLOR_HSV2BGR)', False, 'import cv2\n'), ((18, 25, 18, 47), 'PIL.Image.open', 'Image.open', ({(18, 36, 18, 46): 'image_path'}, {}), '(image_path)', False, 'from PIL import Image\n'), ((87, 33, 87, 65), 'numpy.arange', 'np.arange', ({(87, 43, 87, 53): '-max_value', (87, 55, 87, 64): 'max_value'}, {}), '(-max_value, max_value)', True, 'import numpy as np\n'), ((99, 33, 99, 65), 'numpy.arange', 'np.arange', ({(99, 43, 99, 53): '-max_value', (99, 55, 99, 64): 'max_value'}, {}), '(-max_value, max_value)', True, 'import numpy as np\n'), ((102, 14, 102, 46), 'numpy.clip', 'np.clip', (), '', True, 'import numpy as np\n'), ((89, 20, 89, 67), 'numpy.where', 'np.where', ({(89, 29, 89, 50): '255 - image < value', (89, 51, 89, 54): '255', (89, 55, 89, 66): 'image + value'}, {}), '(255 - image < value, 255, image + value)', True, 'import numpy as np\n'), ((91, 20, 91, 58), 'numpy.where', 'np.where', ({(91, 29, 91, 43): 'image < -value', (91, 44, 91, 45): '0', (91, 46, 91, 57): 'image + value'}, {}), '(image < -value, 0, image + value)', True, 'import numpy as np\n')]