datasets.py

from torchvision import datasets
from torchvision import transforms
from torch.utils.data import DataLoader
import torch
import numpy as np
import os.path as osp


def setup_data_loader(dataset, args):
    if dataset == 'ImageNet-C' or dataset == 'IN-C':
        return setup_imagenetc_data_loader(args)
    elif dataset == 'ImageNet' or dataset == 'IN':
        return setup_imagenet_data_loader(args)
    elif dataset == 'SIN':
        return setup_sin_data_loader(args)
    else:
        raise ValueError(f'Dataset {dataset} is not available')


def setup_imagenet_data_loader(args):
    n_worker = args.workers
    valdir = osp.join(args.datadir_clean, 'val')
    traindir = osp.join(args.datadir_clean, 'train')
    
    if args.arch == 'inceptionV3':
        train_transforms = transforms.Compose([
            transforms.RandomResizedCrop(299),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            ])
        val_transforms = transforms.Compose([
            transforms.Resize(299),
            transforms.CenterCrop(299),
            transforms.ToTensor()
            ])
    else:
        # transforms for a resnet
        train_transforms = transforms.Compose([
            transforms.RandomResizedCrop(224),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            ])
        val_transforms = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor()
            ])
    
    train_dataset = datasets.ImageFolder(traindir, train_transforms)
    train_loader = torch.utils.data.DataLoader(
        train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=n_worker)
    val_dataset = datasets.ImageFolder(valdir, val_transforms)    
    indices_subsample = np.random.choice(len(val_dataset), args.test_subset_size, replace=False)
    test_loader_subsample = torch.utils.data.DataLoader(
        torch.utils.data.Subset(dataset=val_dataset, indices=indices_subsample),
        batch_size=args.test_batch_size, shuffle=False, num_workers=n_worker)
    test_loader = torch.utils.data.DataLoader(val_dataset,
                                              batch_size=args.test_batch_size, shuffle=False, num_workers=n_worker)
    train_loader_retrain = DataLoader(
        train_dataset, batch_size=args.batch_size, shuffle=True,
        num_workers=n_worker)

    return train_loader, test_loader, test_loader_subsample, train_loader_retrain


def setup_imagenetc_data_loader(args):
    n_worker = args.workers
    data_loaders_names = {
        'Brightness': 'brightness',
        'Contrast': 'contrast',
        'Defocus Blur': 'defocus_blur',
        'Elastic Transform': 'elastic_transform',
        'Fog': 'fog',
        'Frost': 'frost',
        'Gaussian Noise': 'gaussian_noise',
        'Glass Blur': 'glass_blur',
        'Impulse Noise': 'impulse_noise',
        'JPEG Compression': 'jpeg_compression',
        'Motion Blur': 'motion_blur',
        'Pixelate': 'pixelate',
        'Shot Noise': 'shot_noise',
        'Snow': 'snow',
        'Zoom Blur': 'zoom_blur'
    }
    
    data_loaders = {}    
    for name, path in data_loaders_names.items():
        data_loaders[name] = {}
        for severity in range(1, 6):
            dset = datasets.ImageFolder(args.imagenetc_path + path +
                                        '/' + str(severity) + '/',
                                        transforms.Compose([transforms.CenterCrop(224), transforms.ToTensor()]))
            data_loaders[name][str(severity)] = torch.utils.data.DataLoader(
                dset, batch_size=args.test_batch_size, shuffle=True, num_workers=n_worker)
    return data_loaders   
    

def setup_sin_data_loader(args):
    n_worker = args.workers
    traindir = osp.join(args.datadir_sin, 'train')
    train_dataset = datasets.ImageFolder(traindir, transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
    ]))
    train_loader = torch.utils.data.DataLoader(
        train_dataset, batch_size=args.sin_batch_size, shuffle=True, num_workers=n_worker)
    
    return train_loader