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}
def __getitem__(self, index):
HR_path, LR_path = None, None
scale = self.opt['scale']
HR_size = self.opt['HR_size']
if HR_size:
LR_size = HR_size // scale
self.znorm = False # Default case: images are in the [0,1] range
if self.opt['znorm']:
if self.opt['znorm'] == True:
# Alternative: images are z-normalized to the [-1,1] range
self.znorm = True
######## Read the images ########
# Check if LR Path is provided
if self.paths_LR:
# If LR is provided, check if 'rand_flip_LR_HR' is enabled
if self.opt['rand_flip_LR_HR'] 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")
# Normal case, no flipping:
else:
LRHRchance = 0.
flip_chance = 0.
#print("No Random Flip")
# get HR and LR images
# If enabled, random chance that LR and HR images are flipped
# Normal case, no flipping
# If img_LR (LR_path) doesn't exist, use img_HR (HR_path)
if LRHRchance < (1 - flip_chance):
HR_path = self.paths_HR[index]
LR_path = self.paths_LR[index]
if LR_path is None:
LR_path = HR_path
#print("HR kept")
# Flipped case:
# If img_HR (LR_path) doesn't exist, use img_HR (LR_path)
else:
HR_path = self.paths_LR[index]
LR_path = self.paths_HR[index]
if HR_path is None:
HR_path = LR_path
#print("HR flipped")
# Read the LR and HR images from the provided paths
img_LR = util.read_img(self.LR_env, LR_path, znorm=self.znorm)
img_HR = util.read_img(self.HR_env, HR_path, znorm=self.znorm)
# Even if LR dataset is provided, force to generate aug_downscale % of downscales OTF from HR
# The code will later make sure img_LR has the correct size
if self.opt['aug_downscale']:
aug_downscale = self.opt['aug_downscale']
if np.random.rand() < aug_downscale:
img_LR = img_HR
# If LR is not provided, use HR and modify on the fly
else:
HR_path = self.paths_HR[index]
img_HR = util.read_img(self.HR_env, HR_path, znorm=self.znorm)
img_LR = img_HR
######## Modify the images ########
# HR modcrop in the validation / test phase
if self.opt['phase'] != 'train':
img_HR = util.modcrop(img_HR, scale)
# change color space if necessary
# Note: Changing the LR colorspace here could make it so some colors are introduced when
# doing the augmentations later (ie: with Gaussian or Speckle noise), may be good if the
# model can learn to remove color noise in grayscale images, otherwise move to before
# converting to tensors
# self.opt['color'] For both LR and HR as in the the original code, kept for compatibility
# self.opt['color_HR'] and self.opt['color_LR'] for independent control
if self.opt['color_HR'] or self.opt['color']: # Only change HR
img_HR = util.channel_convert(img_HR.shape[2], self.opt['color'],
[img_HR])[0]
if self.opt['color_LR'] or self.opt['color']: # Only change LR
img_LR = util.channel_convert(img_LR.shape[2], self.opt['color'],
[img_LR])[0]
######## Augmentations ########
# Augmentations during training
if self.opt['phase'] == 'train':
# Validate there's an img_LR, if not, use img_HR
if img_LR is None:
img_LR = img_HR
print("Image LR: ", LR_path, (
"was not loaded correctly, using HR pair to downscale on the fly."
))
# Check that HR and LR have the same dimensions ratio, else, generate new LR from HR
if img_HR.shape[0] // img_LR.shape[0] != img_HR.shape[
1] // img_LR.shape[1]:
print(
"Warning: img_LR dimensions ratio does not match img_HR dimensions ratio for: ",
HR_path)
img_LR = img_HR
# Random Crop (reduce computing cost and adjust images to correct size first)
if img_HR.shape[0] > HR_size or img_HR.shape[1] > HR_size:
# Here the scale should be in respect to the images, not to the training scale (in case they are being scaled on the fly)
scaleor = img_HR.shape[0] // img_LR.shape[0]
img_HR, img_LR = augmentations.random_crop_pairs(
img_HR, img_LR, HR_size, scaleor)
# Or if the HR images are too small, Resize to the HR_size size and fit LR pair to LR_size too
if img_HR.shape[0] < HR_size or img_HR.shape[1] < HR_size:
print("Warning: Image: ", HR_path,
" size does not match HR size: (", HR_size,
"). The image size is: ", img_HR.shape)
# rescale HR image to the HR_size
img_HR, _ = augmentations.resize_img(np.copy(img_HR),
crop_size=(HR_size,
HR_size),
algo=cv2.INTER_LINEAR)
# rescale LR image to the LR_size (The original code discarded the img_LR and generated a new one on the fly from img_HR)
img_LR, _ = augmentations.resize_img(np.copy(img_LR),
crop_size=(LR_size,
LR_size),
algo=cv2.INTER_LINEAR)
# Randomly scale LR from HR during training if :
# - LR dataset is not provided
# - LR dataset is not in the correct scale
# - Also to check if LR is not at the correct scale already (if img_LR was changed to img_HR)
if img_LR.shape[0] != LR_size or img_LR.shape[1] != LR_size:
ds_algo = 777 # default to matlab-like bicubic downscale
# if manually set and scale algorithms are provided, then:
if self.opt['lr_downscale']:
if self.opt['lr_downscale_types']:
ds_algo = self.opt['lr_downscale_types']
else: # else, if for some reason img_LR is too large, default to matlab-like bicubic downscale
# if not self.opt['aug_downscale']: #only print the warning if not being forced to use HR images instead of LR dataset (which is a known case)
print(
"LR image is too large, auto generating new LR for: ",
LR_path)
img_LR, scale_interpol_algo = augmentations.scale_img(
img_LR, scale, algo=ds_algo)
if self.znorm:
# The generated LR sometimes get slightly out of the [-1,1] range
np.clip(img_LR, -1., 1., out=img_LR)
else:
# The generated LR sometimes get slightly out of the [0,1] range
np.clip(img_LR, 0., 1., out=img_LR)
# """
# Rotations. 'use_flip' = 180 or 270 degrees (mirror), 'use_rot' = 90 degrees, 'HR_rrot' = random rotations +-45 degrees
if (self.opt['use_flip']
or self.opt['use_rot']) and self.opt['hr_rrot']:
if np.random.rand() > 0.5:
img_LR, img_HR = util.augment([img_LR, img_HR],
self.opt['use_flip'],
self.opt['use_rot'])
else:
if np.random.rand(
) > 0.5: # randomize the random rotations, so half the images are the original
img_HR, img_LR = augmentations.random_rotate_pairs(
img_HR, img_LR, HR_size, scale)
elif (self.opt['use_flip']
or self.opt['use_rot']) and not self.opt['hr_rrot']:
# augmentation - flip, rotate
img_LR, img_HR = util.augment([img_LR, img_HR],
self.opt['use_flip'],
self.opt['use_rot'])
elif self.opt['hr_rrot']:
if np.random.rand(
) > 0.5: # randomize the random rotations, so half the images are the original
img_HR, img_LR = augmentations.random_rotate_pairs(
img_HR, img_LR, HR_size, scale)
# Final checks
# if the resulting HR image size so far is too large or too small, resize HR to the correct size and downscale to generate a new LR on the fly
if img_HR.shape[0] != HR_size or img_HR.shape[1] != HR_size:
print("Image: ", HR_path, " size does not match HR size: (",
HR_size, "). The image size is: ", img_HR.shape)
# rescale HR image to the HR_size
img_HR, _ = augmentations.resize_img(np.copy(img_HR),
crop_size=(HR_size,
HR_size),
algo=cv2.INTER_LINEAR)
# if manually provided and scale algorithms are provided, then:
if self.opt['lr_downscale_types']:
ds_algo = self.opt['lr_downscale_types']
else:
# using matlab imresize to generate LR pair
ds_algo = 777
img_LR, _ = augmentations.scale_img(img_HR,
scale,
algo=ds_algo)
# if the resulting LR so far does not have the correct dimensions, also generate a new HR-LR image pair on the fly
if img_LR.shape[0] != LR_size or img_LR.shape[0] != LR_size:
print("Image: ", LR_path, " size does not match LR size: (",
HR_size // scale, "). The image size is: ", img_LR.shape)
# rescale HR image to the HR_size (should not be needed, but something went wrong before, just for sanity)
img_HR, _ = augmentations.resize_img(np.copy(img_HR),
crop_size=(HR_size,
HR_size),
algo=cv2.INTER_LINEAR)
# if manually provided and scale algorithms are provided, then:
if self.opt['lr_downscale_types']:
ds_algo = self.opt['lr_downscale_types']
else:
# using matlab imresize to generate LR pair
ds_algo = 777
img_LR, _ = augmentations.scale_img(img_HR,
scale,
algo=ds_algo)
# Below are the LR On The Fly augmentations
# Add noise to HR if enabled AND noise types are provided (for noise2noise and similar)
if self.opt['hr_noise']:
if self.opt['hr_noise_types']:
img_HR, hr_noise_algo = augmentations.noise_img(
img_HR, noise_types=self.opt['hr_noise_types'])
else:
print(
"Noise types 'hr_noise_types' not defined. Skipping OTF noise for HR."
)
# Create color fringes
# Caution: Can easily destabilize a model
# Only applied to a small % of the images. Around 20% and 50% appears to be stable.
if self.opt['lr_fringes']:
lr_fringes_chance = self.opt['lr_fringes_chance'] if self.opt[
'lr_fringes_chance'] else 0.4
if np.random.rand() > (1. - lr_fringes_chance):
img_LR = augmentations.translate_chan(img_LR)
# """
# v LR blur AND blur types are provided, else will skip
if self.opt['lr_blur']:
if self.opt['lr_blur_types']:
img_LR, blur_algo, blur_kernel_size = augmentations.blur_img(
img_LR, blur_algos=self.opt['lr_blur_types'])
else:
print(
"Blur types 'lr_blur_types' not defined. Skipping OTF blur."
)
# """
# """
# v LR primary noise: Add noise to LR if enabled AND noise types are provided, else will skip
if self.opt['lr_noise']:
if self.opt['lr_noise_types']:
img_LR, noise_algo = augmentations.noise_img(
img_LR, noise_types=self.opt['lr_noise_types'])
else:
print(
"Noise types 'lr_noise_types' not defined. Skipping OTF noise."
)
# v LR secondary noise: Add additional noise to LR if enabled AND noise types are provided, else will skip
if self.opt['lr_noise2']:
if self.opt['lr_noise_types2']:
img_LR, noise_algo2 = augmentations.noise_img(
img_LR, noise_types=self.opt['lr_noise_types2'])
else:
print(
"Noise types 'lr_noise_types2' not defined. Skipping OTF secondary noise."
)
# """
# """
# v LR cutout / LR random erasing (for inpainting/classification tests)
if self.opt['lr_cutout'] and (self.opt['lr_erasing'] != True):
img_LR = augmentations.cutout(img_LR, img_LR.shape[0] // 2)
elif self.opt['lr_erasing'] and (self.opt['lr_cutout'] != True):
img_LR = augmentations.random_erasing(img_LR)
elif self.opt['lr_cutout'] and self.opt['lr_erasing']:
if np.random.rand(
) > 0.5: # only do cutout or erasing, not both at the same time
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])
# """
# Apply "auto levels" to images
# Randomize for augmentation
rand_levels = (1 - self.opt['rand_auto_levels']
) if self.opt['rand_auto_levels'] else 1
if self.opt['auto_levels'] and np.random.rand() > rand_levels:
if self.opt['auto_levels'] == 'HR':
img_HR = augmentations.simplest_cb(img_HR,
znorm=self.znorm)
elif self.opt['auto_levels'] == 'LR':
img_LR = augmentations.simplest_cb(img_LR,
znorm=self.znorm)
elif self.opt['auto_levels'] == True or self.opt[
'auto_levels'] == 'Both':
img_HR = augmentations.simplest_cb(img_HR,
znorm=self.znorm)
img_LR = augmentations.simplest_cb(img_LR,
znorm=self.znorm)
# Apply unsharpening mask to HR images
# img_HR1 = img_HR
# Randomize for augmentation
rand_unsharp = (
1 -
self.opt['rand_unsharp']) if self.opt['rand_unsharp'] else 1
if self.opt['unsharp_mask'] and np.random.rand() > rand_unsharp:
img_HR = augmentations.unsharp_mask(img_HR, znorm=self.znorm)
# For testing and validation
if self.opt['phase'] != 'train':
# Randomly downscale LR if enabled
if self.opt['lr_downscale']:
if self.opt['lr_downscale_types']:
img_LR, scale_interpol_algo = augmentations.scale_img(
img_LR, scale, algo=self.opt['lr_downscale_types'])
else: # Default to matlab-like bicubic downscale
img_LR, scale_interpol_algo = augmentations.scale_img(
img_LR, scale, algo=777)
# Alternative position for changing the colorspace of LR.
# if self.opt['color_LR']: # Only change LR
# img_LR = util.channel_convert(img_LR.shape[2], self.opt['color'], [img_LR])[0]
# Debug
# Save img_LR and img_HR images to a directory to visualize what is the result of the on the fly augmentations
# DO NOT LEAVE ON DURING REAL TRAINING
#self.output_sample_imgs = True
if self.opt['phase'] == 'train':
if self.output_sample_imgs:
import os
# LR_dir, im_name = os.path.split(LR_path)
HR_dir, im_name = os.path.split(HR_path)
baseHRdir, _ = os.path.split(HR_dir)
debugpath = os.path.join(baseHRdir, 'sampleOTFimgs')
print('debugpathhhhhhhh', debugpath)
# debugpath = os.path.join(os.path.split(LR_dir)[0], 'sampleOTFimgs')
# debugpath = os.path.join('D:/tmp_test', 'sampleOTFimgs')
# print(debugpath)
if not os.path.exists(debugpath):
os.makedirs(debugpath)
if self.opt[
'znorm']: # Back from [-1,1] range to [0,1] range for OpenCV2
img_LRn = (img_LR + 1.0) / 2.0
img_HRn = (img_HR + 1.0) / 2.0
# img_HRn1 = (img_HR1 + 1.0) / 2.0
else: # Already in the [0,1] range for OpenCV2
img_LRn = img_LR
img_HRn = img_HR
# img_HRn1 = img_HR1
import uuid
hex = uuid.uuid4().hex
# random name to save + had to multiply by 255, else getting all black image
cv2.imwrite(os.path.join(debugpath, im_name + hex + '_LR.png'),
img_LRn * 255)
# random name to save + had to multiply by 255, else getting all black image
cv2.imwrite(os.path.join(debugpath, im_name + hex + '_HR.png'),
img_HRn * 255)
# cv2.imwrite(debugpath+"\\"+im_name+hex+'_HR1.png',img_HRn1*255) #random name to save + had to multiply by 255, else getting all black image
######## Convert images to PyTorch Tensors ########
"""
if (img_HR.min() < -1):
print("HR.min :", img_HR.min())
print(HR_path)
if (img_HR.max() > 1):
print("HR.max :", img_HR.max())
print(HR_path)
if (img_LR.min() < -1):
print("LR.min :", img_LR.min())
print(LR_path)
if (img_LR.max() > 1):
print("LR.max :", img_LR.max())
print(LR_path)
#"""
# 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]]
# BGRA to RGBA, HWC to CHW, numpy to tensor
elif img_LR.shape[2] == 4:
img_HR = img_HR[:, :, [2, 1, 0, 3]]
img_LR = img_LR[:, :, [2, 1, 0, 3]]
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
}
def __getitem__(self, index):
HR_path, LR_path = None, None
scale = self.opt['scale']
HR_size = self.opt['HR_size']
if HR_size:
LR_size = HR_size // scale
# Check if LR Path is provided
if self.paths_LR:
#If LR is provided, check if 'rand_flip_LR_HR' is enabled (Only will work if HR and LR images have the same initial size) during training
if self.opt['rand_flip_LR_HR'] and (
self.LR_scale
or scale == 1) 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")
# Normal case, no flipping:
else:
LRHRchance = 0.
flip_chance = 0.
#print("No Random Flip")
# get HR and LR images
# If enabled, random chance that LR and HR images are flipped
# Normal case, no flipping
# If img_LR (LR_path) doesn't exist, use img_HR (HR_path)
if LRHRchance < (1 - flip_chance):
HR_path = self.paths_HR[index]
LR_path = self.paths_LR[index]
if LR_path is None:
LR_path = HR_path
#print("HR kept")
# Flipped case:
# If img_HR (LR_path) doesn't exist, use img_HR (LR_path)
else:
HR_path = self.paths_LR[index]
LR_path = self.paths_HR[index]
if HR_path is None:
HR_path = LR_path
#print("HR flipped")
# Read the LR and HR images from the provided paths
img_LR = util.read_img(self.LR_env, LR_path)
img_HR = util.read_img(self.HR_env, HR_path)
# Even if LR dataset is provided, force to generate aug_downscale % of downscales OTF from HR
# The code will later make sure img_LR has the correct size
if self.opt['aug_downscale']:
aug_downscale = self.opt['aug_downscale']
if np.random.rand() < aug_downscale:
img_LR = img_HR
# If LR is not provided, use HR and modify on the fly
else:
HR_path = self.paths_HR[index]
img_HR = util.read_img(self.HR_env, HR_path)
img_LR = img_HR
# HR 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]
img_LR = util.channel_convert(img_LR.shape[2], self.opt['color'],
[img_LR])[0]
#Augmentations during training
if self.opt['phase'] == 'train':
#Validate there's an img_LR,
if img_LR is None:
img_LR = img_HR
print("Image LR: ", LR_path, (
"was not loaded correctly, using HR pair to downscale on the fly."
))
#Random Crop (reduce computing cost and adjust images to correct size first)
if img_HR.shape[0] > HR_size or img_HR.shape[1] > HR_size:
#Here the scale should be in respect to the images, not to the training scale (in case they are being scaled on the fly)
if img_HR.shape[0] // img_LR.shape[0] is not img_HR.shape[
1] // img_LR.shape[1]:
print(
"Warning: img_LR dimensions ratio does not match img_HR dimensions ratio for: ",
HR_path)
img_LR = img_HR
scaleor = img_HR.shape[0] // img_LR.shape[0]
img_HR, img_LR = augmentations.random_crop_pairs(
img_HR, img_LR, HR_size, scaleor)
#or if the HR images are too small, Resize to the HR_size size and fit LR pair to LR_size too
if img_HR.shape[0] < HR_size or img_HR.shape[1] < HR_size:
print("Warning: Image: ", HR_path,
" size does not match HR size: (", HR_size,
"). The image size is: ", img_HR.shape)
# rescale HR image to the HR_size
img_HR, _ = augmentations.resize_img(np.copy(img_HR),
crop_size=(HR_size,
HR_size),
algo=cv2.INTER_LINEAR)
# rescale LR image to the LR_size
img_LR, _ = augmentations.resize_img(np.copy(img_LR),
crop_size=(LR_size,
LR_size),
algo=cv2.INTER_LINEAR)
#"""
# randomly scale LR from HR during training if LR dataset is not provided
# Also check if LR is not at the correct scale already
if img_LR.shape[0] is not LR_size and img_LR.shape[
1] is not LR_size:
if self.LR_scale: # if manually provided and scale algorithms are provided, then:
if self.scale_algos:
ds_algo = self.scale_algos
else:
ds_algo = 777
else: # else, if for some reason img_LR is too large, default to matlab-like bicubic downscale
#if not self.opt['aug_downscale']: #only print the warning if not being forced to use HR images instead of LR dataset (which is a known case)
ds_algo = 777
print(
"LR image is too large, auto generating new LR for: ",
LR_path)
img_LR, scale_interpol_algo = augmentations.scale_img(
img_LR, scale, algo=ds_algo)
#"""
# Rotations. 'use_flip' = 180 or 270 degrees (mirror), 'use_rot' = 90 degrees, 'HR_rrot' = random rotations +-45 degrees
if self.opt['use_flip'] and self.opt['use_rot'] and self.HR_rrot:
if np.random.rand() > 0.5:
img_LR, img_HR = util.augment([img_LR, img_HR], self.opt['use_flip'], \
self.opt['use_rot'])
else:
if np.random.rand(
) > 0.5: # randomize the random rotations, so half the images are the original
img_HR, img_LR = augmentations.random_rotate_pairs(
img_HR, img_LR, HR_size, scale)
elif (self.opt['use_flip']
or self.opt['use_rot']) and not self.HR_rrot:
# augmentation - flip, rotate
img_LR, img_HR = util.augment([img_LR, img_HR], self.opt['use_flip'], \
self.opt['use_rot'])
elif self.HR_rrot:
if np.random.rand(
) > 0.5: # randomize the random rotations, so half the images are the original
img_HR, img_LR = augmentations.random_rotate_pairs(
img_HR, img_LR, HR_size, scale)
# Final checks
# if the resulting HR image size so far is too large or too small, resize HR to the correct size and downscale to generate a new LR on the fly
if img_HR.shape[0] is not HR_size or img_HR.shape[1] is not HR_size:
print("Image: ", HR_path, " size does not match HR size: (",
HR_size, "). The image size is: ", img_HR.shape)
# rescale HR image to the HR_size
img_HR, _ = augmentations.resize_img(np.copy(img_HR),
crop_size=(HR_size,
HR_size),
algo=cv2.INTER_LINEAR)
if self.scale_algos: # if manually provided and scale algorithms are provided, then:
ds_algo = self.scale_algos
else:
## using matlab imresize to generate LR pair
ds_algo = 777
img_LR, _ = augmentations.scale_img(img_HR,
scale,
algo=ds_algo)
# Final checks
# if the resulting LR so far does not have the correct dimensions, also generate a new HR- LR image pair on the fly
if img_LR.shape[0] is not LR_size or img_LR.shape[0] is not LR_size:
print("Image: ", LR_path, " size does not match LR size: (",
HR_size // scale, "). The image size is: ", img_LR.shape)
# rescale HR image to the HR_size
img_HR, _ = augmentations.resize_img(np.copy(img_HR),
crop_size=(HR_size,
HR_size),
algo=cv2.INTER_LINEAR)
if self.scale_algos: # if manually provided and scale algorithms are provided, then:
ds_algo = self.scale_algos
else:
## using matlab imresize to generate LR pair
ds_algo = 777
img_LR, _ = augmentations.scale_img(img_HR,
scale,
algo=ds_algo)
# Add noise to HR if enabled
if self.HR_noise:
img_HR, hr_noise_algo = augmentations.noise_img(
img_HR, noise_types=self.hr_noise_types)
# Below are the LR On The Fly augmentations
#"""
#v LR blur
if self.LR_blur:
img_LR, blur_algo, blur_kernel_size = augmentations.blur_img(
img_LR, blur_algos=self.blur_algos)
#"""
#"""
#v LR primary noise
if self.LR_noise:
img_LR, noise_algo = augmentations.noise_img(
img_LR, noise_types=self.noise_types)
#v LR secondary noise
if self.LR_noise2:
img_LR, noise_algo2 = augmentations.noise_img(
img_LR, noise_types=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])
#"""
#For testing and validation
if self.opt['phase'] != 'train':
#"""
#v randomly downscale LR if enabled
if self.LR_scale:
img_LR, scale_interpol_algo = augmentations.scale_img(
img_LR, scale, algo=self.scale_algos)
#"""
# Debug
# Save img_LR and img_HR images to a directory to visualize what is the result of the on the fly augmentations
# DO NOT LEAVE ON DURING REAL TRAINING
# self.output_sample_imgs = True
if self.opt['phase'] == 'train':
if self.output_sample_imgs:
import os
LR_dir, im_name = os.path.split(LR_path)
#baseHRdir, _ = os.path.split(HR_dir)
#debugpath = os.path.join(baseHRdir, os.sep, 'sampleOTFimgs')
debugpath = os.path.join(
os.path.split(LR_dir)[0], 'sampleOTFimgs')
#print(debugpath)
if not os.path.exists(debugpath):
os.makedirs(debugpath)
import uuid
hex = uuid.uuid4().hex
cv2.imwrite(
debugpath + "\\" + im_name + hex + '_LR.png', img_LR * 255
) #random name to save + had to multiply by 255, else getting all black image
cv2.imwrite(
debugpath + "\\" + im_name + hex + '_HR.png', img_HR * 255
) #random name to save + had to multiply by 255, else getting all black image
# 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]]
# BGRA to RGBA, HWC to CHW, numpy to tensor
elif img_LR.shape[2] == 4:
img_HR = img_HR[:, :, [2, 1, 0, 3]]
img_LR = img_LR[:, :, [2, 1, 0, 3]]
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
}