data_setup.py

import torch, os, pandas as pd

from torch.utils.data import Dataset, DataLoader, random_split
from torchvision import transforms
from PIL import Image


NUM_WORKERS = os.cpu_count()

DEFAULT_TRANSFORMER = transforms.Compose([
    transforms.Resize(size=(64, 64)),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
class CustomDataset(Dataset):
    def __init__(self, csv_dataframe_with_label: pd.DataFrame, images_dir: str = "./data/", transformer: transforms.Compose = DEFAULT_TRANSFORMER):
        self.csv_dataframe_with_label = csv_dataframe_with_label
        self.images_dir = images_dir
        self.transformer = transformer

    def __len__(self):
        return self.csv_dataframe_with_label.shape[0]

    def load_image(self, path):
        image = Image.open(self.images_dir + path)
        if self.transformer is not None:
            return self.transformer(image)
        return image

    def __getitem__(self, item):
        file_name, label = self.csv_dataframe_with_label.iloc[item]
        return self.load_image(file_name), label

def train_test_split(dataset: Dataset, test_size: float):
    TOTAL_LEN = len(dataset)
    VAL_LEN = int(TOTAL_LEN * test_size)
    TRAIN_LEN = int(TOTAL_LEN - VAL_LEN)
    return random_split(dataset, lengths=[TRAIN_LEN, VAL_LEN])

def filter_condition(inputs):
    x, y = inputs
    return x.shape[0] == 3
def my_collate_fn(batch):
    batch = list(filter(filter_condition, batch))
    return torch.utils.data.dataloader.default_collate(batch)
def data_load(dataset: Dataset, batch_size, num_workers):
    return DataLoader(dataset, batch_size, collate_fn=my_collate_fn)
def create_dataloaders(dataset: Dataset, batch_size: int=32, test_size: float=0.1 , split_test: bool=False, num_workers: int=NUM_WORKERS):
    if split_test:
        train_dataset, test_dataset = train_test_split(dataset, test_size)
        return data_load(train_dataset, batch_size, num_workers), data_load(test_dataset, batch_size, num_workers)
    return data_load(dataset, batch_size, num_workers)