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)

[((238, 260), 'os.path.dirname', 'path.dirname', (['__file__'], {}), '(__file__)\n', (250, 260), False, 'from os import listdir, path\n'), ((283, 327), 'os.path.join', 'path.join', (['currentDir', '"""image_data/abstract"""'], {}), "(currentDir, 'image_data/abstract')\n", (292, 327), False, 'from os import listdir, path\n'), ((351, 396), 'os.path.join', 'path.join', (['currentDir', '"""image_data/realistic"""'], {}), "(currentDir, 'image_data/realistic')\n", (360, 396), False, 'from os import listdir, path\n'), ((422, 447), 'os.path.join', 'path.join', (['abstractDir', 'f'], {}), '(abstractDir, f)\n', (431, 447), False, 'from os import listdir, path\n'), ((560, 586), 'os.path.join', 'path.join', (['realisticDir', 'f'], {}), '(realisticDir, f)\n', (569, 586), False, 'from os import listdir, path\n'), ((457, 477), 'os.listdir', 'listdir', (['abstractDir'], {}), '(abstractDir)\n', (464, 477), False, 'from os import listdir, path\n'), ((596, 617), 'os.listdir', 'listdir', (['realisticDir'], {}), '(realisticDir)\n', (603, 617), False, 'from os import listdir, path\n'), ((933, 953), 'PIL.Image.open', 'Image.open', (['filename'], {}), '(filename)\n', (943, 953), False, 'from PIL import Image\n'), ((985, 1018), 'torchvision.transforms.ToTensor', 'torchvision.transforms.ToTensor', ([], {}), '()\n', (1016, 1018), False, 'import torchvision\n'), ((506, 531), 'os.path.join', 'path.join', (['abstractDir', 'f'], {}), '(abstractDir, f)\n', (515, 531), False, 'from os import listdir, path\n'), ((646, 672), 'os.path.join', 'path.join', (['realisticDir', 'f'], {}), '(realisticDir, f)\n', (655, 672), False, 'from os import listdir, path\n')]