data.py

import os
import warnings

import torch
import torchvision as tv
import torchvision.datasets as datasets
import numpy as np
from PIL import Image
from utils.dataset import RandomDatasetGenerator
from utils.misc import _META
from torchvision.datasets.utils import download_url,check_integrity

warnings.filterwarnings("ignore", "(Possibly )?corrupt EXIF data", UserWarning)

class CIFAR10_1(datasets.VisionDataset):
    source_url = 'https://github.com/modestyachts/CIFAR-10.1/raw/master/datasets/'
    files = [['cifar10.1_v6_data.npy','4fcae82cb1326aec9ed1dc1fc62345b8'],
             ['cifar10.1_v6_labels.npy','09a97fb7c430502fcbd69d95093a3f85']]

    def __init__(self, root, train=True, transform=None, target_transform=None,
                 download=False):
        super(CIFAR10_1, self).__init__(root, transform=transform,
                                      target_transform=target_transform)
        self.train = train  # training set or test set

        if download:
            self.download()

        if not self._check_integrity():
            raise RuntimeError('Dataset not found or corrupted.' +
                               ' You can use download=True to download it')

        self.data = []
        self.targets = []

        # now load the picked numpy arrays
        for file_name, checksum in CIFAR10_1.files:
            file_path = os.path.join(self.root, file_name)
            np_array = np.load(file_path)
            if 'data' in file_name:
                self.data.append(np_array)
            else:
                self.targets.append(np_array)

        self.data = np.vstack(self.data).reshape(-1, 32, 32, 3 )
        #self.data = self.data.transpose((0, 2, 3, 1))  # convert to HWC
        self.targets = np.int64(np.vstack(self.targets).reshape(-1)).tolist()

    def __getitem__(self, index):
        """
        Args:
            index (int): Index

        Returns:
            tuple: (image, target) where target is index of the target class.
        """
        img, target = self.data[index], self.targets[index]

        # doing this so that it is consistent with all other datasets
        # to return a PIL Image
        img = Image.fromarray(img)

        if self.transform is not None:
            img = self.transform(img)

        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, target

    def __len__(self):
        return len(self.data)

    def _check_integrity(self):
        root = self.root
        for fentry in CIFAR10_1.files:
            filename, md5 = fentry[0], fentry[1]
            fpath = os.path.join(root, filename)
            if not check_integrity(fpath, md5):
                return False
        return True

    def download(self):
        if self._check_integrity():
            print('Files already downloaded and verified')
            return
        for file,md in CIFAR10_1.files:
            download_url(CIFAR10_1.source_url+file, self.root, md5=md)

class _DS_META(_META):
    _ATTRS = ['nclasses', 'shape', 'mean', 'std']

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def get_normalization(self):
        return {'mean': self.mean, 'std': self.std}

__DATASETS_DEFAULT_PATH = 'Datasets'
_CIFAR10=_DS_META(nclasses=10,shape=(3,32,32),mean=[.491, .482, .446],std=[.247, .243, .261])
_CIFAR100=_DS_META(nclasses=100,shape=(3,32,32),mean=[.491, .482, .446],std=[.247, .243, .261])
_STL10=_DS_META(nclasses=10,shape=(3,32,32),mean=[.491, .482, .446],std=[.247, .243, .261])
_IMAGENET=_DS_META(nclasses=1000,shape=(3,224,224),mean=[0.485,0.456,0.406],std=[0.229,0.224,0.225])
_SVHN=_DS_META(nclasses=10,shape=(3,32,32),mean=[0.437,0.444,0.473],std=[0.198,0.201,0.197])
_MNIST=_DS_META(nclasses=10,shape=(3,32,32),mean=[0.5],std=[0.5])
_MNIST_3C=_DS_META(nclasses=10,shape=(3,32,32),mean=[0.131]*3,std=[0.308]*3)


_DATASET_META_DATA={
    'cifar10':_CIFAR10,
    'cifar100':_CIFAR100,
    'stl10':_STL10,
    'imagenet':_IMAGENET,
    'SVHN':_SVHN,
    'mnist':_MNIST,
    'mnist_3c':_MNIST_3C,
}

