load_data.py

from torchvision import datasets, transforms
from torch.utils.data import DataLoader, random_split
from torchsampler import ImbalancedDatasetSampler


def load_data(
    train_dir="train_images",
    test_dir="test_images",
    batch_size=32,
    val_ratio=0.15,
    use_sampler=True,
):
    """Load and preprocess the training, validation, and test datasets.

    Args:
        train_dir (str, optional): The directory containing the training images. Defaults to 'train_images'.
        test_dir (str, optional): The directory containing the test images. Defaults to 'test_images'.
        batch_size (int, optional): The batch size for the data loaders. Defaults to 32.
        val_ratio (float, optional): The ratio of validation samples to the total training samples. Defaults to 0.15.
        use_sampler (bool, optional): Whether to use ImbalancedDatasetSampler for training. Defaults to True.

    Returns:
        tuple: A tuple containing the train_loader, val_loader, and test_loader.
    """
    # process images and load data
    transform = transforms.Compose(
        [
            transforms.Grayscale(),
            transforms.Resize((36, 36)),
            transforms.ToTensor(),
            transforms.Normalize(mean=(0.5,), std=(0.5,)),  # normalize to (-1,1)
        ]
    )

    # load train dataset and split train data into train and validation sets
    full_train_data = datasets.ImageFolder(train_dir, transform=transform)

    val_len = int(len(full_train_data) * val_ratio)
    train_len = len(full_train_data) - val_len
    train_set, val_set = random_split(full_train_data, [train_len, val_len])

    print(f"Train set size: {len(train_set)}, Validation set size: {len(val_set)}")

    # load test dataset
    test_dataset = datasets.ImageFolder(test_dir, transform=transform)

    # create data loaders
    train_loader = DataLoader(
        train_set,
        batch_size=batch_size,
        sampler=ImbalancedDatasetSampler(train_set) if use_sampler else None,
        shuffle=False,
    )
    val_loader = DataLoader(
        val_set,
        batch_size=batch_size,
        shuffle=False,
    )
    test_loader = DataLoader(
        test_dataset,
        batch_size=batch_size,
        shuffle=False,
    )

    return train_loader, val_loader, test_loader