File size: 4,900 Bytes
3ea5987 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | 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
# num_class = 100
MEAN = (0.48145446, 0.4578275, 0.40821073)
VAR = (0.26862954, 0.26130258, 0.27577711)
class Joint:
# class Noise: ### noise_data_1005
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 ############################################# twice!!!
assert self.num_pic <= 150, "please provide less than 150 images per class"
# self.classnames = get_classnames(classnames)[:num_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
### ImageNet
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
### Noise
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_dataset = self.get_test_datasset()
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(self):
# return "joint/train"
def name_ImageNet(self):
return "imagenet/train"
def name_Noise(self):
return 'noise_data_1000100'
|