| import json |
| from torchvision.datasets import ImageFolder |
| import torch |
| import os |
| from PIL import Image |
| import collections |
| import torchvision.transforms as transforms |
| from label_str_to_imagenet_classes import label_str_to_imagenet_classes |
|
|
| torch.manual_seed(0) |
|
|
| ImageItem = collections.namedtuple('ImageItem', ('image_name', 'tag')) |
|
|
| normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5], |
| std=[0.5, 0.5, 0.5]) |
|
|
| transform = transforms.Compose([ |
| transforms.Resize(256), |
| transforms.CenterCrop(224), |
| transforms.ToTensor(), |
| normalize, |
| ]) |
|
|
| class RobustnessDataset(ImageFolder): |
| def __init__(self, imagenet_path, folder, imagenet_classes_path='imagenet_classes.json', isV2=False, isSI=False): |
| self._isV2 = isV2 |
| self._isSI = isSI |
| self._folder = folder |
| self._imagenet_path = imagenet_path |
| with open(imagenet_classes_path, 'r') as f: |
| self._imagenet_classes = json.load(f) |
| self._all_images = [] |
|
|
| base_dir = os.path.join(self._imagenet_path, folder) |
| for i, file in enumerate(os.listdir(base_dir)): |
| self._all_images.append(ImageItem(file, folder)) |
|
|
|
|
| def __getitem__(self, item): |
| image_item = self._all_images[item] |
| image_path = os.path.join(self._imagenet_path, image_item.tag, image_item.image_name) |
| image = Image.open(image_path) |
| image = image.convert('RGB') |
| image = transform(image) |
|
|
| if self._isV2: |
| class_name = int(image_item.tag) |
| elif self._isSI: |
| class_name = int(label_str_to_imagenet_classes[image_item.tag]) |
| else: |
| class_name = int(self._imagenet_classes[image_item.tag]) |
|
|
| return image, class_name |
|
|
| def __len__(self): |
| return len(self._all_images) |
|
|
| def get_classname(self): |
| if self._isV2: |
| class_name = int(self._folder) |
| elif self._isSI: |
| class_name = int(label_str_to_imagenet_classes[self._folder]) |
| else: |
| class_name = int(self._imagenet_classes[self._folder]) |
| return class_name |
