from torchvision import transforms as T

class CLAHE:
    def __init__(self, clipLimit=2.0, tileGridSize=(8, 8)):
        self.clipLimit = clipLimit
        self.tileGridSize = tileGridSize

    def __call__(self, im):
        im = np.array(im) 
        img_yuv = cv2.cvtColor(im, cv2.COLOR_RGB2YUV)
        clahe = cv2.createCLAHE(clipLimit=self.clipLimit, tileGridSize=self.tileGridSize)
        img_yuv[:, :, 0] = clahe.apply(img_yuv[:, :, 0])
        img_output = cv2.cvtColor(img_yuv, cv2.COLOR_YUV2BGR)
        return img_output

from torchvision import transforms as T

size = 50 #32
im_size = (size, size)

def base_transform():
    transforms = []
    # converts the image, a PIL image, into a PyTorch Tensor
    transforms.append(T.Resize(im_size))
    transforms.append(T.ToTensor())
    return T.Compose(transforms)

def jitter_brightness():
    transforms = [
        T.Resize(im_size),
        T.ColorJitter(brightness=5),
        T.ToTensor()
    ]
    return T.Compose(transforms)

def jitter_saturation():
    transforms = [
        T.Resize(im_size),
        T.ColorJitter(saturation=5),
        T.ToTensor()
    ]
    return T.Compose(transforms)

def jitter_contrast():
    transforms = [
        T.Resize(im_size),
        T.ColorJitter(contrast=5),
        T.ToTensor()
    ]
    return T.Compose(transforms)

def jitter_hue():
    transforms = [
        T.Resize(im_size),
        T.ColorJitter(hue=0.4),
        T.ToTensor()
    ]
    return T.Compose(transforms)

def jitter_rot():
    transforms = [
        T.Resize(im_size),
        T.RandomRotation(15),
        T.ToTensor()
    ]
    return T.Compose(transforms)

def shear():
    transforms = [
        T.Resize(im_size),
        T.RandomAffine(degrees=15, shear=2),
        T.ToTensor()
    ]
    return T.Compose(transforms)

def clahe():
    transforms = [
        T.Resize(im_size),
        CLAHE(),
        T.ToTensor()
    ]
    return T.Compose(transforms)