data_setup.py

import torch.utils.data
from torch.utils.data import DataLoader
from torchvision import datasets, transforms


DEFAULT_TRANSFORMER = transforms.Compose([
  transforms.Resize((64, 64)),
  transforms.ToTensor()
])


def create_dataset(images_path: str, test_train_split_enabled: bool=False, test_size: float=0.3, transformer = DEFAULT_TRANSFORMER):
    dataset = datasets.ImageFolder(images_path, transform=transformer)
    if test_train_split_enabled:
        return train_test_split(dataset, test_size)
    return dataset


def train_test_split(dataset, test_size: float):
    TOTAL = len(dataset)
    VAL = int(TOTAL * test_size)
    TRAIN = int(TOTAL - VAL)
    return torch.utils.data.random_split(dataset=dataset, lengths=[TRAIN, VAL])


def create_dataloader(dataset, batch_size, shuffle: bool=False) -> DataLoader:
    return DataLoader(dataset=dataset, batch_size=batch_size, shuffle=shuffle)