_IMAGINE_CONFIGS=[
    'no_dd_kl','no_dd_mse','no_dd_sym',
    'dd-exp_kl','dd-ce_kl','dd-exp_mse',
    'dd-exp', 'dd-ce']

#todo decuple limit_ds and get_dataset
def get_dataset(name, split='train', transform=None,
                target_transform=None, download=True, datasets_path=__DATASETS_DEFAULT_PATH,
                limit=None,shuffle_before_limit=False,limit_shuffle_seed=None,class_ids=None,
                per_class_limit=True,split_start = 0):
    train = (split == 'train')
    if '+' in name:
        ds=None
        for ds_name in name.split('+'):
            ds_=get_dataset(ds_name, split, transform, target_transform,download, limit=limit,
                            shuffle_before_limit=shuffle_before_limit,datasets_path=__DATASETS_DEFAULT_PATH,split_start=split_start)
            if ds is None:
                ds = ds_
            else:
                ds += ds_
                ds.targets = ds.datasets[0].targets + ds.datasets[1].targets

        if limit or class_ids:
            ds = limit_ds(ds, limit, per_class=per_class_limit, shuffle=shuffle_before_limit,
                          seed=limit_shuffle_seed,
                          allowed_classes=class_ids, split_start=split_start)
        return ds

    if name.endswith('_3c'):
        transform = tv.transforms.Compose(
            [tv.transforms.ToTensor(), lambda x: x.repeat(3, 1, 1), tv.transforms.ToPILImage(), transform])
        name = name[:-3]
    if name.endswith('-raw'):
        ds_dir_name = name[:-4]
    elif name.startswith('folder-'):
        ds_dir_name = name[7:]
    elif name == 'places365_standard-lsun':
        ds_dir_name = 'places365_standard'
        name = ds_dir_name
        class_ids = filter(lambda x: x not in [52, 66, 91, 92, 102, 121, 203, 215, 284, 334], range(365))
    elif name.startswith('DomainNet-'):
        parts = name.split('-')
        domain = parts[1]
        if name.endswith('-measure') and train:
            ds_dir_name = os.path.join('DomainNet', 'measure', domain)
        else:
            ds_dir_name = os.path.join('DomainNet', 'train' if train else 'test', domain)
        name = name.replace('-measure', '')
        if name.endswith('-A') or name.endswith('-B'):
            set = parts[2]
            class_ids = range(173) if set == 'A' else range(173, 345)
    elif name.endswith('-dogs') or name.endswith('-cats'):
        if name.startswith('imagenet-'):
            if name.endswith('dogs'):
                _ids = _imagenet_dogs.keys()
            else:
                _ids = _imagenet_cats.keys()
        else:
            _ids = [1] if name.endswith('dogs') else [0]

        return get_dataset(name[:-5], split, transform, target_transform, download, limit=limit,
                       shuffle_before_limit=True, datasets_path=__DATASETS_DEFAULT_PATH,
                       class_ids=_ids,per_class_limit=False,limit_shuffle_seed=0,split_start=split_start)
    elif name.startswith('imagine-'):
        if train:
            ds_dir_name=None
            for i_cfg in _IMAGINE_CONFIGS:
                idx=name.find(i_cfg)
                if idx >0:
                    ds_dir_name=os.path.join(name[:idx-1],i_cfg,name[idx+len(i_cfg)+1:])
                    #print(ds_dir_name)
                    break
            assert ds_dir_name is not None
        else:
            return get_dataset(name.split('-')[1], split, transform, target_transform,download, limit=limit,
                               shuffle_before_limit=shuffle_before_limit,
                               datasets_path=__DATASETS_DEFAULT_PATH,split_start=split_start)
    else:
        ds_dir_name = name

    root = os.path.join(datasets_path, ds_dir_name)
    if name == 'cifar10':
        return datasets.CIFAR10(root=root,
                                train=train,
                                transform=transform,
                                target_transform=target_transform,
                                download=download)
    elif name == 'cifar100':
        return datasets.CIFAR100(root=root,
                                 train=train,
                                 transform=transform,
                                 target_transform=target_transform,
                                 download=download)
    elif name == 'cifar10.1':
        return CIFAR10_1(root=root,
                         transform=transform,
                         target_transform=target_transform,
                         download=download)
    elif name.startswith('cifar10_custom'):
        ds_ = []
        if 'val' in split:
            ds_.append(get_dataset('cifar10',split='val',transform=transform,download=download,
                                target_transform=target_transform))
        if '10.1' in split:
            ds_.append(get_dataset('cifar10.1', split='', transform=transform, download=download,
                             target_transform=target_transform))
        if 'train' in split:
            ds_.append(get_dataset('cifar10', split='train', transform=transform, download=download,
                                   target_transform=target_transform))
        if 'ext' in split:
            ds_.append(get_dataset('folder-cifar10_ext', split='val', transform=transform, download=download,
                                   target_transform=target_transform))
        ds = ds_[0]
        for d in ds_[1:]:
            ds += d

        return limit_ds(ds, limit=limit, split_start=split_start, per_class=per_class_limit,
                        shuffle=shuffle_before_limit,
                        seed=limit_shuffle_seed, allowed_classes=class_ids)
    elif name.lower().startswith('svhn_custom'):
        ds_ = []
        if 'val' in split or 'test' in split:
            ds_.append(get_dataset('SVHN', split='test', transform=transform, download=download,
                                   target_transform=target_transform))
        if 'train' in split:
            ds_.append(get_dataset('SVHN', split='train', transform=transform, download=download,
                                   target_transform=target_transform))

        ds = ds_[0]
        for d in ds_[1:]:
            ds += d
        return limit_ds(ds, limit=limit, split_start=split_start, per_class=per_class_limit,
                        shuffle=shuffle_before_limit,
                        seed=limit_shuffle_seed, allowed_classes=class_ids)
    elif name.startswith('cifar100_custom'):
        ds_ = []
        if 'val' in split or 'test' in split:
            ds_.append(get_dataset('cifar100', split='test', transform=transform, download=download,
                                   target_transform=target_transform))
        if 'train' in split:
            ds_.append(get_dataset('cifar100', split='train', transform=transform, download=download,
                                   target_transform=target_transform))

        ds = ds_[0]
        for d in ds_[1:]:
            ds += d
        return limit_ds(ds, limit=limit, split_start=split_start, per_class=per_class_limit,
                        shuffle=shuffle_before_limit,
                        seed=limit_shuffle_seed, allowed_classes=class_ids)

    elif name == 'mnist' or name == 'mnist_3c':
        return datasets.MNIST(root=root,
                              train=train,
                              transform=transform,
                              target_transform=target_transform,
                              download=download)
    elif name == 'SVHN':
        return datasets.SVHN(root=root,
                             split='test' if not train else 'train',
                             transform=transform,
                             target_transform=target_transform,
                             download=download)
    elif 'stl10' in name:
        if train and name.endswith('train_test'):
            return datasets.STL10(root=root,
                                  split='train',
                                  transform=transform,
                                  target_transform=target_transform,
                                  download=download) + datasets.STL10(root=root,
                                                                      split='test',
                                                                      transform=transform,
                                                                      target_transform=target_transform,
                                                                      download=download)
        return datasets.STL10(root=root,
                              split=split,
                              transform=transform,
                              target_transform=target_transform,
                              download=download)
    elif name == 'LSUN':
        return datasets.LSUN(root=root,
                             classes=split,
                             transform=transform,
                             target_transform=target_transform)
    elif name.startswith('folder'):
        ds = datasets.ImageFolder(root=root,
                                  transform=transform,
                                  target_transform=target_transform)
        if limit or class_ids:
            ds = limit_ds(ds, limit, per_class=per_class_limit, shuffle=shuffle_before_limit, seed=limit_shuffle_seed,
                          allowed_classes=class_ids,split_start=split_start)
        return ds
    elif name in ['imagenet', 'cats_vs_dogs', 'places365_standard'] or any(i in name for i in ['imagine-', '-raw']):
        if train:
            root = os.path.join(root, 'train')
        else:
            root = os.path.join(root, 'val')
        ds = datasets.ImageFolder(root=root,
                                  transform=transform,
                                  target_transform=target_transform)
        if limit or class_ids:
            if 'no_dd' in name:
                ds = limit_ds(ds, limit * len(ds.classes), per_class=False, shuffle=shuffle_before_limit,
                                             seed=limit_shuffle_seed,split_start=split_start)
            else:
                ds=limit_ds(ds, limit,per_class=per_class_limit,shuffle=shuffle_before_limit,seed=limit_shuffle_seed,allowed_classes=class_ids,split_start=split_start)
        return ds
    elif name.startswith('DomainNet-'):
        ds = datasets.ImageFolder(root=root,
                                    transform=transform,
                                    target_transform=target_transform)
        if limit or class_ids:
            return limit_ds(ds, limit=limit, per_class=per_class_limit, shuffle=shuffle_before_limit, seed=limit_shuffle_seed, allowed_classes=class_ids)
        return ds

    elif name.startswith('random-'):
        np.random.seed(limit_shuffle_seed)
        n_samples = limit or 10000
        dummy_targets = torch.ones(n_samples)

        ds_name = name[7:]
        if name.endswith('-normal'):
            mean, std = [0., 0., 0.], [1., 1., 1.]
            return RandomDatasetGenerator([3, 512, 512], mean, std, limit=n_samples, transform=transform,
                                          train=train)
        use_random_test = False
        if name.endswith('-rt'):
            use_random_test = True
            ds_name = ds_name[:-3]

        if ds_name in _DATASET_META_DATA:
            meta = _DATASET_META_DATA[ds_name]
            nclasses, data_shape, mean, std = meta.get_attrs().values()

        ## borrowed from https://github.com/hendrycks/outlier-exposure/
        elif ds_name == 'gaussian':
            samples = torch.from_numpy(
                np.clip(np.random.normal(size=(n_samples, 3, 32, 32),
                                         loc=0.5, scale=0.5).astype(np.float32), 0, 1))
            return torch.utils.data.TensorDataset(samples, dummy_targets)
        elif ds_name == 'rademacher' or ds_name == 'bernoulli':
            samples = torch.from_numpy(np.random.binomial(n=1, p=0.5, size=(n_samples, 3, 32, 32)).astype(np.float32))
            if ds_name == 'rademacher':
                samples = samples * 2 - 1
            return torch.utils.data.TensorDataset(samples, dummy_targets)
        elif ds_name == 'blob':
            from skimage.filters import gaussian as gblur
            samples = np.float32(np.random.binomial(n=1, p=0.7, size=(n_samples, 32, 32, 3)))
            for i in range(n_samples):
                samples[i] = gblur(samples[i], sigma=1.5, multichannel=False)
                samples[i][samples[i] < 0.75] = 0.0

            samples = torch.from_numpy(samples.transpose((0, 3, 1, 2)))
            return torch.utils.data.TensorDataset(samples, dummy_targets)
        else:
            raise NotImplementedError

        limit = limit or 1000
        if per_class_limit:
            limit = limit * nclasses

        if train or use_random_test:
            return RandomDatasetGenerator(data_shape, mean, std, limit=limit, transform=transform, train=train)
        else:
            return get_dataset(ds_name, split, transform, target_transform, download, limit=limit,
                               shuffle_before_limit=shuffle_before_limit, datasets_path=__DATASETS_DEFAULT_PATH,
                               split_start=split_start)

    elif hasattr(datasets, name):
        return getattr(datasets, name)(root=datasets_path,
                                       train=train,
                                       transform=transform,
                                       target_transform=target_transform,
                                       download=download)


