def _augment(self, batch_image, batch_target):
"""
if self._enable_augmentation:
rand_args = [i for i in range(len(batch_image))]
random.shuffle(rand_args)
for i in range(len(batch_image)):
image, image_mix = batch_image[i], batch_image[rand_args[i]]
image = cv2.resize(image, self._image_size)
image_mix = cv2.resize(image_mix, self._image_size)
target, target_mix = batch_target[i], batch_target[rand_args[i]]
image, target = AUG.random_mixup(image, image_mix,
target, target_mix)
batch_image[i] = image
batch_target[i] = target
"""
output_batch_image = []
output_batch_target = []
for image, target in zip(batch_image, batch_target):
image = cv2.resize(image, self._image_size)
if self._enable_augmentation:
image = AUG.random_shift(image)
image = AUG.random_rotate(image)
image = AUG.random_cutout(image)
output_batch_image.append(image)
output_batch_target.append(target)
return output_batch_image, output_batch_target
def generate_data():
model.train()
global variations_per_image
global surrogate_classes
images_to_process = surrogate_classes / variations_per_image
images_processed = 0
augmented_images = []
surrogate_labels = []
for batch_idx, (data, _) in enumerate(train_unlabeled_loader):
for i in range(len(data)):
images_processed += 1
if images_processed > images_to_process:
break
new_data = []
for j in range(variations_per_image):
new_image = transforms.Compose([
transforms.Lambda(lambda x: augmentations.to_npimg(x)),
transforms.Lambda(lambda x: augmentations.random_translate(x)),
transforms.Lambda(lambda x: augmentations.random_rotate(x)),
transforms.Lambda(lambda x: augmentations.random_scale(x)),
transforms.Lambda(lambda x: augmentations.to_tensor(x)),
])(data[i])
c, h, w = new_image.size()
new_image.resize_(1, c, h, w)
new_data.append(new_image)
surrogate_labels.append(images_processed)
augmented_images.append(torch.cat(new_data, 0))
augmented_images_tensor = torch.cat(augmented_images, 0)
surrogate_labels_tensor = torch.LongTensor(surrogate_labels)
return (augmented_images_tensor, surrogate_labels_tensor)
def __getitem__(self, index):
HR_path, LR_path = None, None
scale = self.opt['scale']
HR_size = self.opt['HR_size']
#v
if self.opt['rand_flip_LR_HR'] and self.LR_scale and self.opt['phase'] == 'train':
LRHRchance = random.uniform(0, 1)
if self.opt['flip_chance']:
flip_chance = self.opt['flip_chance']
else:
flip_chance = 0.05
#print("Random Flip Enabled")
else:
LRHRchance = 0.
flip_chance = 0.
#print("No Random Flip")
# get HR image
if LRHRchance < (1- flip_chance):
HR_path = self.paths_HR[index]
#print("HR kept")
else:
HR_path = self.paths_LR[index]
#print("HR flipped")
#v
img_HR = util.read_img(self.HR_env, HR_path)
# modcrop in the validation / test phase
if self.opt['phase'] != 'train':
img_HR = util.modcrop(img_HR, scale)
# change color space if necessary
if self.opt['color']:
img_HR = util.channel_convert(img_HR.shape[2], self.opt['color'], [img_HR])[0]
#v
if self.HR_crop and (self.HR_rrot != True):
crop_size = (HR_size, HR_size)
img_HR, _ = augmentations.random_resize_img(img_HR, crop_size)
elif self.HR_rrot and (self.HR_crop != True):
img_HR, _ = augmentations.random_rotate(img_HR)
elif self.HR_crop and self.HR_rrot:
if np.random.rand() > 0.5:
crop_size = (HR_size, HR_size)
img_HR, _ = augmentations.random_resize_img(img_HR, crop_size)
else:
img_HR, _ = augmentations.random_rotate(img_HR)
#v
#v
if self.HR_noise:
img_HR, hr_noise_algo = augmentations.noise_img(img_HR, self.hr_noise_types)
#v
# get LR image
if self.paths_LR:
if self.HR_crop or self.HR_rrot: #v
img_LR = img_HR
else:
if LRHRchance < (1- flip_chance):
LR_path = self.paths_LR[index]
#print("LR kept")
else:
LR_path = self.paths_HR[index]
#print("LR flipped")
img_LR = util.read_img(self.LR_env, LR_path)
#"""
#v scale
if self.LR_scale:
img_LR, scale_interpol_algo = augmentations.scale_img(img_LR, scale)
#"""
#"""
#v blur
if self.LR_blur:
img_LR, blur_algo, blur_kernel_size = augmentations.blur_img(img_LR, self.blur_algos)
#"""
#"""
#v noise
if self.LR_noise:
img_LR, noise_algo = augmentations.noise_img(img_LR, self.noise_types)
if self.LR_noise2:
img_LR, noise_algo2 = augmentations.noise_img(img_LR, self.noise_types2)
#"""
#"""
#v LR cutout / LR random erasing
if self.LR_cutout and (self.LR_erasing != True):
img_LR = augmentations.cutout(img_LR, img_LR.shape[0] // 2)
elif self.LR_erasing and (self.LR_cutout != True): #only do cutout or erasing, not both
img_LR = augmentations.random_erasing(img_LR)
elif self.LR_cutout and self.LR_erasing:
if np.random.rand() > 0.5:
img_LR = augmentations.cutout(img_LR, img_LR.shape[0] // 2, p=0.5)
else:
img_LR = augmentations.random_erasing(img_LR, p=0.5, modes=[3])
#"""
else: # down-sampling on-the-fly
# randomly scale during training
if self.opt['phase'] == 'train':
random_scale = random.choice(self.random_scale_list)
H_s, W_s, _ = img_HR.shape
def _mod(n, random_scale, scale, thres):
rlt = int(n * random_scale)
rlt = (rlt // scale) * scale
return thres if rlt < thres else rlt
H_s = _mod(H_s, random_scale, scale, HR_size)
W_s = _mod(W_s, random_scale, scale, HR_size)
img_HR = cv2.resize(np.copy(img_HR), (W_s, H_s), interpolation=cv2.INTER_LINEAR)
# force to 3 channels
if img_HR.ndim == 2:
img_HR = cv2.cvtColor(img_HR, cv2.COLOR_GRAY2BGR)
H, W, _ = img_HR.shape
# using matlab imresize
img_LR = util.imresize_np(img_HR, 1 / scale, True)
if img_LR.ndim == 2:
img_LR = np.expand_dims(img_LR, axis=2)
if self.opt['phase'] == 'train':
# if the image size is too small
H, W, _ = img_HR.shape
if H < HR_size or W < HR_size:
img_HR = cv2.resize(
np.copy(img_HR), (HR_size, HR_size), interpolation=cv2.INTER_LINEAR)
# using matlab imresize
img_LR = util.imresize_np(img_HR, 1 / scale, True)
if img_LR.ndim == 2:
img_LR = np.expand_dims(img_LR, axis=2)
H, W, C = img_LR.shape
LR_size = HR_size // scale
# randomly crop
rnd_h = random.randint(0, max(0, H - LR_size))
rnd_w = random.randint(0, max(0, W - LR_size))
img_LR = img_LR[rnd_h:rnd_h + LR_size, rnd_w:rnd_w + LR_size, :]
rnd_h_HR, rnd_w_HR = int(rnd_h * scale), int(rnd_w * scale)
img_HR = img_HR[rnd_h_HR:rnd_h_HR + HR_size, rnd_w_HR:rnd_w_HR + HR_size, :]
# augmentation - flip, rotate
img_LR, img_HR = util.augment([img_LR, img_HR], self.opt['use_flip'], \
self.opt['use_rot'])
# change color space if necessary
if self.opt['color']:
#img_LR = util.channel_convert(C, self.opt['color'], [img_LR])[0] # TODO during val no definetion
img_LR = util.channel_convert(img_LR.shape[2], self.opt['color'], [img_LR])[0] # v appears to work ok
# BGR to RGB, HWC to CHW, numpy to tensor
if img_HR.shape[2] == 3:
img_HR = img_HR[:, :, [2, 1, 0]]
img_LR = img_LR[:, :, [2, 1, 0]]
img_HR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_HR, (2, 0, 1)))).float()
img_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()
if LR_path is None:
LR_path = HR_path
return {'LR': img_LR, 'HR': img_HR, 'LR_path': LR_path, 'HR_path': HR_path}