Custom Dataset to take Images from Various Folders (not based on class)

Hi All,

I’m wondering how I should augment the Dataset class so that I can make a dataset with images from different folders? These folders are not based on classes(which is why I’m not using ImageFolder) - I just needed to store the images in different folders since the amount of files was too large to store in one folder and was causing input/output errors.

I want to specify a few folder locations, and then make a dataset object that contains all images from those folders in one. I am getting the class labels separately by indexing into a csv file.

Any way on doing this?

Thank you!


I think the snippet below will do. Note that I’ve not tested this.

class Dataset(ImageFolder):

    def __init__(self, labels, *args_for_ImageFolder, **kwargs_for_ImageFolder):
        super().__init__(*args_for_ImageFolder, **kwargs_for_ImageFolder)
        self.labels = labels

    def __getitem__(self, index):
        return (super().__getitem__(index)[0], self.labels[index])