# def balance_image_folder_ds(dataset, n_samples=None,per_class=True,shuffle=False,seed=None,class_ids=None):
#     assert isinstance(dataset,datasets.DatasetFolder)
#
#     if shuffle:
#         import random
#         random.seed(seed)
#         print(f'shufflling with seed {seed}')
#     samps = []
#
#     n_samples = n_samples or len(dataset)
#     if per_class or class_ids is not None:
#         samp_reg_per_class={}
#         for s in dataset.samples:
#             if s[1] in samp_reg_per_class:
#                 if shuffle or len(samp_reg_per_class[s[1]])<n_samples:
#                     samp_reg_per_class[s[1]]+=[s]
#             elif class_ids is not None and s[1] not in class_ids:
#                 continue
#             else:
#                 samp_reg_per_class[s[1]]=[s]
#
#         for k in samp_reg_per_class.keys():
#             if shuffle and per_class:
#                 samps += random.sample(samp_reg_per_class[k],n_samples)
#             else:
#                 samps += samp_reg_per_class[k]
#
#         if not per_class and shuffle and len(samps)>n_samples:
#             samps = random.sample(samps, n_samples)
#     else:
#         if shuffle:
#             samps = random.sample(dataset.samples,n_samples)
#         else:
#             samps = dataset.samples[:n_samples]
#
#     if hasattr(dataset,'imgs'):
#         dataset.imgs = samps
#     dataset.samples = samps
#     return dataset

