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, DataLoader

from torchvision import transforms

import data.augmentation as aug

class EtisDataset(Dataset):

    """

        Dataset class that fetches Etis-LaribPolypDB images with the associated segmentation mask.

        Used for testing.

    """

    def __init__(self, path):

        super(EtisDataset, self).__init__()

        self.path = path

        self.len = len(listdir(join(self.path, "ETIS-LaribPolypDB")))

        self.common_transforms = aug.pipeline_tranforms()

    def __len__(self):

        return self.len

    def __getitem__(self, index):

        image = self.common_transforms(

            open(join(self.path, "ETIS-LaribPolypDB/{}.jpg".format(index + 1))).convert("RGB"))

        mask = self.common_transforms(

            open(join(self.path, "GroundTruth/p{}.jpg".format(index + 1))).convert("RGB"))

        mask = (mask > 0.5).float()

        return image, mask, index + 1

def test_etis():

    for x, y, in DataLoader(EtisDataset("Datasets/ETIS-LaribPolypDB")):

        plt.imshow(x)

        plt.show()