Loading original images besides the transformed ones using ImageFolder

Question

I\'m trying to train a GAN to colorize images. For that, I\'m using ImageFolder of torchvision to load grayscale images but I also need the original da

Answer 1

I assume you have 2 dataset variables i.e. dataset_bw and dataset_color that you can load as you mention using ImageFolder. Then you could do the following :

class GAN_dataset(Dataset):
    def __init__(self, dataset_bw, dataset_color):
        self.dataset1 = dataset_bw
        self.dataset2 = dataset_color

    def __getitem__(self, index):
        x1 = self.dataset1[index]
        x2 = self.dataset2[index]

        return x1, x2

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

dataset = GAN_dataset(dataset_bw, dataset_color)
loader = DataLoader(dataset, batch_size = ...)

This way you when you iterate through loader, you will get two images as you require.