I would suggest to write a custom Dataset
and implement the logic to load the corresponding 10 label images in the __getitem__(self, index)
method.
To enable lazy loading, you should pass the file paths in __init__
and only load the current sample images in __getitem__
.
Here is some pseudo code:
class MyDataset(Dataset):
def __init__(self, data_paths, label_paths, transform=None, target_transform=None):
self.data_paths = data_paths # Could be a list: ['./train/input/image_1.bmp', './train/input/image_2.bmp', ...]
self.label_paths = label_paths # Could be a nested list: [['./train/GT/image_1_1.bmp', './train/GT/image_1_2.bmp', ...], ['./train/GT/image_2_1.bmp', './train/GT/image_2_2.bmp', ...]]
self.transform = transforms
self.target_transform = target_transform
def __getitem__(self, index):
x = Image.open(self.data_paths[index])
if self.transform:
x = self.transform(x)
ys = []
for label_path in self.label_paths[index]:
y = Image.open(label_path)
if self.target_transform:
y = self.target_transform(y)
ys.append(y)
return x, ys
def __len__(self):
return len(self.data_paths)
Let me know, if you get stuck somewhere.