from os import listdir
from os.path import join
import matplotlib.pyplot as plt
from PIL.Image import open
from torch.utils.data import Dataset
import data.augmentation as aug

from PIL import Image


class CVC_ClinicDB(Dataset):
    def __init__(self, root_directory):
        super(CVC_ClinicDB, self).__init__()
        self.root = root_directory
        self.mask_fnames = listdir(join(self.root, "Ground Truth"))
        self.mask_locs = [join(self.root, "Ground Truth", i) for i in self.mask_fnames]
        self.img_locs = [join(self.root, "Original", i) for i in
                         self.mask_fnames]
        self.common_transforms = aug.pipeline_tranforms()

    def __getitem__(self, idx):
        mask = self.common_transforms(open(self.mask_locs[idx]))
        image = self.common_transforms(open(self.img_locs[idx]))
        return image, mask, self.mask_fnames[idx]

    def __len__(self):
        return len(self.mask_fnames)


if __name__ == '__main__':
    dataset = CVC_ClinicDB("Datasets/CVC-ClinicDB")
    for img, mask, fname in dataset:
        plt.imshow(img.T)
        plt.imshow(mask.T, alpha=0.5)
        plt.show()
        input()
    print("done")