Dataset class (custom datasets) in PyTorch - ML Experiment: Train & Evaluate

Experiment - Dataset class (custom datasets)

Problem:You want to train a PyTorch model on your own data, but the data is not in a ready-made format. You need to create a custom Dataset class to load and preprocess your data.

Current Metrics:No model training yet because data loading is incomplete.

Issue:Without a proper Dataset class, the DataLoader cannot fetch data correctly, so training cannot start.

Your Task

Create a custom PyTorch Dataset class that loads images and labels from a folder and a CSV file, and returns them properly for training.

Use PyTorch's Dataset class as the base.

Implement __len__ and __getitem__ methods.

Load images using PIL and apply basic transforms.

Labels should be read from a CSV file.

Hint 1

Hint 2

Hint 3

Solution

PyTorch

import os
import pandas as pd
from PIL import Image
from torch.utils.data import Dataset
import torchvision.transforms as transforms

class CustomImageDataset(Dataset):
    def __init__(self, annotations_file, img_dir, transform=None):
        self.img_labels = pd.read_csv(annotations_file)
        self.img_dir = img_dir
        self.transform = transform

    def __len__(self):
        return len(self.img_labels)

    def __getitem__(self, idx):
        img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
        image = Image.open(img_path).convert('RGB')
        label = self.img_labels.iloc[idx, 1]
        if self.transform:
            image = self.transform(image)
        return image, label

# Example usage:
# transform = transforms.Compose([
#     transforms.Resize((128, 128)),
#     transforms.ToTensor(),
# ])
# dataset = CustomImageDataset('labels.csv', 'images/', transform=transform)
# print(len(dataset))
# img, label = dataset[0]
# print(img.shape, label)

Created a class CustomImageDataset inheriting from torch.utils.data.Dataset.

Implemented __init__ to load CSV labels and store image directory and transforms.

Implemented __len__ to return dataset size.

Implemented __getitem__ to load an image, apply transforms, and return image and label.

Results Interpretation

Before: No data loading possible, training cannot start.

After: Dataset class loads images and labels correctly, enabling training.

Creating a custom Dataset class is essential to load your own data in PyTorch. Implementing __len__ and __getitem__ methods correctly allows DataLoader to fetch data for training.

Bonus Experiment

Add data augmentation transforms (like random horizontal flip and random rotation) to the Dataset class to improve model robustness.

💡 Hint

Use torchvision.transforms.RandomHorizontalFlip and RandomRotation in the transform pipeline.