def limit_ds(dataset, limit=None, per_class=True, shuffle=False, seed=0, allowed_classes=None,split_start=0,verbose=0):
    if not hasattr(dataset,'targets'):
        if hasattr(dataset,'labels'):
            dataset.targets=dataset.labels
        else:
            assert 0, 'dataset not supported'

    id_reg_per_class = {}
    global_sample_count = 0
    # map id to class label
    for e, t in enumerate(dataset.targets):
        if allowed_classes and t not in allowed_classes:
            continue
        global_sample_count += 1
        if t in id_reg_per_class:
            id_reg_per_class[t] += [e]
        else:
            id_reg_per_class[t] = [e]
    n_classes=len(id_reg_per_class)
    ids = []
    # shuffle and clip each class
    for t in id_reg_per_class.keys():
        class_ids = torch.tensor(id_reg_per_class[t])
        samples_in_class=len(class_ids)
        if verbose:
            print(f'class {t} total samples: {samples_in_class}, '
              f'proportion in ds:({samples_in_class / global_sample_count:0.3f})')

        if shuffle:
            class_ids = class_ids[torch.randperm(samples_in_class, generator=torch.Generator().manual_seed(seed))]

        if limit is not None:
            lim = limit
            if limit >= 1 and not per_class and type(limit) == int:
                # update limit to per class value
                lim = max(limit // n_classes, 1)
            elif 0 < limit <= 1 and type(limit) == float:
                # limit is given as a ratio from each class
                lim = max(int(limit * samples_in_class), 1)

            start=split_start
            if split_start >= 1 and not per_class:
                # update limit to per class value
                start = max(split_start // n_classes, 1)
            elif 0<split_start<1:
                # limit is given as a ratio from each class
                start = max(int(split_start * samples_in_class), 1)

            class_ids = class_ids[start:start+lim]
            if verbose:
                print(f'clipping data limit: {lim}\n'
                      f'data slice: id start {start}, id stop {lim+start}')
                print(f'actual size of class {t}: {len(class_ids)}\n')

        ids.append(class_ids)


    all_ids = torch.cat(ids)

    ds = torch.utils.data.Subset(dataset,all_ids)
    ds.targets = torch.tensor(dataset.targets)[all_ids]
    ds.classes = dataset.classes if hasattr(dataset, 'classes') else list(range(len(id_reg_per_class)))
    ds.per_class_ids=ids
    ds.all_ids=all_ids
    if allowed_classes:
        ds.classes = [c for i, c in enumerate(ds.classes) if i in allowed_classes]

    return ds

_imagenet_dogs = {
 151: 'Chihuahua',
 152: 'Japanese spaniel',
 153: 'Maltese dog, Maltese terrier, Maltese',
 154: 'Pekinese, Pekingese, Peke',
 155: 'Shih-Tzu',
 156: 'Blenheim spaniel',
 157: 'papillon',
 158: 'toy terrier',
 159: 'Rhodesian ridgeback',
 160: 'Afghan hound, Afghan',
 161: 'basset, basset hound',
 162: 'beagle',
 163: 'bloodhound, sleuthhound',
 164: 'bluetick',
 165: 'black-and-tan coonhound',
 166: 'Walker hound, Walker foxhound',
 167: 'English foxhound',
 168: 'redbone',
 169: 'borzoi, Russian wolfhound',
 170: 'Irish wolfhound',
 171: 'Italian greyhound',
 172: 'whippet',
 173: 'Ibizan hound, Ibizan Podenco',
 174: 'Norwegian elkhound, elkhound',
 175: 'otterhound, otter hound',
 176: 'Saluki, gazelle hound',
 177: 'Scottish deerhound, deerhound',
 178: 'Weimaraner',
 179: 'Staffordshire bullterrier, Staffordshire bull terrier',
 180: 'American Staffordshire terrier, Staffordshire terrier, American pit bull terrier, pit bull terrier',
 181: 'Bedlington terrier',
 182: 'Border terrier',
 183: 'Kerry blue terrier',
 184: 'Irish terrier',
 185: 'Norfolk terrier',
 186: 'Norwich terrier',
 187: 'Yorkshire terrier',
 188: 'wire-haired fox terrier',
 189: 'Lakeland terrier',
 190: 'Sealyham terrier, Sealyham',
 191: 'Airedale, Airedale terrier',
 192: 'cairn, cairn terrier',
 193: 'Australian terrier',
 194: 'Dandie Dinmont, Dandie Dinmont terrier',
 195: 'Boston bull, Boston terrier',
 196: 'miniature schnauzer',
 197: 'giant schnauzer',
 198: 'standard schnauzer',
 199: 'Scotch terrier, Scottish terrier, Scottie',
 200: 'Tibetan terrier, chrysanthemum dog',
 201: 'silky terrier, Sydney silky',
 202: 'soft-coated wheaten terrier',
 203: 'West Highland white terrier',
 204: 'Lhasa, Lhasa apso',
 205: 'flat-coated retriever',
 206: 'curly-coated retriever',
 207: 'golden retriever',
 208: 'Labrador retriever',
 209: 'Chesapeake Bay retriever',
 210: 'German short-haired pointer',
 211: 'vizsla, Hungarian pointer',
 212: 'English setter',
 213: 'Irish setter, red setter',
 214: 'Gordon setter',
 215: 'Brittany spaniel',
 216: 'clumber, clumber spaniel',
 217: 'English springer, English springer spaniel',
 218: 'Welsh springer spaniel',
 219: 'cocker spaniel, English cocker spaniel, cocker',
 220: 'Sussex spaniel',
 221: 'Irish water spaniel',
 222: 'kuvasz',
 223: 'schipperke',
 224: 'groenendael',
 225: 'malinois',
 226: 'briard',
 227: 'kelpie',
 228: 'komondor',
 229: 'Old English sheepdog, bobtail',
 230: 'Shetland sheepdog, Shetland sheep dog, Shetland',
 231: 'collie',
 232: 'Border collie',
 233: 'Bouvier des Flandres, Bouviers des Flandres',
 234: 'Rottweiler',
 235: 'German shepherd, German shepherd dog, German police dog, alsatian',
 236: 'Doberman, Doberman pinscher',
 237: 'miniature pinscher',
 238: 'Greater Swiss Mountain dog',
 239: 'Bernese mountain dog',
 240: 'Appenzeller',
 241: 'EntleBucher',
 242: 'boxer',
 243: 'bull mastiff',
 244: 'Tibetan mastiff',
 245: 'French bulldog',
 246: 'Great Dane',
 247: 'Saint Bernard, St Bernard',
 248: 'Eskimo dog, husky',
 249: 'malamute, malemute, Alaskan malamute',
 250: 'Siberian husky',
 251: 'dalmatian, coach dog, carriage dog',
 252: 'affenpinscher, monkey pinscher, monkey dog',
 253: 'basenji',
 254: 'pug, pug-dog',
 255: 'Leonberg',
 256: 'Newfoundland, Newfoundland dog',
 257: 'Great Pyrenees',
 258: 'Samoyed, Samoyede',
 259: 'Pomeranian',
 260: 'chow, chow chow',
 261: 'keeshond',
 262: 'Brabancon griffon',
 263: 'Pembroke, Pembroke Welsh corgi',
 264: 'Cardigan, Cardigan Welsh corgi',
 265: 'toy poodle',
 266: 'miniature poodle',
 267: 'standard poodle',
 268: 'Mexican hairless'}

_imagenet_cats = {
 281: 'tabby, tabby cat',
 282: 'tiger cat',
 283: 'Persian cat',
 284: 'Siamese cat, Siamese',
 285: 'Egyptian cat'}