def load_CIFAR10():
# C10+
transform_train = transforms.Compose([transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=(0.4914, 0.4822, 0.4465), std=(0.247, 0.243, 0.261)),
])
transforms_test = transforms.Compose([transforms.ToTensor(),
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))])
train_set = datasets.CIFAR10(root='/data', train=True, transform=transform_train, download=True)
train_loader = DataLoader(train_set, batch_size=64, num_workers=4)
test_set = datasets.CIFAR10(root='/data', train=False, transform=transforms_test, download=True)
test_loader = DataLoader(test_set, batch_size=100, num_workers=4)
return train_loader, test_loader
def get_image_dataloader(path_to_data, batch_size=16):
my_transforms = transforms.Compose([
transforms.ToTensor(),
transforms.ToPILImage(),
transforms.RandomHorizontalFlip(),
transforms.RandomGrayscale(),
transforms.RandomRotation([-30, 30]),
# transforms.Normalize([],[]) # get mean and std
transforms.ToTensor(),
])
dataset = ImageFolder(root=path_to_data,
transform=my_transforms,
is_valid_file=check_valid)
print(len(dataset.classes))
dataloader = DataLoader(dataset=dataset,
batch_size=batch_size,
shuffle=True)
return dataloader, dataset
def get_cifar10_transform(size, train=True):
if train:
transform = transforms.Compose([
transforms.RandomCrop(size, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
else:
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
return transform
def get_augmentation(dataset):
if dataset in [DATA.cifar_10, DATA.cifar_100]:
transform_train = tv_transforms.Compose([
tv_transforms.RandomCrop(32, padding=4),
tv_transforms.RandomHorizontalFlip(),
tv_transforms.ToTensor(),
tv_transforms.Normalize(_CIFAR_MEAN, _CIFAR_STD)
])
transform_test = tv_transforms.Compose([
tv_transforms.ToTensor(),
tv_transforms.Normalize(_CIFAR_MEAN, _CIFAR_STD),
])
else:
raise NotImplementedError(
'Not yet implemented for dataset={}.'.format(dataset))
return transform_train, transform_test
def cifar10_train_augmented():
"""
### Augmented CIFAR 10 train dataset
"""
from torchvision.datasets import CIFAR10
from torchvision.transforms import transforms
return CIFAR10(str(lab.get_data_path()),
train=True,
download=True,
transform=transforms.Compose([
# Pad and crop
transforms.RandomCrop(32, padding=4),
# Random horizontal flip
transforms.RandomHorizontalFlip(),
#
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
]))
def __init__(
self,
*,
crop_size,
resize_size,
mean=(0.43216, 0.394666, 0.37645),
std=(0.22803, 0.22145, 0.216989),
hflip_prob=0.5,
):
trans = [
ConvertBHWCtoBCHW(),
transforms.ConvertImageDtype(torch.float32),
transforms.Resize(resize_size),
]
if hflip_prob > 0:
trans.append(transforms.RandomHorizontalFlip(hflip_prob))
trans.extend([transforms.Normalize(mean=mean, std=std), transforms.RandomCrop(crop_size), ConvertBCHWtoCBHW()])
self.transforms = transforms.Compose(trans)
def build_data_loader(self):
if self.train_set:
transform = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.RandomVerticalFlip(),
transforms.ToTensor()
])
else:
transform = transforms.ToTensor()
cifar10_dtst = CIFAR10(root='data',
download=True,
train=self.train_set,
transform=transform)
cifar10_dl = data.DataLoader(cifar10_dtst,
batch_size=self.batch_size,
shuffle=True,
num_workers=self.num_workers)
return cifar10_dl
def __init__(self, data_path, mode):
if mode == "train" or mode == "val":
img_paths = sorted(glob.glob(os.path.join(data_path, '*.jpg')))
random.shuffle(img_paths)
else:
img_paths = sorted(glob.glob(os.path.join(data_path, '*.jpg')),
key=get_name)
split = int(len(img_paths) * 0.75)
if mode == "train":
self.img_paths = img_paths[:split]
elif mode == "val":
self.img_paths = img_paths[split:]
else:
self.img_paths = img_paths
if mode == "train" or mode == "val":
self.labels = [
path.split("/")[-1].split(".")[0] for path in self.img_paths
]
self.labels = [int(label == "cat") for label in self.labels]
else:
self.img_ids = [
path.split("/")[-1].split(".")[0] for path in self.img_paths
]
if mode == "train":
transforms = [
T.Grayscale(),
T.Resize((RE_HEIGHT, RE_WIDTH)),
T.RandomHorizontalFlip(),
T.ToTensor()
]
else: # "val" and "test"
transforms = [
T.Grayscale(),
T.Resize((RE_HEIGHT, RE_WIDTH)),
T.ToTensor()
]
self.transforms = T.Compose(transforms)
self.mode = mode
def train(args):
bsd300_train = load_bsd300('../data', split='train')
bsd300_test = load_bsd300('../data', split='test')
img_transforms = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.RandomRotation(15),
])
train_data = IsrDataset(bsd300_train,
output_size=200,
scale_factor=args.scale_factor,
deterministic=False,
base_image_transform=img_transforms,
transform=transforms.ToTensor(),
target_transform=transforms.ToTensor())
n_train = int(len(train_data) * 0.8)
split = [n_train, len(train_data) - n_train]
train_data, val_data = random_split(train_data, split)
test_data = IsrDataset(bsd300_test,
output_size=200,
scale_factor=args.scale_factor,
deterministic=True,
transform=transforms.ToTensor(),
target_transform=transforms.ToTensor())
model = SubPixelSrCnn(hparams=args)
trainer = Trainer()
trainer.fit(
model,
train_dataloader=DataLoader(test_data,
shuffle=True,
batch_size=32,
num_workers=2),
val_dataloaders=DataLoader(val_data,
shuffle=False,
batch_size=32,
num_workers=2),
)
trainer.test(model,
test_dataloaders=DataLoader(test_data,
shuffle=False,
batch_size=32,
num_workers=2))
def train_images():
from train.images.image_utils import params
data_args = params['data_args']
train_args = params['train_args']
model_args = params['model_args']
if params['use_wandb']:
wandb.login(key=os.environ['wanda_api_key'])
run_wandb = wandb.init(
project='dalle_train_vae',
job_type='train_model',
config=params,
resume=train_args['checkpoint_path'] is not None
)
else:
run_wandb = RunDisabled()
model = VqVae(**model_args).to('cuda')
print('num of trainable parameters: %d' % get_model_size(model))
print(model)
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
unnormalize = NormalizeInverse(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
training_data = ImagesDataset(
data_args['root_dir'],
transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize
])
)
training_loader = torch.utils.data.DataLoader(
training_data, batch_size=data_args['batch_size'], shuffle=True, num_workers=data_args['num_workers'])
train_object = TrainVqVae(model=model, training_loader=training_loader, run_wandb=run_wandb,
unnormalize=unnormalize,
**train_args)
try:
train_object.train()
finally:
run_wandb.finish()
def create_test_dataloaders(model_size, dataset_dir, batch_size):
normalize = transforms.Normalize(mean=(0.5, 0.5, 0.5),
std=(0.5, 0.5, 0.5))
train_transforms = transforms.Compose([
transforms.RandomResizedCrop(model_size),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])
dummy_config = type('dummy', (object,), {'dataset_dir': dataset_dir})()
train_dataset = create_cifar(dummy_config, dataset_config='cifar10', is_train=True, transform=train_transforms)
pin_memory = True
workers = 1
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=False, num_workers=workers,
pin_memory=pin_memory)
return train_loader, train_dataset
class Base:
base = ResNet
args = list()
kwargs = dict()
transform_train = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.Resize(32),
transforms.RandomCrop(32, padding=4),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010))
])
transform_test = transforms.Compose([
transforms.Resize(32),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010))
])
def load_data(batch_size=32, valid_batch_size=32):
train_transformations = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
test_transformations = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
kwargs_dl = {'root': '../data', 'download': True}
train_set = ds.CIFAR10(train=True, transform=train_transformations, **kwargs_dl)
test_set = ds.CIFAR10(train=False, transform=test_transformations, **kwargs_dl)
kwargs_train = {'shuffle': True, 'batch_size': batch_size, 'num_workers': NUM_WORKER}
kwargs_test = {'shuffle': True, 'batch_size': valid_batch_size, 'num_workers': NUM_WORKER}
train_set = DataLoader(train_set, **kwargs_train)
test_set = DataLoader(test_set, **kwargs_test)
return train_set, test_set
def load_CIFAR10(rate, batch_size=128):
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
trainset = torchvision.datasets.CIFAR10(root='.data/cifar10',
train=True,
download=True,
transform=transform_train)
L = len(trainset)
target_array = np.asarray(trainset.targets)
thr = 1 - rate * 10 / 9
target_array = np.where(
np.random.rand(L) <= thr, target_array, np.random.randint(0, 10, L))
trainset.targets = target_array.tolist()
idx = np.random.permutation(L)
train_idx = idx[:int(L * 0.8)]
val_idx = idx[int(L * 0.8):]
trainloader = torch.utils.data.DataLoader(
trainset,
batch_size=batch_size,
sampler=SubsetRandomSampler(train_idx))
valloader = torch.utils.data.DataLoader(
trainset, batch_size=batch_size, sampler=SubsetRandomSampler(val_idx))
testset = torchvision.datasets.CIFAR10(root='.data/cifar10',
train=False,
download=True,
transform=transform_test)
testloader = torch.utils.data.DataLoader(testset,
batch_size=batch_size,
shuffle=False)
return None, trainloader, valloader, testloader
def __init__(self, opt):
BaseDataset.__init__(self, opt)
self.shuffle = True if opt.isTrain else False
self.lr_size = opt.load_size // opt.scale_factor
self.hr_size = opt.load_size
self.img_dir = opt.dataroot
self.img_names = self.get_img_names()
self.aug = transforms.Compose([
transforms.RandomHorizontalFlip(),
Scale((1.0, 1.3), opt.load_size)
])
self.to_tensor = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
def data_transform(data_path, name, train=True):
with open(data_path + 'decathlon_mean_std.pickle', 'rb') as handle:
dict_mean_std = pickle._Unpickler(handle)
dict_mean_std.encoding = 'latin1'
dict_mean_std = dict_mean_std.load()
means = dict_mean_std[name + 'mean']
stds = dict_mean_std[name + 'std']
if name in ['gtsrb', 'omniglot', 'svhn']: # no horz flip
transform_train = transforms.Compose([
transforms.Resize(72),
transforms.CenterCrop(72),
transforms.ToTensor(),
transforms.Normalize(means, stds),
])
else:
transform_train = transforms.Compose([
transforms.Resize(72),
transforms.RandomCrop(72),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(means, stds),
])
if name in ['gtsrb', 'omniglot', 'svhn']: # no horz flip
transform_test = transforms.Compose([
transforms.Resize(72),
transforms.CenterCrop(72),
transforms.ToTensor(),
transforms.Normalize(means, stds),
])
else:
transform_test = transforms.Compose([
transforms.Resize(72),
transforms.CenterCrop(72),
transforms.ToTensor(),
transforms.Normalize(means, stds),
])
if train:
return transform_train
else:
return transform_test
def main():
start_time = time.time()
args = args_parser()
if not os.path.exists(args.save_dir):
os.makedirs(args.save_dir)
write_dict(vars(args), os.path.join(args.save_dir, 'arguments.csv'))
torch.manual_seed(args.seed)
cudnn.benchmark = True
torch.cuda.manual_seed_all(args.seed)
os.environ['CUDA_VISIBLE_DEVICES'] = args.device_ids
train_transformer = transforms.Compose([
transforms.Resize(size=(256, 256)),
transforms.RandomCrop(size=image_size),
transforms.RandomRotation(10),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
# transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
transforms.Lambda(normalize)
])
train_dataset = ChestXRay14Dataset(args.image_dir, args.train_file,
train_transformer)
val_transformer = transforms.Compose([
transforms.Resize(size=image_size),
transforms.ToTensor(),
# transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
transforms.Lambda(normalize)
])
val_dataset = ChestXRay14Dataset(args.image_dir, args.val_file,
val_transformer)
print(vars(args))
if args.model == 'simple':
model = simple_mlc_model(train_dataset.get_tag_size(),
backbone=args.backbone)
else:
model = mlc_model(train_dataset.get_tag_size(), backbone=args.backbone)
train(model, train_dataset, val_dataset, args)
def build_transforms(cfg, split="train"):
if split=="train":
min_size = min(cfg.EXTERNAL.IMAGE.HEIGHT,cfg.EXTERNAL.IMAGE.WIDTH)
max_size = max(cfg.EXTERNAL.IMAGE.HEIGHT,cfg.EXTERNAL.IMAGE.WIDTH)
flip_horizontal_prob = 0.5 # cfg.INPUT.FLIP_PROB_TRAIN
flip_vertical_prob = cfg.INPUT.VERTICAL_FLIP_PROB_TRAIN
brightness = cfg.INPUT.BRIGHTNESS
contrast = cfg.INPUT.CONTRAST
saturation = cfg.INPUT.SATURATION
hue = cfg.INPUT.HUE
else:
min_size = min(cfg.EXTERNAL.IMAGE.HEIGHT, cfg.EXTERNAL.IMAGE.WIDTH)
max_size = max(cfg.EXTERNAL.IMAGE.HEIGHT, cfg.EXTERNAL.IMAGE.WIDTH)
flip_horizontal_prob = 0.0
flip_vertical_prob = 0.0
brightness = 0.0
contrast = 0.0
saturation = 0.0
hue = 0.0
to_bgr255 = cfg.INPUT.TO_BGR255
normalize_transform = T.Normalize(
mean=cfg.INPUT.PIXEL_MEAN, std=cfg.INPUT.PIXEL_STD,
)
color_jitter = T.ColorJitter(
brightness=brightness,
contrast=contrast,
saturation=saturation,
hue=hue,
)
transform = T.Compose(
[
color_jitter,
T.Resize(min_size, max_size),
T.RandomHorizontalFlip(flip_horizontal_prob),
T.RandomVerticalFlip(flip_vertical_prob),
T.ToTensor(),
normalize_transform,
]
)
return transform
def load_cifar_data(self):
if self.params.transform_train:
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
self.normalize,
])
else:
transform_train = transforms.Compose([
transforms.ToTensor(),
self.normalize,
])
transform_test = transforms.Compose([
transforms.ToTensor(),
self.normalize,
])
self.train_dataset = torchvision.datasets.CIFAR10(
root=self.params.data_path,
train=True,
download=True,
transform=transform_train)
if self.params.poison_images:
self.train_loader = self.remove_semantic_backdoors()
else:
self.train_loader = DataLoader(self.train_dataset,
batch_size=self.params.batch_size,
shuffle=True,
num_workers=0)
self.test_dataset = torchvision.datasets.CIFAR10(
root=self.params.data_path,
train=False,
download=True,
transform=transform_test)
self.test_loader = DataLoader(self.test_dataset,
batch_size=self.params.test_batch_size,
shuffle=False,
num_workers=0)
self.classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog',
'horse', 'ship', 'truck')
return True
def predict_image(image_path):
print("prediciton in progress")
image = Image.open(image_path)
transformation = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
image_tensor = transformation(image).float()
image_tensor = image_tensor.unsqueeze_(0)
if torch.cuda.is_available():
image_tensor.cuda()
input = Variable(image_tensor)
output = model(input)
index = output.data.numpy().argmax()
return index
def get_transforms(self) -> tuple:
# TODO: Need to rethink the food101 transforms
MEAN = [0.5451, 0.4435, 0.3436]
STD = [0.2171, 0.2251,
0.2260] # TODO: should be [0.2517, 0.2521, 0.2573]
train_transf = [
transforms.Resize((32, 32)),
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip()
]
# food101 has images of varying sizes and are ~512 each side
test_transf = [transforms.Resize((32, 32))]
normalize = [transforms.ToTensor(), transforms.Normalize(MEAN, STD)]
train_transform = transforms.Compose(train_transf + normalize)
test_transform = transforms.Compose(test_transf + normalize)
return train_transform, test_transform
def gen_transform(input_size):
data_transforms = {
'train':
transforms.Compose([
transforms.ToPILImage(),
transforms.Resize(input_size),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]),
'val':
transforms.Compose([
transforms.ToPILImage(),
transforms.Resize(input_size),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]),
}
return data_transforms
def __init__(
self,
*,
crop_size,
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
interpolation=InterpolationMode.BILINEAR,
hflip_prob=0.5,
auto_augment_policy=None,
random_erase_prob=0.0,
):
trans = [
transforms.RandomResizedCrop(crop_size,
interpolation=interpolation)
]
if hflip_prob > 0:
trans.append(transforms.RandomHorizontalFlip(hflip_prob))
if auto_augment_policy is not None:
if auto_augment_policy == "ra":
trans.append(
autoaugment.RandAugment(interpolation=interpolation))
elif auto_augment_policy == "ta_wide":
trans.append(
autoaugment.TrivialAugmentWide(
interpolation=interpolation))
elif auto_augment_policy == "augmix":
trans.append(autoaugment.AugMix(interpolation=interpolation))
else:
aa_policy = autoaugment.AutoAugmentPolicy(auto_augment_policy)
trans.append(
autoaugment.AutoAugment(policy=aa_policy,
interpolation=interpolation))
trans.extend([
transforms.PILToTensor(),
transforms.ConvertImageDtype(torch.float),
transforms.Normalize(mean=mean, std=std),
])
if random_erase_prob > 0:
trans.append(transforms.RandomErasing(p=random_erase_prob))
self.transforms = transforms.Compose(trans)
def get_transforms(self) -> tuple:
MEAN = [0.485, 0.456, 0.406]
STD = [0.229, 0.224, 0.225]
_IMAGENET_PCA = {
'eigval': [0.2175, 0.0188, 0.0045],
'eigvec': [
[-0.5675, 0.7192, 0.4009],
[-0.5808, -0.0045, -0.8140],
[-0.5836, -0.6948, 0.4203],
]
}
transform_train, transform_test = None, None
transform_train = transforms.Compose([
transforms.RandomResizedCrop(
224,
scale=(
0.08,
1.0), # TODO: these two params are normally not specified
interpolation=Image.BICUBIC),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(brightness=0.4,
contrast=0.4,
saturation=0.4,
hue=0.2),
transforms.ToTensor(),
# TODO: Lighting is not used in original darts paper
# Lighting(0.1, _IMAGENET_PCA['eigval'], _IMAGENET_PCA['eigvec']),
transforms.Normalize(mean=MEAN, std=STD)
])
transform_test = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=MEAN, std=STD)
])
return transform_train, transform_test
def load_data(data_, labels_):
'''
Data loader
:param data_: imput images
:param labels_: labels of images
:return: train and test data transformed
'''
# Import data, scale them to values between [0,1] and cast them to the appropriate types
data = np.load(data_)
labels = np.load(labels_)
# Normalize each channel
img_mean = np.mean(np.swapaxes(data / 255.0, 0, 1).reshape(3, -1), 1)
img_std = np.std(np.swapaxes(data / 255.0, 0, 1).reshape(3, -1), 1)
normalize = transforms.Normalize(mean=list(img_mean), std=list(img_std))
x_train, x_test, y_train, y_test = train_test_split(data,
labels,
test_size=0.15,
random_state=42)
train = MyDataset(
x_train, y_train,
transforms.Compose([
transforms.ToPILImage(),
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(), normalize
]))
test = MyDataset(
x_test, y_test,
transforms.Compose([
transforms.ToPILImage(),
transforms.ToTensor(),
normalize,
]))
return train, test
def load_data(data_path, use_openfire=False, img_size=224, crop_pct=0.8):
# Data loading code
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_transforms = transforms.Compose([
transforms.RandomResizedCrop(size=img_size, scale=(crop_pct, 1.0)),
transforms.RandomRotation(degrees=5),
transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.1),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(), normalize
])
val_transforms = transforms.Compose([
transforms.Resize(size=(img_size)),
transforms.CenterCrop(size=img_size),
transforms.ToTensor(), normalize
])
print("Loading data")
if use_openfire:
train_set = OpenFire(root=data_path,
train=True,
download=True,
transform=train_transforms)
val_set = OpenFire(root=data_path,
train=False,
download=True,
transform=val_transforms)
else:
train_dir = os.path.join(data_path, 'train')
val_dir = os.path.join(data_path, 'val')
train_set = ImageFolder(train_dir,
train_transforms,
target_transform=target_transform)
val_set = ImageFolder(val_dir,
val_transforms,
target_transform=target_transform)
return train_set, val_set
def __init__(self, imsize, data_path, startidx=0):
super(Dataset64, self).__init__()
self.imsize = imsize
self.data_path = data_path
self.startidx = startidx
self.transform = transforms.Compose([lambda x: Image.open(x).convert('RGB'),
transforms.RandomResizedCrop(imsize),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406),
(0.229, 0.224, 0.225))])
self.path_images = os.path.join(data_path, 'miniimagenet', 'images') # image path
csvdata = self.loadCSV(os.path.join(data_path, 'miniimagenet', 'train.csv')) # csv path
self.data = []
self.img2label = {}
for i, (k, v) in enumerate(csvdata.items()):
self.data.extend(v)
self.img2label[k] = i + self.startidx
self.num_classes = len(self.img2label)
def __getitem__(self, index):
rgb = np.array(self.f["rgb"][index])
label = np.array((self.f["labels"][index] - self.f["Mean"])/self.f["Variance"])
t_rgb = torch.zeros(3, 224, 224)
prob = random.uniform(0, 1)
prob2 = random.uniform(0, 1)
if self.transform is not None:
if (prob > 0.5 and not self.test):
flip_transform = transforms.Compose([transforms.ToPILImage(), transforms.RandomHorizontalFlip(1.0)])
rgb[:,:,:] = flip_transform(rgb[:,:,:])
if (prob2 > 0.5 and not self.test):
color_jitter_transform = transforms.Compose([transforms.ToPILImage() ,transforms.ColorJitter(brightness = 0.5, contrast = 0.5, saturation = 0.5, hue = 0.2)])
rgb[:,:,:] = color_jitter_transform(rgb[:,:,:])
t_rgb[:,:,:] = self.transform(rgb[:,:,:])
return t_rgb, label
def get_train_data_loader(dataset_name: str, batch_size: int, num_workers: int,
resize_dim: int) -> DataLoader:
with DatasetTimer():
dataset = ImageFolder(root=os.path.join("/datasets", dataset_name,
"train"),
transform=transforms.Compose([
transforms.RandomResizedCrop(resize_dim),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=IMAGENET_MEAN,
std=IMAGENET_STD),
]))
return DataLoader(
dataset=dataset,
batch_size=batch_size,
sampler=RandomSampler(data_source=dataset),
num_workers=num_workers,
collate_fn=collate_fn,
drop_last=True,
)
def __init__(self, config: BaseConfig): super().__init__(config) self.config = config # self.img_dir = os.path.join(self.config.dataset_config.data_root_dir, self.config.dataset_config.img_dir) train_attr, test_attr, self.label = get_attr_binary(self.config.dataset_config) train, val, test = get_img(self.config.dataset_config) if config.isTrain: # distribution:每个属性的正样本占比 self.img = train['data'] self.img_ids = train['ids'] self.img_attr = train_attr self.transforms = transforms.Compose([ transforms.Resize(size=(288, 144)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) else: self.transforms = transforms.Compose([ transforms.Resize(size=(288, 144)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) self.img_attr = test_attr if self.config.isTest: # test set self.img = test['data'] self.img_ids = test['ids'] else: # val set self.img = val['data'] self.img_ids = val['ids'] self.length = len(self.img)
