import albumentations as albu
import numpy as np
import cv2
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'


class Dataset:
    def __init__(
            self, 
            image_path, 
            augmentation=None, 
            preprocessing=None,
    ):
        self.pil_image = image_path
        self.augmentation = augmentation
        self.preprocessing = preprocessing
    
    def get(self):
        # pil image > numpy array
        image = np.array(self.pil_image)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        # apply augmentations
        if self.augmentation:
            sample = self.augmentation(image=image)
            image = sample['image']
        
        # apply preprocessing
        if self.preprocessing:
            sample = self.preprocessing(image=image)
            image = sample['image']
            
        return image


def get_validation_augmentation():
    """Add paddings to make image shape divisible by 32"""
    test_transform = [
        albu.PadIfNeeded(384, 480)
    ]
    return albu.Compose(test_transform)


def to_tensor(x, **kwargs):
    return x.transpose(2, 0, 1).astype('float32')


def get_preprocessing(preprocessing_fn):
    
    _transform = [
        albu.Lambda(image=preprocessing_fn),
        albu.Lambda(image=to_tensor),
    ]
    return albu.Compose(_transform)