dataset.py

from pathlib import Path

import torch
from torch.utils.data import DataLoader, random_split, Subset
from torchvision import transforms
from torchvision.datasets import ImageFolder


def get_transforms(train: bool = True):
    # CIFAR-10 style preprocessing with 32x32 inputs
    mean = (0.4914, 0.4822, 0.4465)
    std = (0.2470, 0.2435, 0.2616)
    if train:
        return transforms.Compose(
            [
                transforms.Resize((32, 32)),
                transforms.RandomHorizontalFlip(),
                transforms.RandomCrop(32, padding=4),
                transforms.ToTensor(),
                transforms.Normalize(mean, std),
            ]
        )
    return transforms.Compose(
        [
            transforms.Resize((32, 32)),
            transforms.ToTensor(),
            transforms.Normalize(mean, std),
        ]
    )


def create_dataloaders(
    base_dir: str,
    data_dir: str = "Multi-class Weather Dataset",
    batch_size: int = 64,
    val_split: float = 0.2,
    num_workers: int = 0,
    seed: int = 42,
):
    root = str(Path(base_dir) / data_dir)

    train_dataset = ImageFolder(root=root, transform=get_transforms(train=True))
    dataset_size = len(train_dataset)
    val_size = int(dataset_size * val_split)
    train_size = dataset_size - val_size

    generator = torch.Generator().manual_seed(seed)
    train_idx, val_idx = random_split(
        range(dataset_size), [train_size, val_size], generator=generator
    )

    val_dataset = ImageFolder(root=root, transform=get_transforms(train=False))
    train_set = Subset(train_dataset, train_idx.indices)
    val_set = Subset(val_dataset, val_idx.indices)

    train_loader = DataLoader(
        train_set, batch_size=batch_size, shuffle=True, num_workers=num_workers
    )
    val_loader = DataLoader(
        val_set, batch_size=batch_size, shuffle=False, num_workers=num_workers
    )
    return train_loader, val_loader, train_dataset.class_to_idx