| """ |
| Code adapted from https://github.com/mlfoundations/wise-ft/blob/master/src/datasets/objectnet.py |
| Thanks to the authors of wise-ft |
| """ |
|
|
| import json |
| import os |
| from pathlib import Path |
|
|
| import numpy as np |
| import PIL |
| import torch |
| from torchvision import datasets |
| from torchvision.transforms import Compose |
|
|
|
|
| def get_metadata(folder): |
| metadata = Path(folder) |
|
|
| with open(metadata / "folder_to_objectnet_label.json", "r") as f: |
| folder_map = json.load(f) |
| folder_map = {v: k for k, v in folder_map.items()} |
| with open(metadata / "objectnet_to_imagenet_1k.json", "r") as f: |
| objectnet_map = json.load(f) |
|
|
| with open(metadata / "pytorch_to_imagenet_2012_id.json", "r") as f: |
| pytorch_map = json.load(f) |
| pytorch_map = {v: k for k, v in pytorch_map.items()} |
|
|
| with open(metadata / "imagenet_to_label_2012_v2", "r") as f: |
| imagenet_map = {v.strip(): str(pytorch_map[i]) for i, v in enumerate(f)} |
|
|
| folder_to_ids, class_sublist = {}, [] |
| classnames = [] |
| for objectnet_name, imagenet_names in objectnet_map.items(): |
| imagenet_names = imagenet_names.split("; ") |
| imagenet_ids = [ |
| int(imagenet_map[imagenet_name]) for imagenet_name in imagenet_names |
| ] |
| class_sublist.extend(imagenet_ids) |
| folder_to_ids[folder_map[objectnet_name]] = imagenet_ids |
|
|
| class_sublist = sorted(class_sublist) |
| class_sublist_mask = [(i in class_sublist) for i in range(1000)] |
| classname_map = {v: k for k, v in folder_map.items()} |
| return class_sublist, class_sublist_mask, folder_to_ids, classname_map |
|
|
|
|
| class ObjectNetDataset(datasets.ImageFolder): |
|
|
| def __init__(self, root, transform): |
| ( |
| self._class_sublist, |
| self.class_sublist_mask, |
| self.folders_to_ids, |
| self.classname_map, |
| ) = get_metadata(root) |
| subdir = os.path.join(root, "objectnet-1.0", "images") |
| label_map = { |
| name: idx |
| for idx, name in enumerate(sorted(list(self.folders_to_ids.keys()))) |
| } |
| self.label_map = label_map |
| super().__init__(subdir, transform=transform) |
| self.samples = [ |
| d |
| for d in self.samples |
| if os.path.basename(os.path.dirname(d[0])) in self.label_map |
| ] |
| self.imgs = self.samples |
| self.classes = sorted(list(self.folders_to_ids.keys())) |
| self.classes = [self.classname_map[c].lower() for c in self.classes] |
|
|
| def __len__(self): |
| return len(self.samples) |
|
|
| def __getitem__(self, index): |
| path, target = self.samples[index] |
| sample = self.loader(path) |
| if self.transform is not None: |
| sample = self.transform(sample) |
| label = os.path.basename(os.path.dirname(path)) |
| return sample, self.label_map[label] |
|
|
