import numpy as np
from albumentations import (Compose, HorizontalFlip, VerticalFlip, Rotate, RandomRotate90,
                            ShiftScaleRotate, ElasticTransform,
                            GridDistortion, RandomSizedCrop, RandomCrop, CenterCrop,
                            RandomBrightnessContrast, HueSaturationValue, IAASharpen,
                            RandomGamma, RandomBrightness, RandomBrightnessContrast,
                            GaussianBlur,CLAHE,
                            Cutout, CoarseDropout, GaussNoise, ChannelShuffle, ToGray, OpticalDistortion,
                            Normalize, OneOf, NoOp)
from albumentations.pytorch import ToTensorV2 as ToTensor
from get_config import get_config
config = get_config()

MEAN = np.array([0.485, 0.456, 0.406])
STD  = np.array([0.229, 0.224, 0.225])

def get_transforms_train():
    transform_train = Compose([
        #Basic
        RandomRotate90(p=1),
        HorizontalFlip(p=0.5),
        
        #Morphology
        ShiftScaleRotate(shift_limit=0, scale_limit=(-0.2,0.2), rotate_limit=(-30,30), 
                         interpolation=1, border_mode=0, value=(0,0,0), p=0.5),
        GaussNoise(var_limit=(0,50.0), mean=0, p=0.5),
        GaussianBlur(blur_limit=(3,7), p=0.5),
        
        #Color
        RandomBrightnessContrast(brightness_limit=0.35, contrast_limit=0.5, 
                                 brightness_by_max=True,p=0.5),
        HueSaturationValue(hue_shift_limit=30, sat_shift_limit=30, 
                           val_shift_limit=0, p=0.5),
        
        CoarseDropout(max_holes=2, 
                      max_height=config['input_resolution'][0]//4, max_width=config['input_resolution'][1]//4, 
                      min_holes=1,
                      min_height=config['input_resolution'][0]//16, min_width=config['input_resolution'][1]//16, 
                      fill_value=0, mask_fill_value=0, p=0.5),
        
        Normalize(mean=(MEAN[0], MEAN[1], MEAN[2]), 
                  std=(STD[0], STD[1], STD[2])),
        ToTensor(),
    ])
    return transform_train


def get_transforms_valid():
    transform_valid = Compose([
        Normalize(mean=(MEAN[0], MEAN[1], MEAN[2]), 
                  std=(STD[0], STD[1], STD[2])),
        ToTensor(),
    ] )
    return transform_valid


def denormalize(z, mean=MEAN.reshape(-1,1,1), std=STD.reshape(-1,1,1)):
    return std*z + mean

