| import os |
|
|
| import numpy as np |
| import torch |
| import torchvision.transforms as T |
|
|
| from .common import ImageFolderWithPaths, SubsetSampler |
| from .imagenet_classnames import get_classnames |
| from ..templates.openai_imagenet_template import openai_imagenet_template |
|
|
| |
| MEAN = (0.48145446, 0.4578275, 0.40821073) |
| VAR = (0.26862954, 0.26130258, 0.27577711) |
|
|
| class Joint: |
| |
| def __init__(self, |
| preprocess, |
| location=os.path.expanduser('~/data'), |
| batch_size=32, |
| num_workers=32, |
| classnames='openai'): |
| self.preprocess_noise = T.Compose([T.ToTensor(), T.Normalize(MEAN, VAR)]) |
| self.preprocess = preprocess |
| self.location = location |
| self.batch_size = batch_size |
| self.num_workers = num_workers |
|
|
| self.num_class = 10 |
| self.num_pic = 100 |
| assert self.num_pic <= 150, "please provide less than 150 images per class" |
|
|
| |
| self.classnames = get_classnames(classnames)[:self.num_class] |
| self.template = openai_imagenet_template |
|
|
| self.populate_train() |
| self.populate_test() |
| |
| def populate_train(self): |
| SUBCLASS = np.arange(1000)[:self.num_class].tolist() |
| ImageNetDir = os.path.join(self.location, self.name_ImageNet()) |
| NoiseDir = os.path.join(self.location, self.name_Noise()) |
|
|
| self.train_dataset = ImageFolderWithPaths( |
| ImageNetDir, |
| transform=self.preprocess) |
|
|
| self.noise_dataset = ImageFolderWithPaths( |
| NoiseDir, |
| transform=self.preprocess_noise) |
|
|
| samples = [] |
| targets = [] |
| classes = [] |
| storage = {} |
| dic = self.train_dataset.class_to_idx |
| noise_dic = self.noise_dataset.class_to_idx |
| |
| |
| for cla in self.train_dataset.classes: |
| if dic[cla] in SUBCLASS: |
| classes.append(cla) |
|
|
| for i in np.arange(len(self.train_dataset.samples)): |
| sample = self.train_dataset.samples[i] |
| target = self.train_dataset.targets[i] |
| if sample[1] in SUBCLASS: |
| if sample[1] not in storage.keys(): |
| storage[sample[1]] = 1 |
| samples.append(sample) |
| targets.append(target) |
| else: |
| if storage[sample[1]] < self.num_pic/2: |
| samples.append(sample) |
| targets.append(target) |
| storage[sample[1]] += 1 |
|
|
| |
| for cla in self.noise_dataset.classes: |
| if noise_dic[cla] in SUBCLASS: |
| classes.append(cla) |
|
|
| for i in np.arange(len(self.noise_dataset.samples)): |
| sample = self.noise_dataset.samples[i] |
| target = self.noise_dataset.targets[i] |
| if sample[1] in SUBCLASS: |
| if sample[1] not in storage.keys(): |
| storage[sample[1]] = 1 |
| samples.append(sample) |
| targets.append(target) |
| print("error") |
| breakpoint() |
| else: |
| if storage[sample[1]] < self.num_pic: |
| samples.append(sample) |
| targets.append(target) |
| storage[sample[1]] += 1 |
|
|
|
|
|
|
| self.train_dataset.classes = classes |
| self.train_dataset.samples = samples |
| self.train_dataset.targets = targets |
|
|
| self.train_loader = torch.utils.data.DataLoader( |
| self.train_dataset, |
| batch_size=self.batch_size, |
| num_workers=self.num_workers, |
| shuffle=False |
| ) |
|
|
| def populate_test(self): |
| |
| self.test_loader = torch.utils.data.DataLoader( |
| self.train_dataset, |
| batch_size=self.batch_size, |
| num_workers=self.num_workers, |
| sampler=self.get_test_sampler() |
| ) |
|
|
| def get_test_path(self): |
| test_path = os.path.join(self.location, self.name(), 'val_in_folder') |
| if not os.path.exists(test_path): |
| test_path = os.path.join(self.location, self.name(), 'val') |
| return test_path |
|
|
| def get_train_sampler(self): |
| return None |
|
|
| def get_test_sampler(self): |
| return None |
|
|
| def get_test_dataset(self): |
| return ImageFolderWithPaths(self.get_test_path(), transform=self.preprocess) |
|
|
| |
| |
| def name_ImageNet(self): |
| return "imagenet/train" |
|
|
| def name_Noise(self): |
| return 'noise_data_1000100' |
| |
|
|
|
|
