I am writing a simple custom DataLoader (which I will add more features to later) for a segmentation dataset but the (image, mask) pair I return using __getitem()__ method are different; the returned mask belongs to a different image than the one which is returned. My directory structure is /home/bohare/data/images and /home/bohare/data/masks .
Following is the code I have:
import torch
from torch.utils.data.dataset import Dataset
from PIL import Image
import glob
import os
import matplotlib.pyplot as plt
class CustomDataset(Dataset):
def __init__(self, folder_path):
self.img_files = glob.glob(os.path.join(folder_path,'images','*.png'))
self.mask_files = glob.glob(os.path.join(folder_path,'masks','*.png'))
def __getitem__(self, index):
image = Image.open(self.img_files[index])
mask = Image.open(self.mask_files[index])
return image, mask
def __len__(self):
return len(self.img_files)
data = CustomDataset(folder_path = '/home/bohare/data')
len(data)
This code correctly gives out the total size of the dataset.
But when I use:
img, msk = data.__getitem__(n) where n is the index of any (image, mask) pair and I plot the image and mask, they do not correspond to one another.
How can I modify/what can I add to the code to make sure the (image, mask) pair are returned correctly? Thanks for the help.