ood/datasets.py

# Copyright (c) Alibaba, Inc. and its affiliates.
import os, ast
import numpy as np
from PIL import Image

import torch
import torchvision
from torch.utils.data import DataLoader, Subset, dataloader
from torchvision import transforms


class CLIPFeatDataset(torch.utils.data.Dataset):
    def __init__(self, data_dir, epoch=0, nshot=-1, split='train'):
        super().__init__()
        self.data_dir = data_dir
        self.nshot = nshot
        self.split = split
        self.load_data(epoch)
        self.epoch = epoch

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

    def __getitem__(self, idx):
        if self.paths is None:  # for OOD 
            return self.features[idx], self.targets[idx]
        else:
            output = {
                'img': self.features[idx],
                "label": self.targets[idx],
                "domain": '',
                "impath": self.paths[idx],
                "index": idx
            }
            return output 

    def load_data(self, epoch=None):
        if hasattr(self, 'epoch') and epoch == self.epoch:
            return
        prefix = ('' if epoch is None else f'ep{epoch}_') + self.split
        if self.nshot > 0:
            prefix = f'{self.nshot}shot_' + prefix
        self.features = torch.load(f'{self.data_dir}/{prefix}_image_features.pt', map_location='cpu')
        self.targets = torch.load(f'{self.data_dir}/{prefix}_labels.pt', map_location='cpu').long()

        path_file = f'{self.data_dir}/{prefix}_paths.txt'
        if os.path.exists(path_file):
            with open(path_file, 'r') as f:
                self.paths = f.read().splitlines()
            self.epoch = epoch
            if 'train' in prefix:
                print(f'\nLoaded dataset from epoch {epoch}\n')
        else:
            self.paths = None


class LargeOODDataset(torchvision.datasets.ImageFolder):
    def __init__(self, root, id_name, ood_name, transform):
        ood_name_dict = {'texture': 'dtd', 'inaturalist': 'iNaturalist', 'places': 'Places', 'sun': 'SUN'}
        data_root = f'{root}/{ood_name_dict[ood_name]}'
        target_transform = lambda x: x * 0 - 1  # all -1

        super().__init__(data_root, transform, target_transform=target_transform)
        print("LargeOODDataset (id %s, ood %s) Contain %d images" % (id_name, ood_name, len(self)))


class SemanticOODDataset(torchvision.datasets.ImageFolder):
    def __init__(self, root, id_name, ood_name, transform):
        assert 'imagenet' in id_name and ood_name in ['easy', 'rand', 'hard']
        
        data_root = f'{root}/{id_name}/ood_{ood_name}'
        target_transform = lambda x: x * 0 - 1  # all -1

        super().__init__(data_root, transform, target_transform=target_transform)
        print("SemanticOODDataset (id %s, ood %s) Contain %d images" % (id_name, ood_name, len(self)))


class ClassOODDataset(torchvision.datasets.ImageFolder):
    def __init__(self, root, id_name, ood_name, transform):
        assert 'imagenet' in id_name
        
        if 'severity' in ood_name:
            data_root = f'{root}/{id_name}/{ood_name}'
        else:
            assert 'imagenet' in ood_name and '-o' in ood_name
            data_root = f'{root}/../{ood_name}/images'
        target_transform = lambda x: x * 0 - 1  # all -1

        super().__init__(data_root, transform, target_transform=target_transform)
        print("ClassOODDataset (id %s, ood %s) Contain %d images" % (id_name, ood_name, len(self)))


class SCOODDataset(torch.utils.data.Dataset):

    def __init__(self, root, id_name, ood_name, transform):

        super(SCOODDataset, self).__init__()

        assert id_name in ['cifar10', 'cifar100', 'imagenet', 'imagenet100']
        if 'imagenet' in id_name:
            id_name = 'cifar10'
        if ood_name == 'cifar':
            assert 'cifar' in id_name
            if id_name == 'cifar10':
                ood_name = 'cifar100'
            else:
                ood_name = 'cifar10'

        imglist_path = os.path.join(root, 'data/imglist/benchmark_%s' % id_name, 'test_%s.txt' % ood_name)

        with open(imglist_path) as fp:
            self.imglist = fp.readlines()
        
        self.transform = transform
        self.root = root

        print("SCOODDataset (id %s, ood %s) Contain %d images" % (id_name, ood_name, len(self.imglist)))

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

    def __getitem__(self, index):
        # parse the string in imglist file:
        line = self.imglist[index].strip("\n")
        tokens = line.split(" ", 1)
        image_name, extra_str = tokens[0], tokens[1]
        extras = ast.literal_eval(extra_str)
        sc_label = extras['sc_label'] # the ood label is here. -1 means ood.

        # read image according to image name:
        img_path = os.path.join(self.root, 'data', 'images', image_name)
        with open(img_path, 'rb') as f:
            img = Image.open(f).convert('RGB')
        if self.transform is not None:
            img = self.transform(img)

        return img, sc_label


class TinyImages(torch.utils.data.Dataset):

    def __init__(self, root, transform):

        super(TinyImages, self).__init__()

        self.data = np.load(os.path.join(root, 'tinyimages80m', '300K_random_images.npy'))
        self.transform = transform

        print("TinyImages Contain {} images".format(len(self.data)))

    def __getitem__(self, index):
        img = self.data[index]
        img = Image.fromarray(img)
        if self.transform is not None:
            img = self.transform(img)

        return img, -1  # -1 is the class

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


def tinyimages300k_dataloaders(num_samples=300000, train_batch_size=64, num_workers=8, data_root_path='/ssd1/haotao/datasets'):

    num_samples = int(num_samples)

    data_dir = os.path.join(data_root_path)

    train_transform = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
    ])

    train_set = Subset(TinyImages(data_dir, train=True, transform=train_transform, download=True), list(range(num_samples)))

    train_loader = DataLoader(train_set, batch_size=train_batch_size, shuffle=True, num_workers=num_workers,
                                drop_last=True, pin_memory=True)

    return train_loader
    

class _RepeatSampler:
    """ Sampler that repeats forever
    Args:
        sampler (Sampler)
    """

    def __init__(self, sampler):
        self.sampler = sampler

    def __iter__(self):
        while True:
            yield from iter(self.sampler)


class InfiniteDataLoader(dataloader.DataLoader):
    """ Dataloader that reuses workers
    Uses same syntax as vanilla DataLoader
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        object.__setattr__(self, 'batch_sampler', _RepeatSampler(self.batch_sampler))
        self.iterator = super().__iter__()

    def __len__(self):
        return len(self.batch_sampler.sampler)

    def __iter__(self):
        for _ in range(len(self)):
            yield next(self.iterator)