Alternatively you could specify your own ImageFolder which is also completly straight forward.
An example is given below and it should work quite simple if you shuffle imgs
in the __init__
. This way you can also do some fancy preprocessing on numpy etc by specifying your own load-funktion and pass it to loader
class ImageFolder(data.Dataset):
"""Class for handling image load process and transformations"""
def __init__(self, image_path, options, transform=None, return_paths=True,
loader=default_loader_unaligned):
"""
Function to create the dataset and initialize the class variables
:param image_path: path containing image-files
:param options: class containing all options (args of BaseOptions or subclass)
:param transform: transformation to apply on the Image after loading it
:param return_paths: Boolean, True if paths should be returned alongside images , False if only images
:param loader: function to load and resize images
"""
imgs = make_dataset(image_path)
if len(imgs) == 0:
raise(RuntimeError("Found 0 images in: " + image_path + "\n"
"Supported image extensions are: " + ",".join(IMG_EXTENSIONS)))
self.root = image_path
self.imgs = imgs
self.transform = transform
self.return_paths = return_paths
self.loader = loader
self.options = options
def __getitem__(self, index):
"""
Function to get certain item in dataset
:param index: index of dataset-list
:return: item in dataset with given index
"""
path = self.imgs[index]
img = self.loader(path, self.options.imageSize, self.options.inputNc)
if self.transform is not None:
img = self.transform(img)
if self.return_paths:
return img, path
else:
return img
def __len__(self):
"""Function to get number of items in dataset"""
return len(self.imgs)