import torch
from os import listdir, path
from PIL import Image
import torchvision


class DiscriminatorDataset(torch.utils.data.Dataset):

    def __init__(self):
        super(DiscriminatorDataset, self).__init__()
        currentDir = path.dirname(__file__)
        abstractDir = path.join(currentDir, 'image_data/abstract')
        realisticDir = path.join(currentDir, 'image_data/realistic')
        abstractFiles = [path.join(abstractDir, f) for f in listdir(
            abstractDir) if path.isfile(path.join(abstractDir, f))]
        realisticFiles = [path.join(realisticDir, f) for f in listdir(
            realisticDir) if path.isfile(path.join(realisticDir, f))]
        self.abstractFilesLen = len(abstractFiles)
        self.allFiles = abstractFiles + realisticFiles

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

    def __getitem__(self, index):
        filename = self.allFiles[index]
        pilImage = Image.open(filename).convert("RGB")
        return (torchvision.transforms.ToTensor()(pilImage), 1 if index < self.abstractFilesLen else 0)

[((11, 21, 11, 43), 'os.path.dirname', 'path.dirname', ({(11, 34, 11, 42): '__file__'}, {}), '(__file__)', False, 'from os import listdir, path\n'), ((12, 22, 12, 66), 'os.path.join', 'path.join', ({(12, 32, 12, 42): 'currentDir', (12, 44, 12, 65): '"""image_data/abstract"""'}, {}), "(currentDir, 'image_data/abstract')", False, 'from os import listdir, path\n'), ((13, 23, 13, 68), 'os.path.join', 'path.join', ({(13, 33, 13, 43): 'currentDir', (13, 45, 13, 67): '"""image_data/realistic"""'}, {}), "(currentDir, 'image_data/realistic')", False, 'from os import listdir, path\n'), ((14, 25, 14, 50), 'os.path.join', 'path.join', ({(14, 35, 14, 46): 'abstractDir', (14, 48, 14, 49): 'f'}, {}), '(abstractDir, f)', False, 'from os import listdir, path\n'), ((16, 26, 16, 52), 'os.path.join', 'path.join', ({(16, 36, 16, 48): 'realisticDir', (16, 50, 16, 51): 'f'}, {}), '(realisticDir, f)', False, 'from os import listdir, path\n'), ((14, 60, 15, 24), 'os.listdir', 'listdir', ({(15, 12, 15, 23): 'abstractDir'}, {}), '(abstractDir)', False, 'from os import listdir, path\n'), ((16, 62, 17, 25), 'os.listdir', 'listdir', ({(17, 12, 17, 24): 'realisticDir'}, {}), '(realisticDir)', False, 'from os import listdir, path\n'), ((26, 19, 26, 39), 'PIL.Image.open', 'Image.open', ({(26, 30, 26, 38): 'filename'}, {}), '(filename)', False, 'from PIL import Image\n'), ((27, 16, 27, 49), 'torchvision.transforms.ToTensor', 'torchvision.transforms.ToTensor', ({}, {}), '()', False, 'import torchvision\n'), ((15, 40, 15, 65), 'os.path.join', 'path.join', ({(15, 50, 15, 61): 'abstractDir', (15, 63, 15, 64): 'f'}, {}), '(abstractDir, f)', False, 'from os import listdir, path\n'), ((17, 41, 17, 67), 'os.path.join', 'path.join', ({(17, 51, 17, 63): 'realisticDir', (17, 65, 17, 66): 'f'}, {}), '(realisticDir, f)', False, 'from os import listdir, path\n')]