How to Load A Custom Dataset On PyTorch | Towards Data Science

When you want to build a machine learning model, the first thing that you have to do is to prepare the dataset. When your data is on tabular format, it’s easy to prepare them. But what about data like images?
Images don’t have the same format with tabular data. There are so many data representations for this format.
Some people put the images to a folder based on its corresponding class, and some people make the metadata on tabular format that describes the image file name and its labels.
When the dataset on the first format, we can load the dataset easier by using a class called ImageFolder from torch.data.utils library.
But most of the time, the image datasets have the second format, where it consists of the metadata and the image folder. Therefore, we have to give some effort for preparing the dataset.
For example, you want to build an image classifier using deep learning, and it consists of a metadata that looks like this,
As you can see here, the dataset consists of image ids and labels. In this case, the image ids also represent the filename on .jpg format, and the labels are on one-hot encoded format.
How can we load the dataset so the model can read the images and their labels? What you can do is to build an object that can contain them.
In this article, I will show you on how to load image dataset that contains metadata using PyTorch. For the dataset, we will use a dataset from Kaggle competition called Plant Pathology 2020 — FGVC7, which you can access the data here.
The first thing that we have to do is to preprocess the metadata. As we can see from the image above, the dataset does not consists the image file name. Also, the label still on one-hot format.
But thankfully, the image ids also represent the image file name by adding .jpg to the ids. The code to generate image file names looks like this,
import pandas as pd# ASSUME THAT YOU RUN THE CODE ON KAGGLE NOTEBOOKpath = '/kaggle/input/plant-pathology-2020-fgvc7/'img_path = path + 'images'# LOAD THE DATASETtrain_df = pd.read_csv(path + 'train.csv')test_df = pd.read_csv(path + 'test.csv')sample = pd.read_csv(path + 'sample_submission.csv')# GET THE IMAGE FILE NAMEtrain_df['img_path'] = train_df['image_id'] + '.jpg'test_df['img_path'] = test_df['image_id'] + '.jpg'train_df.head()
The result looks like this,
Right after we get the image file names, now we can unpivot the labels to become a single column. The code looks like this,
# UNPIVOT / MELT THE LABELStrain_label = train_df.melt(id_vars=['image_id', 'img_path'])# FILTER THE DATAtrain_label = train_label[train_label['value'] == 1]# GET THE IMAGE ID NUMBERtrain_label['id'] = [int(i[1]) for i in train_label['image_id'].str.split('_')]# RESET THE INDEXtrain_label = train_label.sort_values('id').reset_index()# ADD THE LABEL TO THE DATASETtrain_df['label'] = train_label['variable']# REFORMAT THE DATASETtrain_df = train_df[train_df.columns[[0, 5, 1, 2, 3, 4, 6]]]print(train_label.shape)train_df.head()
The result looks like this,
Because the machine learning model can only read numbers, we have to encode the label to numbers. The code looks like this,
from sklearn.preprocessing import LabelEncoder# Encode the labelle = LabelEncoder()label_encoded = le.fit_transform(train_df['label'])train_df['label_encoded'] = label_encoded# Taking the class namelabel_names = label_encoded.classes_train_df.head()
The result looks like this,
Right after we preprocess the metadata, now we can move to the next step.
The next step is to build a container object for our images and labels. The reason why we need to build that object is to make our task for loading the data to the deep learning model much easier. Therefore, we can access the image and its label by using an index.
To create the object, we can use a class called Dataset from torch.utils.data library. This class is an abstract class because it consists of functions or methods that are not yet being implemented. Therefore, we can implement those functions by our own that suits to our needs.
The functions that we need to implement are
__init__ function,
__len__ function, and
__getitem__ function.
The __init__ function will initialize an object from its class and collect parameters from the user. The __len__function will return the length of the dataset. Lastly, the __getitem__ function, which is the most important one, will help us to return data observation by using an index.
By understanding the class and its corresponding functions, now we can implement the code. In this case, I will use the class name called PathologyPlantsDataset that will inherit functions from Dataset class. The code looks like this,
class PathologyPlantsDataset(Dataset):  """  The Class will act as the container for our dataset. It will take your dataframe, the root path, and also the transform function for transforming the dataset.  """    def __init__(self, data_frame, root_dir, transform=None):        self.data_frame = data_frame        self.root_dir = root_dir        self.transform = transform        def __len__(self):        # Return the length of the dataset        return len(self.data_frame)        def __getitem__(self, idx):        # Return the observation based on an index. Ex. dataset[0] will return the first element from the dataset, in this case the image and the label.        if torch.is_tensor(idx):            idx = idx.tolist()                img_name = os.path.join(self.root_dir, self.data_frame.iloc[idx, 1])        image = Image.open(img_name)        label = self.data_frame.iloc[idx, -1]                if self.transform:            image = self.transform(image)            return (image, label)
After we create the class, now we can build the object from it. When we create the object, we will set parameters that consist of the dataset, the root directory, and the transform function. The code looks like this,
# INSTANTIATE THE OBJECTpathology_train = PathologyPlantsDataset(    data_frame=train_part,    root_dir=path + 'images',    transform=transforms.Compose([        transforms.RandomResizedCrop(224),        transforms.RandomHorizontalFlip(),        transforms.ToTensor(),        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])    ]))
Now, we can extract the image and its label by using the object. As I’ve mentioned above, for accessing the observation from the data, we can use an index.
For example, when we want to access the third row of the dataset, which the index is 2, we can access it by using pathology_train[2].
Let me show you the example on how to visualize the result using pathology_train variable. The code looks like this,
temp_img, temp_lab = pathology_train[2]plt.imshow(temp_img.numpy().transpose((1, 2, 0)))plt.title(label_names[temp_lab])plt.axis('off')plt.show()
The result looks like this,
Well done! Now we have implemented the object that can load the dataset for our deep learning model much easier. I hope you can try it with your dataset.
If you want to discuss more, you can connect with me on LinkedIn and have a discussion on it.
Also, you can follow my Medium to read more of my articles, thank you!
[1] https://pytorch.org/tutorials/beginner/data_loading_tutorial.html[2] https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html