Topic: Python – Reading Images from Datasets and Organizing Them (Part 3): Custom Dataset Class and Data Augmentation

---

1. Creating a Custom Dataset Class (PyTorch)

Sometimes you need more control over how images and labels are loaded and processed. You can create a custom dataset class by extending torch.utils.data.Dataset.

import os
from PIL import Image
from torch.utils.data import Dataset

class CustomImageDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.image_paths = []
        self.labels = []
        self.class_to_idx = {}

        classes = sorted(os.listdir(root_dir))
        self.class_to_idx = {cls_name: idx for idx, cls_name in enumerate(classes)}

        for cls_name in classes:
            cls_dir = os.path.join(root_dir, cls_name)
            for img_name in os.listdir(cls_dir):
                img_path = os.path.join(cls_dir, img_name)
                self.image_paths.append(img_path)
                self.labels.append(self.class_to_idx[cls_name])

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

    def __getitem__(self, idx):
        img_path = self.image_paths[idx]
        image = Image.open(img_path).convert("RGB")
        label = self.labels[idx]

        if self.transform:
            image = self.transform(image)

        return image, label

---

2. Using Data Augmentation with `transforms`

Data augmentation helps improve model generalization by artificially increasing dataset diversity.

from torchvision import transforms

transform = transforms.Compose([
    transforms.Resize((128, 128)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(10),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

Pass this transform to the custom dataset:

dataset = CustomImageDataset(root_dir='dataset/', transform=transform)

---

3. Loading Dataset with DataLoader

from torch.utils.data import DataLoader

dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

---

4. Summary

• Custom dataset classes offer flexibility in how data is loaded and labeled.

• Data augmentation techniques such as flipping and rotation can be applied using torchvision transforms.

• Use DataLoader for batching and shuffling during training.

---

Exercise

• Extend the custom dataset to handle grayscale images and apply a random brightness adjustment transform.

---

#Python #DatasetHandling #PyTorch #DataAugmentation #ImageProcessing

https://t.me/DataScience4