# 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)

[((538, 589), 'cv2.resize', 'cv2.resize', (['image', '(width, heigh)', 'cv2.INTER_LINEAR'], {}), '(image, (width, heigh), cv2.INTER_LINEAR)\n', (548, 589), False, 'import cv2\n'), ((379, 401), 'cv2.imread', 'cv2.imread', (['image_path'], {}), '(image_path)\n', (389, 401), False, 'import cv2\n'), ((1257, 1275), 'numpy.random.random', 'np.random.random', ([], {}), '()\n', (1273, 1275), True, 'import numpy as np\n'), ((1297, 1313), 'numpy.fliplr', 'np.fliplr', (['image'], {}), '(image)\n', (1306, 1313), True, 'import numpy as np\n'), ((1329, 1344), 'numpy.fliplr', 'np.fliplr', (['mask'], {}), '(mask)\n', (1338, 1344), True, 'import numpy as np\n'), ((1475, 1493), 'numpy.random.random', 'np.random.random', ([], {}), '()\n', (1491, 1493), True, 'import numpy as np\n'), ((1515, 1531), 'numpy.flipud', 'np.flipud', (['image'], {}), '(image)\n', (1524, 1531), True, 'import numpy as np\n'), ((1547, 1562), 'numpy.flipud', 'np.flipud', (['mask'], {}), '(mask)\n', (1556, 1562), True, 'import numpy as np\n'), ((1898, 1948), 'numpy.arange', 'np.arange', (['(-max_shift_pixels)', '(max_shift_pixels + 1)'], {}), '(-max_shift_pixels, max_shift_pixels + 1)\n', (1907, 1948), True, 'import numpy as np\n'), ((1955, 1973), 'numpy.random.random', 'np.random.random', ([], {}), '()\n', (1971, 1973), True, 'import numpy as np\n'), ((1995, 2024), 'numpy.roll', 'np.roll', (['image', 'shift'], {'axis': '(1)'}), '(image, shift, axis=1)\n', (2002, 2024), True, 'import numpy as np\n'), ((2040, 2068), 'numpy.roll', 'np.roll', (['mask', 'shift'], {'axis': '(1)'}), '(mask, shift, axis=1)\n', (2047, 2068), True, 'import numpy as np\n'), ((2395, 2445), 'numpy.arange', 'np.arange', (['(-max_shift_pixels)', '(max_shift_pixels + 1)'], {}), '(-max_shift_pixels, max_shift_pixels + 1)\n', (2404, 2445), True, 'import numpy as np\n'), ((2452, 2470), 'numpy.random.random', 'np.random.random', ([], {}), '()\n', (2468, 2470), True, 'import numpy as np\n'), ((2496, 2525), 'numpy.roll', 'np.roll', (['image', 'shift'], {'axis': '(0)'}), '(image, shift, axis=0)\n', (2503, 2525), True, 'import numpy as np\n'), ((2545, 2573), 'numpy.roll', 'np.roll', (['mask', 'shift'], {'axis': '(0)'}), '(mask, shift, axis=0)\n', (2552, 2573), True, 'import numpy as np\n'), ((2694, 2712), 'numpy.random.random', 'np.random.random', ([], {}), '()\n', (2710, 2712), True, 'import numpy as np\n'), ((3039, 3057), 'numpy.random.random', 'np.random.random', ([], {}), '()\n', (3055, 3057), True, 'import numpy as np\n'), ((3083, 3115), 'numpy.percentile', 'np.percentile', (['(255 - image)'], {'q': '(25)'}), '(255 - image, q=25)\n', (3096, 3115), True, 'import numpy as np\n'), ((3696, 3714), 'numpy.random.random', 'np.random.random', ([], {}), '()\n', (3712, 3714), True, 'import numpy as np\n'), ((3801, 3839), 'cv2.cvtColor', 'cv2.cvtColor', (['image', 'cv2.COLOR_BGR2HSV'], {}), '(image, cv2.COLOR_BGR2HSV)\n', (3813, 3839), False, 'import cv2\n'), ((3960, 3996), 'cv2.cvtColor', 'cv2.cvtColor', (['hsv', 'cv2.COLOR_HSV2BGR'], {}), '(hsv, cv2.COLOR_HSV2BGR)\n', (3972, 3996), False, 'import cv2\n'), ((437, 459), 'PIL.Image.open', 'Image.open', (['image_path'], {}), '(image_path)\n', (447, 459), False, 'from PIL import Image\n'), ((3210, 3242), 'numpy.arange', 'np.arange', (['(-max_value)', 'max_value'], {}), '(-max_value, max_value)\n', (3219, 3242), True, 'import numpy as np\n'), ((3753, 3785), 'numpy.arange', 'np.arange', (['(-max_value)', 'max_value'], {}), '(-max_value, max_value)\n', (3762, 3785), True, 'import numpy as np\n'), ((3894, 3926), 'numpy.clip', 'np.clip', (['hsv'], {'a_min': '(0)', 'a_max': '(255)'}), '(hsv, a_min=0, a_max=255)\n', (3901, 3926), True, 'import numpy as np\n'), ((3286, 3335), 'numpy.where', 'np.where', (['(255 - image < value)', '(255)', '(image + value)'], {}), '(255 - image < value, 255, image + value)\n', (3294, 3335), True, 'import numpy as np\n'), ((3385, 3427), 'numpy.where', 'np.where', (['(image < -value)', '(0)', '(image + value)'], {}), '(image < -value, 0, image + value)\n', (3393, 3427), True, 'import numpy as np\n')]