[Beginner] How do I write a custom dataset that allows me to return image, and its target image(not label) as a pair?

vision
#1

Right now, most of the examples just give out image, lables based upon the folder sturcture. How do i write a dataloader/dataset, such that i can use it for a network that take input image and outputs image?

#2

You could create your own Dataset like this:

class MyDataset(Dataset):
    def __init__(self, data_paths):
        self.paths = data_paths 

    def __getitem(self, index):
        image, target_image = load_images(self.paths[index]) # load as np.array or PIL.Image
        # transform your images here
        x = transformations.ToTensor()(image)
        y = transformations.ToTensor()(target_image)

        return x, y

    def __len__(self):
        return len(self.paths)
#3

Thanks @ptrblck, I was able to implement a custom loader with the same idea on the lines of ImageFolder!

#4

Very nice!
I was writing from my phone and just realized, that my code snippet has some typos (e.g. transformations should be transforms), sorry for that and apparently you fixed these issues anyway. :wink:

#5

Yea, thats the benefit of using an advanced IDE like PyCharm :sweat_smile:

1 Like