Pytorch Loading two Images from Dataloader

Question

I'm trying to make a GAN which takes a lo-res image, and tries to create a hi-res image from it. To do this, I need to user a Dataloader which has both the hi-res and low-res training images stored in it.

     data_transform = transforms.Compose([transforms.Resize(imageSize),
                                         transforms.Grayscale(num_output_channels=1),
                                         transforms.ToTensor()])

    dataset_hi = "./hi-res-train"
    dataset_lo = "./low-res-train"

    img_data_hi = dset.ImageFolder(root=dataset_hi,transform=data_transform)
    img_data_lo = dset.ImageFolder(root=dataset_lo,transform=data_transform)

    dataloader_hi = torch.utils.data.DataLoader(img_data_hi, batch_size = batchSize, shuffle = True, num_workers = 2) 
    dataloader_lo = torch.utils.data.DataLoader(img_data_lo, batch_size = batchSize, shuffle = True, num_workers = 2) 

I've tried using two separate data loaders (shown above) but when they are shuffled, I cant enumerate through them both because the hi-res and low-res images are not matched up. How can I make it so I can enumerate and shuffle both with pytorch?

Answer 1

Assuming you have similar names for hi & low resolution images (say img01_hi & img01_low), one option is to create a custom Dataloader that returns both images by overriding __getitem__ method.

As both images are returned in one call, you can make sure they match by appending _hi & _low to the filename.

You may need to create a "cue" text file containing list of all your image file names to make sure you are processing each image file only once.