[((737, 749), 'get_config.get_config', 'get_config', ([], {}), '()\n', (747, 749), False, 'from get_config import get_config\n'), ((758, 789), 'numpy.array', 'np.array', (['[0.485, 0.456, 0.406]'], {}), '([0.485, 0.456, 0.406])\n', (766, 789), True, 'import numpy as np\n'), ((797, 828), 'numpy.array', 'np.array', (['[0.229, 0.224, 0.225]'], {}), '([0.229, 0.224, 0.225])\n', (805, 828), True, 'import numpy as np\n'), ((913, 932), 'albumentations.RandomRotate90', 'RandomRotate90', ([], {'p': '(1)'}), '(p=1)\n', (927, 932), False, 'from albumentations import Compose, HorizontalFlip, VerticalFlip, Rotate, RandomRotate90, ShiftScaleRotate, ElasticTransform, GridDistortion, RandomSizedCrop, RandomCrop, CenterCrop, RandomBrightnessContrast, HueSaturationValue, IAASharpen, RandomGamma, RandomBrightness, RandomBrightnessContrast, GaussianBlur, CLAHE, Cutout, CoarseDropout, GaussNoise, ChannelShuffle, ToGray, OpticalDistortion, Normalize, OneOf, NoOp\n'), ((942, 963), 'albumentations.HorizontalFlip', 'HorizontalFlip', ([], {'p': '(0.5)'}), '(p=0.5)\n', (956, 963), False, 'from albumentations import Compose, HorizontalFlip, VerticalFlip, Rotate, RandomRotate90, ShiftScaleRotate, ElasticTransform, GridDistortion, RandomSizedCrop, RandomCrop, CenterCrop, RandomBrightnessContrast, HueSaturationValue, IAASharpen, RandomGamma, RandomBrightness, RandomBrightnessContrast, GaussianBlur, CLAHE, Cutout, CoarseDropout, GaussNoise, ChannelShuffle, ToGray, OpticalDistortion, Normalize, OneOf, NoOp\n'), ((1002, 1142), 'albumentations.ShiftScaleRotate', 'ShiftScaleRotate', ([], {'shift_limit': '(0)', 'scale_limit': '(-0.2, 0.2)', 'rotate_limit': '(-30, 30)', 'interpolation': '(1)', 'border_mode': '(0)', 'value': '(0, 0, 0)', 'p': '(0.5)'}), '(shift_limit=0, scale_limit=(-0.2, 0.2), rotate_limit=(-30,\n    30), interpolation=1, border_mode=0, value=(0, 0, 0), p=0.5)\n', (1018, 1142), False, 'from albumentations import Compose, HorizontalFlip, VerticalFlip, Rotate, RandomRotate90, ShiftScaleRotate, ElasticTransform, GridDistortion, RandomSizedCrop, RandomCrop, CenterCrop, RandomBrightnessContrast, HueSaturationValue, IAASharpen, RandomGamma, RandomBrightness, RandomBrightnessContrast, GaussianBlur, CLAHE, Cutout, CoarseDropout, GaussNoise, ChannelShuffle, ToGray, OpticalDistortion, Normalize, OneOf, NoOp\n'), ((1170, 1216), 'albumentations.GaussNoise', 'GaussNoise', ([], {'var_limit': '(0, 50.0)', 'mean': '(0)', 'p': '(0.5)'}), '(var_limit=(0, 50.0), mean=0, p=0.5)\n', (1180, 1216), False, 'from albumentations import Compose, HorizontalFlip, VerticalFlip, Rotate, RandomRotate90, ShiftScaleRotate, ElasticTransform, GridDistortion, RandomSizedCrop, RandomCrop, CenterCrop, RandomBrightnessContrast, HueSaturationValue, IAASharpen, RandomGamma, RandomBrightness, RandomBrightnessContrast, GaussianBlur, CLAHE, Cutout, CoarseDropout, GaussNoise, ChannelShuffle, ToGray, OpticalDistortion, Normalize, OneOf, NoOp\n'), ((1225, 1263), 'albumentations.GaussianBlur', 'GaussianBlur', ([], {'blur_limit': '(3, 7)', 'p': '(0.5)'}), '(blur_limit=(3, 7), p=0.5)\n', (1237, 1263), False, 'from albumentations import Compose, HorizontalFlip, VerticalFlip, Rotate, RandomRotate90, ShiftScaleRotate, ElasticTransform, GridDistortion, RandomSizedCrop, RandomCrop, CenterCrop, RandomBrightnessContrast, HueSaturationValue, IAASharpen, RandomGamma, RandomBrightness, RandomBrightnessContrast, GaussianBlur, CLAHE, Cutout, CoarseDropout, GaussNoise, ChannelShuffle, ToGray, OpticalDistortion, Normalize, OneOf, NoOp\n'), ((1296, 1398), 'albumentations.RandomBrightnessContrast', 'RandomBrightnessContrast', ([], {'brightness_limit': '(0.35)', 'contrast_limit': '(0.5)', 'brightness_by_max': '(True)', 'p': '(0.5)'}), '(brightness_limit=0.35, contrast_limit=0.5,\n    brightness_by_max=True, p=0.5)\n', (1320, 1398), False, 'from albumentations import Compose, HorizontalFlip, VerticalFlip, Rotate, RandomRotate90, ShiftScaleRotate, ElasticTransform, GridDistortion, RandomSizedCrop, RandomCrop, CenterCrop, RandomBrightnessContrast, HueSaturationValue, IAASharpen, RandomGamma, RandomBrightness, RandomBrightnessContrast, GaussianBlur, CLAHE, Cutout, CoarseDropout, GaussNoise, ChannelShuffle, ToGray, OpticalDistortion, Normalize, OneOf, NoOp\n'), ((1437, 1526), 'albumentations.HueSaturationValue', 'HueSaturationValue', ([], {'hue_shift_limit': '(30)', 'sat_shift_limit': '(30)', 'val_shift_limit': '(0)', 'p': '(0.5)'}), '(hue_shift_limit=30, sat_shift_limit=30, val_shift_limit=\n    0, p=0.5)\n', (1455, 1526), False, 'from albumentations import Compose, HorizontalFlip, VerticalFlip, Rotate, RandomRotate90, ShiftScaleRotate, ElasticTransform, GridDistortion, RandomSizedCrop, RandomCrop, CenterCrop, RandomBrightnessContrast, HueSaturationValue, IAASharpen, RandomGamma, RandomBrightness, RandomBrightnessContrast, GaussianBlur, CLAHE, Cutout, CoarseDropout, GaussNoise, ChannelShuffle, ToGray, OpticalDistortion, Normalize, OneOf, NoOp\n'), ((1568, 1849), 'albumentations.CoarseDropout', 'CoarseDropout', ([], {'max_holes': '(2)', 'max_height': "(config['input_resolution'][0] // 4)", 'max_width': "(config['input_resolution'][1] // 4)", 'min_holes': '(1)', 'min_height': "(config['input_resolution'][0] // 16)", 'min_width': "(config['input_resolution'][1] // 16)", 'fill_value': '(0)', 'mask_fill_value': '(0)', 'p': '(0.5)'}), "(max_holes=2, max_height=config['input_resolution'][0] // 4,\n    max_width=config['input_resolution'][1] // 4, min_holes=1, min_height=\n    config['input_resolution'][0] // 16, min_width=config[\n    'input_resolution'][1] // 16, fill_value=0, mask_fill_value=0, p=0.5)\n", (1581, 1849), False, 'from albumentations import Compose, HorizontalFlip, VerticalFlip, Rotate, RandomRotate90, ShiftScaleRotate, ElasticTransform, GridDistortion, RandomSizedCrop, RandomCrop, CenterCrop, RandomBrightnessContrast, HueSaturationValue, IAASharpen, RandomGamma, RandomBrightness, RandomBrightnessContrast, GaussianBlur, CLAHE, Cutout, CoarseDropout, GaussNoise, ChannelShuffle, ToGray, OpticalDistortion, Normalize, OneOf, NoOp\n'), ((1937, 2010), 'albumentations.Normalize', 'Normalize', ([], {'mean': '(MEAN[0], MEAN[1], MEAN[2])', 'std': '(STD[0], STD[1], STD[2])'}), '(mean=(MEAN[0], MEAN[1], MEAN[2]), std=(STD[0], STD[1], STD[2]))\n', (1946, 2010), False, 'from albumentations import Compose, HorizontalFlip, VerticalFlip, Rotate, RandomRotate90, ShiftScaleRotate, ElasticTransform, GridDistortion, RandomSizedCrop, RandomCrop, CenterCrop, RandomBrightnessContrast, HueSaturationValue, IAASharpen, RandomGamma, RandomBrightness, RandomBrightnessContrast, GaussianBlur, CLAHE, Cutout, CoarseDropout, GaussNoise, ChannelShuffle, ToGray, OpticalDistortion, Normalize, OneOf, NoOp\n'), ((2039, 2049), 'albumentations.pytorch.ToTensorV2', 'ToTensor', ([], {}), '()\n', (2047, 2049), True, 'from albumentations.pytorch import ToTensorV2 as ToTensor\n'), ((2155, 2228), 'albumentations.Normalize', 'Normalize', ([], {'mean': '(MEAN[0], MEAN[1], MEAN[2])', 'std': '(STD[0], STD[1], STD[2])'}), '(mean=(MEAN[0], MEAN[1], MEAN[2]), std=(STD[0], STD[1], STD[2]))\n', (2164, 2228), False, 'from albumentations import Compose, HorizontalFlip, VerticalFlip, Rotate, RandomRotate90, ShiftScaleRotate, ElasticTransform, GridDistortion, RandomSizedCrop, RandomCrop, CenterCrop, RandomBrightnessContrast, HueSaturationValue, IAASharpen, RandomGamma, RandomBrightness, RandomBrightnessContrast, GaussianBlur, CLAHE, Cutout, CoarseDropout, GaussNoise, ChannelShuffle, ToGray, OpticalDistortion, Normalize, OneOf, NoOp\n'), ((2257, 2267), 'albumentations.pytorch.ToTensorV2', 'ToTensor', ([], {}), '()\n', (2265, 2267), True, 'from albumentations.pytorch import ToTensorV2 as ToTensor\n')]