def __getitem__(self, index): LR_path = None # get LR image LR_path = self.paths_LR[index] #img_LR = util.read_img(self.LR_env, LR_path) img_LR = util.read_img(self.LR_env, LR_path) H, W, C = img_LR.shape # change color space if necessary if self.opt['color']: img_LR = util.channel_convert(C, self.opt['color'], [img_LR])[0] # BGR to RGB, HWC to CHW, numpy to tensor img_LR = util.np2tensor(img_LR, normalize=self.znorm, add_batch=False) return {'LR': img_LR, 'LR_path': LR_path}
def __getitem__(self, index): """Return a data point and its metadata information. Parameters: index (int): a random integer for data indexing Returns a dictionary that contains A and A_paths or LR and LR_path A (tensor): an image in the input domain A_paths (str): the path of the image """ # get single image img_A, A_path = read_single_dataset(self.opt, index, self.A_paths, self.A_env) # change color space if necessary # TODO: move to get_transform() if self.opt.get('color', None): img_A = channel_convert(image_channels(img_A), self.opt['color'], [img_A])[0] # # apply transforms if any # default_int_method = get_default_imethod(image_type(img_A)) # transform_params = get_params(self.opt, image_size(img_A)) # A_transform = get_transform( # self.opt, # transform_params, # # grayscale=(input_nc == 1), # method=default_int_method) # img_A = A_transform(img_A) ######## Convert images to PyTorch Tensors ######## img_A = self.tensor_transform(img_A) if self.vars == 'A': return {'A': img_A, 'A_path': A_path} else: return {'LR': img_A, 'LR_path': A_path}
def __getitem__(self, index): HR_path, LR_path = None, None scale = self.opt.get('scale', 4) HR_size = self.opt.get('HR_size', 128) if HR_size: LR_size = HR_size // scale # Default case: tensor will result in the [0,1] range # Alternative: tensor will be z-normalized to the [-1,1] range znorm = self.opt.get('znorm', False) ######## Read the images ######## #TODO: check cases where default of 3 channels will be troublesome image_channels = self.opt.get('image_channels', 3) # 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.get('rand_flip_LR_HR', None) and self.opt['phase'] == 'train': LRHRchance = random.uniform(0, 1) flip_chance = self.opt.get('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, out_nc=image_channels) img_HR = util.read_img(self.HR_env, HR_path, out_nc=image_channels) # 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.get('aug_downscale', None): #aug_downscale = self.opt['aug_downscale'] if np.random.rand() < self.opt['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, out_nc=image_channels) img_LR = img_HR LR_path = HR_path # tmp_vis(img_HR, False) # tmp_vis(img_LR, False) ######## 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.get('color', None): # Change both 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] if self.opt.get('color_HR', None): # Only change HR img_HR = util.channel_convert(img_HR.shape[2], self.opt['color_HR'], [img_HR])[0] if self.opt.get('color_LR', None): # Only change LR img_LR = util.channel_convert(img_LR.shape[2], self.opt['color_LR'], [img_LR])[0] ######## Augmentations ######## #Augmentations during training if self.opt['phase'] == 'train': # Note: this is NOT recommended, HR should not be exposed to degradation, as it will # degrade the model's results, only added because it exist as an option in downstream forks # HR downscale if self.opt.get('hr_downscale', None): ds_algo = self.opt.get('hr_downscale_types', 777) hr_downscale_amt = self.opt.get('hr_downscale_amt', 2) if isinstance(hr_downscale_amt, list): hr_downscale_amt = random.choice(hr_downscale_amt) # will ignore if 1 or if result is smaller than hr size if hr_downscale_amt > 1 and img_HR.shape[ 0] // hr_downscale_amt >= HR_size and img_HR.shape[ 1] // hr_downscale_amt >= HR_size: #img_HR, hr_scale_interpol_algo = augmentations.scale_img(img_HR, hr_downscale_amt, algo=ds_algo) img_HR, _ = Scale(img=img_HR, scale=hr_downscale_amt, algo=ds_algo) # Downscales LR to match new size of HR if scale does not match after if img_LR is not None and ( img_HR.shape[0] // scale != img_LR.shape[0] or img_HR.shape[1] // scale != img_LR.shape[1]): #img_LR, lr_scale_interpol_algo = augmentations.scale_img(img_LR, hr_downscale_amt, algo=ds_algo) img_LR, _ = Scale(img=img_LR, scale=hr_downscale_amt, algo=ds_algo) # 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 use an option to process #TODO: add options variable shape_change = self.opt.get('shape_change', 'reshape_hr') #"reshape_lr" if img_HR.shape[0] // img_LR.shape[0] != img_HR.shape[ 1] // img_LR.shape[1]: if shape_change == "reshape_lr": img_LR = transforms.Resize( (int(img_HR.shape[0] / scale), int(img_HR.shape[1] / scale)), interpolation="BILINEAR")(np.copy(img_LR)) elif shape_change == "reshape_hr": # print("pre: ", img_HR.shape[0], img_HR.shape[1]) nh = img_HR.shape[0] * (2 * img_LR.shape[1]) // ( img_LR.shape[0] + img_LR.shape[1]) nw = img_HR.shape[1] * (2 * img_LR.shape[0]) // ( img_LR.shape[0] + img_LR.shape[1]) #TODO: temporary change to test contextual loss with unaligned LR-HR pairs, forcing them to have the correct scale #img_HR, _ = augmentations.resize_img(np.copy(img_HR), newdim=(img_LR.shape[1]*scale,img_LR.shape[0]*scale), algo=cv2.INTER_LINEAR) #img_HR = transforms.Resize((img_LR.shape[0]*scale,img_LR.shape[1]*scale), interpolation="BILINEAR")(np.copy(img_HR)) img_HR = transforms.Resize( (nh, nw), interpolation="BILINEAR")(np.copy(img_HR)) # print("post: ", img_HR.shape[0], img_HR.shape[1]) img_LR = transforms.Resize( (int(img_HR.shape[0] / scale), int(img_HR.shape[1] / scale)), interpolation="BILINEAR")(np.copy(img_LR)) #TODO: check, this may not make sense # elif shape_change == "pad_lr": # LR_pad, _ = get_pad(img_LR, HR_size//scale, fill=0, padding_mode='constant') # img_LR = LR_pad(np.copy(img_LR)) else: # generate new LR from HR #TODO: disabled to test cx loss #print("Warning: img_LR dimensions ratio does not match img_HR dimensions ratio for: ", HR_path) img_LR = img_HR # centercrop image if self.opt['center_crop'] == True: if img_HR.shape[0] > img_HR.shape[1]: img_HR = transforms.CenterCrop( (img_HR.shape[1], img_HR.shape[1]))(np.copy(img_HR)) img_LR = transforms.CenterCrop( (img_HR.shape[1], img_HR.shape[1]))(np.copy(img_LR)) else: img_HR = transforms.CenterCrop( (img_HR.shape[0], img_HR.shape[0]))(np.copy(img_HR)) img_LR = transforms.CenterCrop( (img_HR.shape[0], img_HR.shape[0]))(np.copy(img_LR)) # resize lr and hr image to given hr dimension if self.opt['resize_HR_dimension'] == True: img_HR = transforms.Resize( (self.opt['HR_size'], self.opt['HR_size']), interpolation="BILINEAR")(np.copy(img_HR)) img_LR = transforms.Resize( (self.opt['HR_size'], self.opt['HR_size']), interpolation="BILINEAR")(np.copy(img_LR)) # 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) #TODO: test to be removed # img_HR = transforms.Resize((HR_size-100,img_HR.shape[1]), interpolation="BILINEAR")(np.copy(img_HR)) # img_LR = transforms.Resize((LR_size-(100//scale),img_LR.shape[1]), interpolation="BILINEAR")(np.copy(img_LR)) # Or if the HR images are too small, Resize to the HR_size size and fit LR pair to LR_size too #TODO: add options variable dim_change = self.opt.get('dim_change', 'pad') if img_HR.shape[0] < HR_size or img_HR.shape[1] < HR_size: if dim_change == "resize": #TODO: temp disabled to test #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), newdim=(HR_size,HR_size), algo=cv2.INTER_LINEAR) img_HR = transforms.Resize( (HR_size, HR_size), interpolation="BILINEAR")(np.copy(img_HR)) # 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), newdim=(LR_size,LR_size), algo=cv2.INTER_LINEAR) img_LR = transforms.Resize( (LR_size, LR_size), interpolation="BILINEAR")(np.copy(img_LR)) elif dim_change == "pad": # if img_LR is img_HR, padding will be wrong, downscaling LR before padding if img_LR.shape[ 0] != img_HR.shape[0] // scale or img_LR.shape[ 1] != img_HR.shape[1] // scale: ds_algo = 777 # default to matlab-like bicubic downscale if self.opt.get( 'lr_downscale', None ): # if manually set and scale algorithms are provided, then: ds_algo = self.opt.get('lr_downscale_types', 777) if self.opt.get('lr_downscale', None) and self.opt.get( 'dataroot_kernels', None ) and 999 in self.opt["lr_downscale_types"]: ds_kernel = self.ds_kernels else: ds_kernel = None img_LR, _ = Scale(img=img_LR, scale=scale, algo=ds_algo, ds_kernel=ds_kernel) HR_pad, fill = get_pad(img_HR, HR_size, fill='random', padding_mode=self.opt.get( 'pad_mode', 'constant')) img_HR = HR_pad(np.copy(img_HR)) LR_pad, _ = get_pad(img_LR, HR_size // scale, fill=fill, padding_mode=self.opt.get( 'pad_mode', 'constant')) img_LR = LR_pad(np.copy(img_LR)) # (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 self.opt.get( 'lr_downscale', None ): # if manually set and scale algorithms are provided, then: ds_algo = self.opt.get('lr_downscale_types', 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) print( "LR image is too large, auto generating new LR for: ", LR_path) if self.opt.get('lr_downscale', None) and self.opt.get( 'dataroot_kernels', None) and 999 in self.opt["lr_downscale_types"]: ds_kernel = self.ds_kernels #KernelDownscale(scale, self.kernel_paths, self.num_kernel) else: ds_kernel = None #img_LR, scale_interpol_algo = augmentations.scale_img(img_LR, scale, algo=ds_algo) #print("algo") #print(ds_algo) img_LR, _ = Scale(img=img_LR, scale=scale, algo=ds_algo, ds_kernel=ds_kernel) # resize = get_resize(size=LR_size, scale=scale, ds_algo=ds_algo) # img_LR = resize(np.copy(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.get( 'hr_rrot', None): 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.get('hr_rrot', None): # augmentation - flip, rotate img_LR, img_HR = util.augment([img_LR, img_HR], self.opt['use_flip'], \ self.opt['use_rot']) elif self.opt.get('hr_rrot', None): 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 sizes 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 the resulting LR so far does not have the correct dimensions, also generate a new HR-LR image pair on the fly if img_HR.shape[0] != HR_size or img_HR.shape[ 1] != HR_size or img_LR.shape[ 0] != LR_size or img_LR.shape[1] != LR_size: #if img_HR.shape[0] != HR_size or img_HR.shape[1] != HR_size: #TODO: temp disabled to test #print("Image: ", HR_path, " size does not match HR size: (", HR_size,"). The image size is: ", img_HR.shape) #if img_LR.shape[0] != LR_size or img_LR.shape[0] != LR_size: #TODO: temp disabled to test #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 in LR case, but something went wrong before, just for sanity) #img_HR, _ = augmentations.resize_img(np.copy(img_HR), newdim=(HR_size,HR_size), algo=cv2.INTER_LINEAR) img_HR = transforms.Resize( (HR_size, HR_size), interpolation="BILINEAR")(np.copy(img_HR)) # if manually provided and scale algorithms are provided, then use it, else use matlab imresize to generate LR pair ds_algo = self.opt.get('lr_downscale_types', 777) #img_LR, _ = augmentations.scale_img(img_HR, scale, algo=ds_algo) img_LR, _ = Scale(img_HR, scale, algo=ds_algo) # Below are the On The Fly augmentations # Apply "auto levels" to images #rand_levels = (1 - self.opt.get('rand_auto_levels', 0)) # Randomize for augmentation rand_levels = self.opt.get('rand_auto_levels', 0) #rand_percent = self.opt.get('rand_auto_per', 10) #if self.opt.get('auto_levels', None) and np.random.rand() > rand_levels: if self.opt.get('auto_levels', None) == 'HR' or self.opt.get( 'auto_levels', None) == 'Both': #img_HR = augmentations.simplest_cb(img_HR, znorm=znorm) #TODO: now images are processed in the [0,255] range #img_HR = augmentations.simplest_cb(img_HR) img_HR = transforms.FilterColorBalance( p=rand_levels, percent=10, random_params=True)(img_HR) if self.opt.get('auto_levels', None) == 'LR' or self.opt.get( 'auto_levels', None) == 'Both': #img_LR = augmentations.simplest_cb(img_LR, znorm=znorm) #TODO: now images are processed in the [0,255] range #img_LR = augmentations.simplest_cb(img_LR) img_LR = transforms.FilterColorBalance( p=rand_levels, percent=10, random_params=True)(img_LR) # Apply unsharpening mask to LR images ''' rand_unsharp = (1 - self.opt.get('lr_rand_unsharp', 0)) # Randomize for augmentation if self.opt.get('lr_unsharp_mask', None) and np.random.rand() > rand_unsharp: #img_LR = augmentations.unsharp_mask(img_LR, znorm=znorm) #TODO: now images are processed in the [0,255] range img_LR = augmentations.unsharp_mask(img_LR) ''' if self.opt.get('lr_unsharp_mask', None): lr_rand_unsharp = self.opt.get('lr_rand_unsharp', 0) img_LR = transforms.FilterUnsharp(p=lr_rand_unsharp)(img_LR) # Apply unsharpening mask to HR images ''' rand_unsharp = (1 - self.opt.get('hr_rand_unsharp', 0)) # Randomize for augmentation if self.opt.get('hr_unsharp_mask', None) and np.random.rand() > rand_unsharp: #img_HR = augmentations.unsharp_mask(img_HR, znorm=znorm) #TODO: now images are processed in the [0,255] range img_HR = augmentations.unsharp_mask(img_HR) ''' if self.opt.get('hr_unsharp_mask', None): hr_rand_unsharp = self.opt.get('hr_rand_unsharp', 0) img_HR = transforms.FilterUnsharp(p=hr_rand_unsharp)(img_HR) # Add noise to HR if enabled AND noise types are provided (for noise2noise and similar) if self.opt.get('hr_noise', None): noise_option = get_noise(self.opt.get('hr_noise_types', None)) if noise_option: #img_HR, hr_noise_algo = augmentations.noise_img(img_HR, noise_types=self.opt['hr_noise_types']) img_HR = noise_option(img_HR) #TODO: check if necessary # 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.get('lr_fringes', None): 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.get('lr_blur', None): blur_option = get_blur(self.opt.get('lr_blur_types', None)) if blur_option: #img_LR, blur_algo, blur_kernel_size = augmentations.blur_img(img_LR, blur_algos=blur_algos) img_LR = blur_option(img_LR) #TODO: check if necessary # 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.get('lr_noise', None): noise_option = get_noise(self.opt.get('lr_noise_types', None), self.noise_patches) if noise_option: #img_LR, noise_algo = augmentations.noise_img(img_LR, noise_types=self.opt['lr_noise_types']) img_LR = noise_option(img_LR) #TODO: check if necessary # 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.get('lr_noise2', None): noise_option = get_noise(self.opt.get('lr_noise_types2', None), self.noise_patches) if noise_option: #img_LR, noise_algo2 = augmentations.noise_img(img_LR, noise_types=self.opt['lr_noise_types2']) img_LR = noise_option(img_LR) #TODO: check if necessary # 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.get( 'lr_cutout', None) and (self.opt.get('lr_erasing', None) != True): # img_LR = augmentations.cutout(img_LR, img_LR.shape[0] // 2) img_LR = transforms.Cutout(p=1, mask_size=img_LR.shape[0] // 2)(img_LR) elif self.opt.get( 'lr_erasing', None) and (self.opt.get('lr_cutout', None) != True): # img_LR = augmentations.random_erasing(img_LR) img_LR = transforms.RandomErasing(p=1)(img_LR, mode=[3]) elif self.opt.get('lr_cutout', None) and self.opt.get( 'lr_erasing', None): 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) img_LR = transforms.Cutout(p=1, mask_size=img_LR.shape[0] // 2)(img_LR) else: #img_LR = augmentations.random_erasing(img_LR, p=0.5, modes=[3]) img_LR = transforms.RandomErasing(p=1)(img_LR, mode=[3]) #""" # 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) img_LR, _ = Scale(img_LR, scale, algo=self.opt.get('lr_downscale_types', 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 #TODO: use the debugging functions to visualize or save images instead # 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, os.sep, 'sampleOTFimgs') # 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) import uuid hex = uuid.uuid4().hex cv2.imwrite(debugpath + "\\" + im_name + hex + '_LR.png', img_LR) #random name to save cv2.imwrite(debugpath + "\\" + im_name + hex + '_HR.png', img_HR) #random name to save # cv2.imwrite(debugpath+"\\"+im_name+hex+'_HR1.png',img_HRn1) #random name to save ######## Convert images to PyTorch Tensors ######## """ # for debugging if (img_HR.min() < -1): describe_numpy(img_HR, all=True) print(HR_path) if (img_HR.max() > 1): describe_numpy(img_HR, all=True) print(HR_path) if (img_LR.min() < -1): describe_numpy(img_LR, all=True) print(LR_path) if (img_LR.max() > 1): describe_numpy(img_LR, all=True) print(LR_path) #""" # check for grayscale images #TODO: should not be needed anymore if len(img_HR.shape) < 3: img_HR = img_HR[..., np.newaxis] if len(img_LR.shape) < 3: img_LR = img_LR[..., np.newaxis] # tmp_vis(img_HR, False) # tmp_vis(img_LR, False) # edge-connect if self.opt['training_with_canny'] == True: img_HR_gray = cv2.cvtColor(img_HR, cv2.COLOR_BGR2GRAY) img_HR_canny = cv2.Canny(img_HR_gray, 100, 150) img_HR_gray = torch.from_numpy(img_HR_gray).unsqueeze(0) img_HR_canny = torch.from_numpy(img_HR_canny).unsqueeze(0) if self.opt['training_with_canny_SR'] == True: img_LR_gray = cv2.cvtColor(img_LR, cv2.COLOR_BGR2GRAY) img_LR_canny = cv2.Canny(img_LR_gray, 100, 150) img_LR_canny = torch.from_numpy(img_LR_canny).unsqueeze(0) if self.opt['noise_estimation'] == True: ds_kernel = torch.from_numpy(getattr( ds_kernel, 'used_kernel')).unsqueeze(0).float() noise_est = noise_estimate(img_LR, 4) sigma = torch.tensor(noise_est).float().view([1, 1, 1]) img_HR = util.np2tensor(img_HR, normalize=znorm, add_batch=False) #.astype('uint8').clip(0,255) img_LR = util.np2tensor(img_LR, normalize=znorm, add_batch=False) if LR_path is None: LR_path = HR_path if self.opt['training_with_canny'] == True: return { 'LR': img_LR, 'HR': img_HR, 'LR_path': LR_path, 'HR_path': HR_path, 'img_HR_gray': img_HR_gray, 'img_HR_canny': img_HR_canny } elif self.opt['training_with_canny_SR'] == True: return { 'LR': img_LR, 'HR': img_HR, 'LR_path': LR_path, 'HR_path': HR_path, 'img_LR_canny': img_LR_canny } if self.opt['noise_estimation'] == True: return { 'LR': img_LR, 'HR': img_HR, 'LR_path': LR_path, 'HR_path': HR_path, 'ds_kernel': ds_kernel, 'sigma': sigma } else: return { 'LR': img_LR, 'HR': img_HR, 'LR_path': LR_path, 'HR_path': HR_path }
def __getitem__(self, index): """Return a data point and its metadata information. Parameters: index (int): a random integer for data indexing Returns a dictionary that contains A, B, A_paths and B_paths (or LR, HR, LR_paths and HR_paths) A (tensor): an image in the input domain B (tensor): its corresponding image in the target domain A_paths (str): paths A images B_paths (str): paths B images (can be same as A_paths if using single images) """ scale = self.opt.get('scale') ######## Read the images ######## if self.AB_paths: img_A, img_B, A_path, B_path = read_split_single_dataset( self.opt, index, self.AB_paths, self.AB_env) else: img_A, img_B, A_path, B_path = read_imgs_from_path( self.opt, index, self.A_paths, self.B_paths, self.A_env, self.B_env) ######## Modify the images ######## # for the validation / test phases if self.opt['phase'] != 'train': img_type = image_type(img_B) # B/HR modcrop img_B = modcrop(img_B, scale=scale, img_type=img_type) # modcrop and downscale A/LR if enabled if self.opt['lr_downscale']: img_A = modcrop(img_A, scale=scale, img_type=img_type) # TODO: 'pil' images will use default method for scaling img_A, _ = Scale(img_A, scale, algo=self.opt.get('lr_downscale_types', 777), img_type=img_type) # change color space if necessary # TODO: move to get_transform() color_B = self.opt.get('color', None) or self.opt.get('color_HR', None) if color_B: img_B = channel_convert(image_channels(img_B), color_B, [img_B])[0] color_A = self.opt.get('color', None) or self.opt.get('color_LR', None) if color_A: img_A = channel_convert(image_channels(img_A), color_A, [img_A])[0] ######## Augmentations ######## #Augmentations during training if self.opt['phase'] == 'train': default_int_method = get_default_imethod(image_type(img_A)) # random HR downscale img_A, img_B = random_downscale_B(img_A=img_A, img_B=img_B, opt=self.opt) # validate there's an img_A, if not, use img_B if img_A is None: img_A = img_B print( f"Image A: {A_path} was not loaded correctly, using B pair to downscale on the fly." ) # validate proper dimensions between paired images, generate A if needed img_A, img_B = paired_imgs_check(img_A, img_B, opt=self.opt, ds_kernels=self.ds_kernels) # get and apply the paired transformations below transform_params = get_params(scale_opt(self.opt, scale), image_size(img_A)) A_transform = get_transform( scale_opt(self.opt, scale), transform_params, # grayscale=(input_nc == 1), method=default_int_method) B_transform = get_transform( self.opt, scale_params(transform_params, scale), # grayscale=(output_nc == 1), method=default_int_method) img_A = A_transform(img_A) img_B = B_transform(img_B) # Below are the On The Fly augmentations # get and apply the unpaired transformations below a_aug_params, b_aug_params = get_unpaired_params(self.opt) a_augmentations = get_augmentations( self.opt, params=a_aug_params, noise_patches=self.noise_patches, ) b_augmentations = get_augmentations( self.opt, params=b_aug_params, noise_patches=self.noise_patches, ) img_A = a_augmentations(img_A) img_B = b_augmentations(img_B) # Alternative position for changing the colorspace of A/LR. # color_A = self.opt.get('color', None) or self.opt.get('color_A', None) # if color_A: # img_A = channel_convert(image_channels(img_A), color_A, [img_A])[0] ######## Convert images to PyTorch Tensors ######## tensor_transform = get_totensor(self.opt, params=self.totensor_params, toTensor=True, grayscale=False) img_A = tensor_transform(img_A) img_B = tensor_transform(img_B) if A_path is None: A_path = B_path if self.vars == 'AB': return {'A': img_A, 'B': img_B, 'A_path': A_path, 'B_path': B_path} else: return { 'LR': img_A, 'HR': img_B, 'LR_path': A_path, 'HR_path': B_path }
def __getitem__(self, index): """Return a data point and its metadata information. Parameters: index (int): a random integer for data indexing Returns a dictionary that contains A, B, A_paths and B_paths A (tensor): an image in the input domain B (tensor): its corresponding image in the target domain A_paths (str): paths A images B_paths (str): paths B images """ scale = self.opt.get('scale') ######## Read the images ######## img_A, A_path = read_single_dataset(opt=self.opt, index=index, paths=self.A_paths, env=self.A_env, idx_case='inrange', d_size=self.A_size) img_B, B_path = read_single_dataset(opt=self.opt, index=index, paths=self.B_paths, env=self.B_env, idx_case=self.idx_case, d_size=self.B_size) ######## Modify the images ######## # change color space if necessary # TODO: move to get_transform() color_B = self.opt.get('color', None) or self.opt.get('color_B', None) if color_B: img_B = channel_convert(image_channels(img_B), color_B, [img_B])[0] color_A = self.opt.get('color', None) or self.opt.get('color_A', None) if color_A: img_A = channel_convert(image_channels(img_A), color_A, [img_A])[0] # apply image transformation default_int_method = get_default_imethod(image_type(img_A)) # get first set of random params transform_params_A = get_params(scale_opt(self.opt, scale), image_size(img_A)) # get second set of random params transform_params_B = get_params(self.opt, image_size(img_B)) A_transform = get_transform( scale_opt(self.opt, scale), transform_params_A, # grayscale=(input_nc == 1), method=default_int_method) B_transform = get_transform( self.opt, transform_params_B, # grayscale=(output_nc == 1), method=default_int_method) img_A = A_transform(img_A) img_B = B_transform(img_B) #TODO: not needed initially, but available # get and apply the unpaired transformations below # a_aug_params, b_aug_params = get_unpaired_params(self.opt) # a_augmentations = get_augmentations( # self.opt, # params=a_aug_params, # noise_patches=self.noise_patches, # ) # b_augmentations = get_augmentations( # self.opt, # params=b_aug_params, # noise_patches=self.noise_patches, # ) # img_A = a_augmentations(img_A) # img_B = b_augmentations(img_B) ######## Convert images to PyTorch Tensors ######## tensor_transform = get_totensor(self.opt, params=self.totensor_params, toTensor=True, grayscale=False) img_A = tensor_transform(img_A) img_B = tensor_transform(img_B) if self.vars == 'AB': return {'A': img_A, 'B': img_B, 'A_path': A_path, 'B_path': B_path} else: return { 'LR': img_A, 'HR': img_B, 'LR_path': A_path, 'HR_path': B_path }
def __getitem__(self, index): HR_path, LR_path = None, None scale = self.opt['scale'] HR_size = self.opt['HR_size'] # get HR image HR_path = self.paths_HR[index] 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] # get LR image if self.paths_LR: LR_path = self.paths_LR[index] img_LR = util.read_img(self.LR_env, LR_path) 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 # 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]] 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 }