| import os |
|
|
| import numpy as np |
| import torch |
| import torchvision |
| from torchvision.datasets import CIFAR10 as PyTorchCIFAR10 |
| from torchvision.datasets import CIFAR100 as PyTorchCIFAR100 |
| from torchvision.datasets import VisionDataset |
|
|
| cifar10_template = [ |
| lambda c : f'a photo of a {c}.', |
| lambda c : f'a blurry photo of a {c}.', |
| lambda c : f'a black and white photo of a {c}.', |
| lambda c : f'a low contrast photo of a {c}.', |
| lambda c : f'a high contrast photo of a {c}.', |
| lambda c : f'a bad photo of a {c}.', |
| lambda c : f'a good photo of a {c}.', |
| lambda c : f'a photo of a small {c}.', |
| lambda c : f'a photo of a big {c}.', |
| lambda c : f'a photo of the {c}.', |
| lambda c : f'a blurry photo of the {c}.', |
| lambda c : f'a black and white photo of the {c}.', |
| lambda c : f'a low contrast photo of the {c}.', |
| lambda c : f'a high contrast photo of the {c}.', |
| lambda c : f'a bad photo of the {c}.', |
| lambda c : f'a good photo of the {c}.', |
| lambda c : f'a photo of the small {c}.', |
| lambda c : f'a photo of the big {c}.', |
| ] |
|
|
|
|
| DOWNLOAD = True |
| cifar_classnames = [ |
| "airplane", |
| "automobile", |
| "bird", |
| "cat", |
| "deer", |
| "dog", |
| "frog", |
| "horse", |
| "ship", |
| "truck", |
| ] |
|
|
|
|
| class CIFAR10: |
| def __init__( |
| self, |
| preprocess, |
| location="./data", |
| batch_size=128, |
| num_workers=16, |
| classnames=None, |
| ): |
|
|
| self.train_dataset = PyTorchCIFAR10( |
| root=location, download=DOWNLOAD, train=True, transform=preprocess |
| ) |
|
|
| self.train_loader = torch.utils.data.DataLoader( |
| self.train_dataset, |
| batch_size=batch_size, |
| shuffle=True, |
| num_workers=num_workers, |
| ) |
|
|
| self.test_dataset = PyTorchCIFAR10( |
| root=location, download=DOWNLOAD, train=False, transform=preprocess |
| ) |
|
|
| self.test_loader = torch.utils.data.DataLoader( |
| self.test_dataset, |
| batch_size=batch_size, |
| shuffle=False, |
| num_workers=num_workers, |
| ) |
|
|
| self.classnames = self.test_dataset.classes |
| self.template = cifar10_template |
|
|
|
|
| def convert(x): |
| if isinstance(x, np.ndarray): |
| return torchvision.transforms.functional.to_pil_image(x) |
| return x |
|
|
|
|
| class BasicVisionDataset(VisionDataset): |
| def __init__(self, images, targets, transform=None, target_transform=None): |
| if transform is not None: |
| transform.transforms.insert(0, convert) |
| super(BasicVisionDataset, self).__init__( |
| root=None, transform=transform, target_transform=target_transform |
| ) |
| assert len(images) == len(targets) |
|
|
| self.images = images |
| self.targets = targets |
|
|
| def __getitem__(self, index): |
| return self.transform(self.images[index]), self.targets[index] |
|
|
| def __len__(self): |
| return len(self.targets) |
|
|
|
|
| class CIFAR101: |
| def __init__( |
| self, |
| preprocess, |
| location=os.path.expanduser("~/data"), |
| batch_size=128, |
| num_workers=16, |
| classnames=None, |
| ): |
|
|
| data_root = os.path.join(location, "CIFAR-10.1") |
| data = np.load( |
| os.path.join(data_root, "cifar10.1_v6_data.npy"), allow_pickle=True |
| ) |
| labels = np.load( |
| os.path.join(data_root, "cifar10.1_v6_labels.npy"), allow_pickle=True |
| ) |
|
|
| use_cuda = torch.cuda.is_available() |
|
|
| |
| kwargs = {"num_workers": num_workers, "pin_memory": True} if use_cuda else {} |
|
|
| self.train_loader = None |
| self.location = location |
|
|
| self.test_dataset = BasicVisionDataset( |
| images=data, |
| targets=torch.Tensor(labels).long(), |
| transform=preprocess, |
| ) |
|
|
| self.test_loader = torch.utils.data.DataLoader( |
| self.test_dataset, batch_size=batch_size, shuffle=False, **kwargs |
| ) |
|
|
| self.classnames = cifar_classnames |
| self.template = cifar10_template |
|
|
|
|
| class CIFAR102: |
| def __init__( |
| self, |
| preprocess, |
| location=os.path.expanduser("~/data"), |
| batch_size=32, |
| num_workers=16, |
| classnames=None, |
| ): |
|
|
| train_data = np.load( |
| os.path.join(location, "CIFAR-10.2", "cifar102_train.npy"), |
| allow_pickle=True, |
| ).item() |
| test_data = np.load( |
| os.path.join(location, "CIFAR-10.2", "cifar102_test.npy"), allow_pickle=True |
| ).item() |
|
|
| train_data_images = train_data["images"] |
| train_data_labels = train_data["labels"] |
|
|
| test_data_images = test_data["images"] |
| test_data_labels = test_data["labels"] |
|
|
| use_cuda = torch.cuda.is_available() |
|
|
| |
| kwargs = {"num_workers": num_workers, "pin_memory": True} if use_cuda else {} |
|
|
| self.test_dataset = BasicVisionDataset( |
| images=test_data_images, |
| targets=torch.Tensor(test_data_labels).long(), |
| transform=preprocess, |
| ) |
|
|
| self.test_loader = torch.utils.data.DataLoader( |
| self.test_dataset, batch_size=batch_size, shuffle=False, **kwargs |
| ) |
|
|
| self.train_dataset = BasicVisionDataset( |
| images=train_data_images, |
| targets=torch.Tensor(train_data_labels).long(), |
| transform=preprocess, |
| ) |
|
|
| self.train_loader = torch.utils.data.DataLoader( |
| self.train_dataset, batch_size=batch_size, shuffle=False, **kwargs |
| ) |
|
|
| self.classnames = cifar_classnames |
| self.template = cifar10_template |
|
|
|
|
| class CIFAR100: |
| def __init__( |
| self, |
| preprocess, |
| location="./data", |
| batch_size=128, |
| num_workers=16, |
| classnames=None, |
| ): |
|
|
| self.train_dataset = PyTorchCIFAR100( |
| root=location, download=DOWNLOAD, train=True, transform=preprocess |
| ) |
|
|
| self.train_loader = torch.utils.data.DataLoader( |
| self.train_dataset, batch_size=batch_size, num_workers=num_workers |
| ) |
|
|
| self.test_dataset = PyTorchCIFAR100( |
| root=location, download=DOWNLOAD, train=False, transform=preprocess |
| ) |
|
|
| self.test_loader = torch.utils.data.DataLoader( |
| self.test_dataset, |
| batch_size=batch_size, |
| shuffle=False, |
| num_workers=num_workers, |
| ) |
|
|
| self.classnames = self.test_dataset.classes |
| self.template = cifar10_template |
|
|
