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 }