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


class EndoCV2020(Dataset):
    def __init__(self, root_directory):
        super(EndoCV2020, self).__init__()
        self.root = root_directory
        self.mask_fnames = listdir(join(self.root, "masksPerClass", "polyp"))
        self.mask_locs = [join(self.root, "masksPerClass", "polyp", i) for i in self.mask_fnames]
        self.img_locs = [join(self.root, "originalImages", i.replace("_polyp", "").replace(".tif", ".jpg")) 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 = EndoCV2020("Datasets/EndoCV2020")
    for img, mask, fname in dataset:
        plt.imshow(img.T)
        plt.imshow(mask.T, alpha=0.5)
        plt.show()
        input()
    print("done")