def __init__(self, opt): super(MixDataset, self).__init__() self.opt = opt # temporal augmentation self.HQ_roots = opt["dataroots_HQ"] self.LQ_roots = opt["dataroots_LQ"] self.use_identical = opt["identical_loss"] dataset_weights = opt["dataset_weights"] self.data_type = "lmdb" # directly load image keys self.HQ_envs, self.LQ_envs = None, None self.paths_HQ = [] for idx, (HQ_root, LQ_root) in enumerate(zip(self.HQ_roots, self.LQ_roots)): paths_HQ, _ = util.get_image_paths(self.data_type, HQ_root) self.paths_HQ += list(zip([idx] * len(paths_HQ), paths_HQ)) * dataset_weights[idx] random.shuffle(self.paths_HQ) logger.info("Using lmdb meta info for cache keys.")
def __init__(self, opt): super(REDSDataset, self).__init__() self.opt = opt # temporal augmentation self.interval_list = opt['interval_list'] self.random_reverse = opt['random_reverse'] logger.info('Temporal augmentation interval list: [{}], with random reverse is {}.'.format( ','.join(str(x) for x in opt['interval_list']), self.random_reverse)) self.half_N_frames = opt['N_frames'] // 2 self.GT_root, self.LQ_root = opt['dataroot_GT'], opt['dataroot_LQ'] self.data_type = self.opt['data_type'] self.LR_input = False if opt['GT_size'] == opt['LQ_size'] else True # low resolution inputs #### directly load image keys if self.data_type == 'lmdb': self.paths_GT, _ = util.get_image_paths(self.data_type, opt['dataroot_GT']) logger.info('Using lmdb meta info for cache keys.') elif opt['cache_keys']: logger.info('Using cache keys: {}'.format(opt['cache_keys'])) self.paths_GT = pickle.load(open(opt['cache_keys'], 'rb'))['keys'] else: raise ValueError( 'Need to create cache keys (meta_info.pkl) by running [create_lmdb.py]') # remove the REDS4 for testing self.paths_GT = [ v for v in self.paths_GT if v.split('_')[0] not in ['000', '011', '015', '020'] ] assert self.paths_GT, 'Error: GT path is empty.' if self.data_type == 'lmdb': self.GT_env, self.LQ_env = None, None elif self.data_type == 'mc': # memcached self.mclient = None elif self.data_type == 'img': pass else: raise ValueError('Wrong data type: {}'.format(self.data_type))
def __init__(self, opt): super(MetaREDSDatasetOnline, self).__init__() self.opt = opt # temporal augmentation self.interval_list = opt['interval_list'] self.random_reverse = opt['random_reverse'] logger.info( 'Temporal augmentation interval list: [{}], with random reverse is {}.' .format(','.join(str(x) for x in opt['interval_list']), self.random_reverse)) self.half_N_frames = opt['N_frames'] // 2 self.GT_root = opt['dataroot_GT'] self.data_type = self.opt['data_type'] #### Load image keys if self.data_type == 'lmdb': self.paths_GT, self.GT_size_tuple = util.get_image_paths( self.data_type, self.GT_root) logger.info('Using lmdb meta info for cache keys.') self.paths_GT = [v for v in self.paths_GT if v.split('_')[0] not in \ ['000', '011', '015', '020']] else: seqs = sorted(os.listdir(self.GT_root)) self.paths_GT = [] for seq in seqs: if not seq in ['000', '011', '015', '020']: names = os.listdir(osp.join(self.GT_root, seq)) self.paths_GT.extend([seq + '_' + x[:-4] for x in names]) assert self.paths_GT, 'Error: GT path is empty.' if self.data_type == 'lmdb': self.GT_env = None self.scales = self.opt['scale'] assert (len(self.scales) >= 1) self.GT_sizes = self.opt['GT_size'] self.LQ_size = self.opt['LQ_size']
def __init__(self, opt, transform=DataTransform2(), acc_fac=None): super(HDF5Dataset, self).__init__() self.opt = opt self.acc_fac = acc_fac self.transform = transform self.paths_h5 = None self.data_type = 'h5' if self.acc_fac is not None: # Use both if self.acc_fac is None. assert self.acc_fac in (4, 8), 'Invalid acceleration factor' _, self.paths_h5 = util.get_image_paths(self.data_type, opt['dataroot_HR']) file_names = self.paths_h5 data_path = Path(opt['dataroot_HR']) #print(f'Initializing {data_path.stem}. This might take a minute') slice_counts = [ self.get_slice_number(file_name) for file_name in file_names ] self.num_slices = sum(slice_counts) names_and_slices = list() if self.acc_fac is not None: for name, slice_num in zip(file_names, slice_counts): names_and_slices += [[name, s_idx, self.acc_fac] for s_idx in range(slice_num)] else: for name, slice_num in zip(file_names, slice_counts): names_and_slices += [[name, s_idx, choice((4, 8))] for s_idx in range(slice_num)] self.names_and_slices = names_and_slices assert self.num_slices == len(names_and_slices), 'Error in length'
def __init__(self, opt, specific_image=None): super(JpegDataset, self).__init__() self.opt = opt self.paths_Uncomp = None self.Uncomp_env = None self.block_size = 8 if opt['input_downsampling'] is not None: self.block_size *= opt['input_downsampling'] self.quality_factors = opt['jpeg_quality_factor'] if not isinstance(self.quality_factors, list): self.quality_factors = [self.quality_factors] self.QF_probs = opt['QF_probs'] if self.QF_probs is None: self.QF_probs = np.ones([len(self.quality_factors)]) else: assert len(self.QF_probs) == len(self.quality_factors) self.QF_probs /= self.QF_probs.sum() # read image list from subset list txt if opt['subset_file'] is not None and opt['phase'] == 'train': with open(opt['subset_file']) as f: self.paths_Uncomp = sorted([os.path.join(opt['dataroot_Uncomp'], line.rstrip('\n')) \ for line in f]) if opt['dataroot_LR'] is not None: raise NotImplementedError( 'Now subset only supports generating LR on-the-fly.') else: # read image list from lmdb or image files self.Uncomp_env, self.paths_Uncomp = util.get_image_paths( opt['data_type'], opt['dataroot_Uncomp']) # self.Uncomp_env, self.paths_Uncomp = util.get_image_paths(opt['data_type'], # opt['dataroot_Uncomp'].replace('GrayScale','HRx4') if '_chroma' in opt['mode'] else opt['dataroot_Uncomp']) if opt['scales'] is not None: assert len(opt['scales']) == 3 new_paths_list = [] for scale_num, prob_ratio in enumerate(opt['scales']): if prob_ratio == 0: continue new_paths_list += prob_ratio * [ path for path in self.paths_Uncomp if '_scale%d_' % (scale_num) in path ] self.paths_Uncomp = new_paths_list assert self.paths_Uncomp, 'Error: Uncomp path is empty.' if self.opt['phase'] == 'train': assert not self.opt[ 'patch_size'] % 8, 'Training for JPEG compression artifacts removal - Training images should have an integer number of 8x8 blocks.' else: # self.per_index_QF = np.round(np.linspace(start=self.quality_factors[0][0],stop=self.quality_factors[0][1]-1,num=len(self))).astype(int) if len(self.quality_factors) >= len(self): sampled_QFs = np.round( np.linspace(start=0, stop=len(self.quality_factors), num=len(self))).astype(int) per_range_len = [ 1 if (QF in sampled_QFs) else 0 for QF in self.quality_factors ] else: num_exact_values = sum( [not isinstance(QF, list) for QF in self.quality_factors]) per_range_len = [ ((len(self) - num_exact_values) // (len(self.quality_factors) - num_exact_values)) if isinstance(QF, list) else 1 for QF in self.quality_factors ] if any([isinstance(QF, list) for QF in self.quality_factors]): per_range_len[np.argwhere([ isinstance(QF, list) for QF in self.quality_factors ])[0][0]] += len(self) - sum(per_range_len) else: per_range_len[0] += len(self) - sum(per_range_len) self.per_index_QF = [] for i, QF_range_len in enumerate(per_range_len): if isinstance(self.quality_factors[i], list): self.per_index_QF += list( np.round( np.linspace(start=self.quality_factors[i][0], stop=self.quality_factors[i][1] - 1, num=QF_range_len)).astype(int)) else: self.per_index_QF += [self.quality_factors[i] ] * QF_range_len self.random_scale_list = [1]
def __init__(self, opt): super(LRHRRefDataset, self).__init__() self.opt = opt self.paths_LR = [] self.paths_HR = [] self.paths_ref = [] if opt['data_type'] == 'lmdb': # only used in train phase import lmdb if opt['dataroot_LR'] is not None: self.LR_env = lmdb.open(opt['dataroot_LR'], readonly=True, \ lock=False, readahead=False, meminit=False) # get keys keys_cache_file = os.path.join(opt['dataroot_LR'], '_keys_cache.p') if os.path.isfile(keys_cache_file): print('read lmdb keys from cache: {}'.format( keys_cache_file)) keys = pickle.load(open(keys_cache_file, "rb")) else: with self.LR_env.begin(write=False) as txn: print('creating lmdb keys cache: {}'.format( keys_cache_file)) keys = [key.decode('ascii') for key, _ in txn.cursor()] pickle.dump(keys, open(keys_cache_file, "wb")) self.paths_LR = sorted( [key for key in keys if not key.endswith('.meta')]) if opt['dataroot_HR'] is not None: self.HR_env = lmdb.open(opt['dataroot_HR'], readonly=True, \ lock=False, readahead=False, meminit=False) # get keys keys_cache_file = os.path.join(opt['dataroot_HR'], '_keys_cache.p') if os.path.isfile(keys_cache_file): print('read lmdb keys from cache: {}'.format( keys_cache_file)) keys = pickle.load(open(keys_cache_file, "rb")) else: with self.HR_env.begin(write=False) as txn: print('creating lmdb keys cache: {}'.format( keys_cache_file)) keys = [key.decode('ascii') for key, _ in txn.cursor()] pickle.dump(keys, open(keys_cache_file, "wb")) self.paths_HR = sorted( [key for key in keys if not key.endswith('.meta')]) if self.paths_LR and self.paths_HR: assert len(self.paths_LR) == len(self.paths_HR), \ 'HR and LR lmdb datasets have different number of images.' # TODO lmdb does not support ref image now else: # read image from files if opt['phase'] == 'train' and opt['subset_file'] is not None: # get HR image paths from list with open(opt['subset_file']) as f: self.paths_HR = sorted([os.path.join(opt['dataroot_HR'], line.rstrip('\n')) \ for line in f]) if opt['dataroot_LR'] is not None: raise NotImplementedError( 'Now subset only support generating LR on-the-fly.') else: if opt['dataroot_LR'] is not None: self.paths_LR = sorted(get_image_paths(opt['dataroot_LR'])) if opt['dataroot_HR'] is not None: self.paths_HR = sorted(get_image_paths(opt['dataroot_HR'])) assert self.paths_LR or self.paths_HR, 'Both LR and HR paths are empty.' if self.paths_LR and self.paths_HR: assert len(self.paths_LR) == len(self.paths_HR), \ 'HR and LR datasets have different number of images.' # ref images if opt['dataroot_ref']: self.paths_ref = sorted(get_image_paths(opt['dataroot_ref']))
def __init__(self, opt): super(LRHRDataset, self).__init__() self.opt = opt self.paths_LR = None self.paths_HR = None self.LR_env = None # environment for lmdb self.HR_env = None self.HR_crop = None #v self.HR_rrot = None #v self.LR_scale = None #v self.LR_blur = None #v self.HR_noise = None #v self.LR_noise = None #v self.LR_noise2 = None #v self.LR_cutout = None #v self.LR_erasing = None #v # read image list from subset list txt if opt['subset_file'] is not None and opt['phase'] == 'train': with open(opt['subset_file']) as f: self.paths_HR = sorted([os.path.join(opt['dataroot_HR'], line.rstrip('\n')) \ for line in f]) if opt['dataroot_LR'] is not None: raise NotImplementedError('Now subset only supports generating LR on-the-fly.') else: # read image list from lmdb or image files self.HR_env, self.paths_HR = util.get_image_paths(opt['data_type'], opt['dataroot_HR']) self.LR_env, self.paths_LR = util.get_image_paths(opt['data_type'], opt['dataroot_LR']) assert self.paths_HR, 'Error: HR path is empty.' if self.paths_LR and self.paths_HR: assert len(self.paths_LR) == len(self.paths_HR), \ 'HR and LR datasets have different number of images - {}, {}.'.format(\ len(self.paths_LR), len(self.paths_HR)) #v parse on the fly options if opt['hr_crop']: #v variable to activate automatic crop of HR image to correct size and generate LR self.HR_crop = True print("Automatic crop of HR images enabled") if opt['hr_rrot']: #v variable to activate automatic rotate HR image and generate LR self.HR_rrot = True print("HR random rotation enabled") if opt['hr_noise']: #v variable to activate adding noise to HR image self.HR_noise = True self.hr_noise_types = opt['hr_noise_types'] print("HR_noise enabled") print(self.hr_noise_types) if opt['lr_downscale']: #v variable to activate automatic downscale of HR images to LR pair, controlled by the scale of the model self.LR_scale = True self.scale_algos = opt['lr_downscale_types'] print("LR_scale enabled") print(self.scale_algos) if opt['lr_blur']: #v variable to activate automatic blur of LR images self.LR_blur = True self.blur_algos = opt['lr_blur_types'] print("LR_blur enabled") print(self.blur_algos) if opt['lr_noise']: #v variable to activate adding noise to LR image self.LR_noise = True self.noise_types = opt['lr_noise_types'] print("LR_noise enabled") print(self.noise_types) if opt['lr_noise2']: #v variable to activate adding a secondary noise to LR image self.LR_noise2 = True self.noise_types2 = opt['lr_noise_types2'] print("LR_noise 2 enabled") print(self.noise_types2) if opt['lr_cutout']: #v variable to activate random cutout self.LR_cutout = True print("LR cutout enabled") if opt['lr_erasing']: #v variable to activate random erasing self.LR_erasing = True print("LR random erasing enabled") #v parse on the fly options self.random_scale_list = [1]
def __init__(self, opt): self.opt = opt input_folder = opt['input_folder'] self.images = data_util.get_image_paths('img', input_folder)[0] print("Found %i images" % (len(self.images), ))
def __init__(self, opt): super(REDSDataset, self).__init__() self.opt = opt # temporal augmentation self.interval_list = opt['interval_list'] self.random_reverse = opt['random_reverse'] logger.info('Temporal augmentation interval list: [{}], with random reverse is {}.'.format( ','.join(str(x) for x in opt['interval_list']), self.random_reverse)) self.half_N_frames = opt['N_frames'] // 2 self.GT_root, self.LQ_root = opt['dataroot_GT'], opt['dataroot_LQ'] self.data_type = self.opt['data_type'] self.LR_input = False if opt['GT_size'] == opt['LQ_size'] else True # low resolution inputs #### directly load image keys if self.data_type == 'lmdb': self.paths_GT, _ = util.get_image_paths(self.data_type, opt['dataroot_GT']) logger.info('Using lmdb meta info for cache keys.') elif opt['cache_keys']: logger.info('Using cache keys: {}'.format(opt['cache_keys'])) self.paths_GT = pickle.load(open(opt['cache_keys'], 'rb'))['keys'] else: raise ValueError( 'Need to create cache keys (meta_info.pkl) by running [create_lmdb.py]') # # remove the REDS4 for testing # self.paths_GT = [ # v for v in self.paths_GT if v.split('_')[0] not in ['000', '011', '015', '020'] # ] # remove the 4Khdr for testing if not opt['all_dataset_for_train'] == True: print("exclude validation set for training") self.paths_GT = [ v for v in self.paths_GT if v.split('_')[0] not in ['00000', '00001', '00002', '00003', '00004', '00005', '00006', '00007', '00008', '00009'] ] if len(opt['baddata_list']) > 0: print("exclude bad training data") print(opt['baddata_list']) self.paths_GT = [ v for v in self.paths_GT if v.split('_')[0] not in opt['baddata_list'] ] print("Include validation set for training") assert self.paths_GT, 'Error: GT path is empty.' if self.data_type == 'lmdb': self.GT_env, self.LQ_env = None, None elif self.data_type == 'mc': # memcached self.mclient = None elif self.data_type == 'img': pass else: raise ValueError('Wrong data type: {}'.format(self.data_type)) # load screen notation with open(opt['frame_notation']) as f: self.frame_notation = json.load(f) print("Training Dataset Initialized")
def __init__(self, opt): super(VideoTestDataset, self).__init__() self.opt = opt self.cache_data = opt['cache_data'] self.half_N_frames = opt['N_frames'] // 2 self.GT_root, self.LQ_root = opt['dataroot_GT'], opt['dataroot_LQ'] self.data_type = self.opt['data_type'] self.data_info = { 'path_LQ': [], 'path_GT': [], 'folder': [], 'idx': [], 'border': [] } #### Generate data info and cache data self.imgs_LQ, self.imgs_GT = {}, {} if opt['name'].lower() in ['vid4', 'reds4']: if self.data_type == 'lmdb': self.lmdb_paths_GT, _ = util.get_image_paths( self.data_type, self.GT_root) self.lmdb_paths_LQ, _ = util.get_image_paths( self.data_type, self.LQ_root) self.GT_env, self.LQ_env = None, None previous_name_a = None previous_name_b = None for lmdb_path_GT, lmdb_path_LQ in zip(self.lmdb_paths_GT, self.lmdb_paths_LQ): GT_name_a, GT_name_b = lmdb_path_GT.split('_') assert lmdb_path_GT == lmdb_path_LQ, 'GT path and LQ path in lmdb is not matched' if previous_name_a != GT_name_a and previous_name_a is not None: max_idx = int(previous_name_b) + 1 for i in range(max_idx): self.data_info['idx'].append('{}/{}'.format( i, max_idx)) border_l = [0] * max_idx for i in range(self.half_N_frames): border_l[i] = 1 border_l[max_idx - i - 1] = 1 self.data_info['border'].extend(border_l) self.data_info['folder'].append(GT_name_a) previous_name_a = GT_name_a previous_name_b = GT_name_b max_idx = int(previous_name_b) + 1 for i in range(max_idx): self.data_info['idx'].append('{}/{}'.format(i, max_idx)) border_l = [0] * max_idx for i in range(self.half_N_frames): border_l[i] = 1 border_l[max_idx - i - 1] = 1 self.data_info['border'].extend(border_l) else: subfolders_LQ = util.glob_file_list(self.LQ_root) subfolders_GT = util.glob_file_list(self.GT_root) for subfolder_LQ, subfolder_GT in zip(subfolders_LQ, subfolders_GT): subfolder_name = osp.basename(subfolder_GT) img_paths_LQ = util.glob_file_list(subfolder_LQ) img_paths_GT = util.glob_file_list(subfolder_GT) max_idx = len(img_paths_LQ) assert max_idx == len( img_paths_GT ), 'Different number of images in LQ and GT folders' self.data_info['path_LQ'].extend(img_paths_LQ) self.data_info['path_GT'].extend(img_paths_GT) self.data_info['folder'].extend([subfolder_name] * max_idx) for i in range(max_idx): self.data_info['idx'].append('{}/{}'.format( i, max_idx)) border_l = [0] * max_idx for i in range(self.half_N_frames): border_l[i] = 1 border_l[max_idx - i - 1] = 1 self.data_info['border'].extend(border_l) if self.cache_data: self.imgs_LQ[subfolder_name] = util.read_img_seq( img_paths_LQ) self.imgs_GT[subfolder_name] = util.read_img_seq( img_paths_GT) else: self.imgs_LQ[subfolder_name] = img_paths_LQ self.imgs_GT[subfolder_name] = img_paths_GT elif opt['name'].lower() in ['vimeo90k-test']: pass # TODO else: raise ValueError( 'Not support video test dataset. Support Vid4, REDS4 and Vimeo90k-Test.' )
def __init__(self, opt): super(LRHRDataset, self).__init__() self.opt = opt self.paths_LR = None self.paths_HR = None self.LR_env = None # environment for lmdb self.HR_env = None #self.HR_crop = None #v self.HR_rrot = None #v self.LR_scale = None #v self.scale_algos = None #v self.LR_blur = None #v self.HR_noise = None #v self.LR_noise = None #v self.LR_noise2 = None #v self.LR_cutout = None #v self.LR_erasing = None #v self.output_sample_imgs = None # read image list from subset list txt if opt['subset_file'] is not None and opt['phase'] == 'train': with open(opt['subset_file']) as f: self.paths_HR = sorted([os.path.join(opt['dataroot_HR'], line.rstrip('\n')) \ for line in f]) if opt['dataroot_LR'] is not None: raise NotImplementedError( 'Now subset only supports generating LR on-the-fly.') else: # read image list from lmdb or image files self.HR_env, self.paths_HR = util.get_image_paths( opt['data_type'], opt['dataroot_HR']) self.LR_env, self.paths_LR = util.get_image_paths( opt['data_type'], opt['dataroot_LR']) assert self.paths_HR, 'Error: HR path is empty.' if self.paths_LR and self.paths_HR: # Modify to allow using HR and LR folders with different amount of images # - If an LR image pair is not found, downscale HR on the fly, else, use the LR # - If all LR are provided and 'lr_downscale' is enabled, randomize use of provided LR and OTF LR for augmentation """ assert len(self.paths_LR) == len(self.paths_HR), \ 'HR and LR datasets have different number of images - {}, {}.'.format(\ len(self.paths_LR), len(self.paths_HR)) """ #""" assert len(self.paths_HR) >= len(self.paths_LR), \ 'HR dataset contains less images than LR dataset - {}, {}.'.format(\ len(self.paths_LR), len(self.paths_HR)) if len(self.paths_LR) < len(self.paths_HR): print( 'LR contains less images than HR dataset - {}, {}. Will generate missing images on the fly.' .format(len(self.paths_LR), len(self.paths_HR))) import os i = 0 tmp = [] for idx in range(0, len(self.paths_HR)): _, HRtail = os.path.split(self.paths_HR[idx]) if i < len(self.paths_LR): LRhead, LRtail = os.path.split(self.paths_LR[i]) if LRtail == HRtail: LRimg_path = os.path.join(LRhead, LRtail) tmp.append(LRimg_path) i += 1 else: LRimg_path = None tmp.append(LRimg_path) else: #if the last image is missing LRimg_path = None tmp.append(LRimg_path) self.paths_LR = tmp #""" #v parse on the fly options if opt['hr_rrot']: #v variable to activate automatic rotate HR image and generate LR self.HR_rrot = True print("HR random rotation enabled") if opt['hr_noise']: #v variable to activate adding noise to HR image self.HR_noise = True self.hr_noise_types = opt['hr_noise_types'] print("HR_noise enabled") print(self.hr_noise_types) if opt['lr_downscale']: #v variable to activate automatic downscale of HR images to LR pair, controlled by the scale of the model self.LR_scale = True self.scale_algos = opt['lr_downscale_types'] print("LR_scale enabled") print(self.scale_algos) if opt['lr_blur']: #v variable to activate automatic blur of LR images self.LR_blur = True self.blur_algos = opt['lr_blur_types'] print("LR_blur enabled") print(self.blur_algos) if opt['lr_noise']: #v variable to activate adding noise to LR image self.LR_noise = True self.noise_types = opt['lr_noise_types'] print("LR_noise enabled") print(self.noise_types) if opt['lr_noise2']: #v variable to activate adding a secondary noise to LR image self.LR_noise2 = True self.noise_types2 = opt['lr_noise_types2'] print("LR_noise 2 enabled") print(self.noise_types2) if opt['lr_cutout']: #v variable to activate random cutout self.LR_cutout = True print("LR cutout enabled") if opt['lr_erasing']: #v variable to activate random erasing self.LR_erasing = True print("LR random erasing enabled")
def __init__(self, opt): super(AI4KDataset, self).__init__() self.opt = opt if opt['video_class']: self.video_class = opt[ 'video_class'] # all | movie | cartoon | lego else: self.video_class = 'all' # temporal augmentation self.interval_list = opt['interval_list'] self.random_reverse = opt['random_reverse'] logger.info( 'Temporal augmentation interval list: [{}], with random reverse is {}.' .format(','.join(str(x) for x in opt['interval_list']), self.random_reverse)) self.half_N_frames = opt['N_frames'] // 2 self.GT_root, self.LQ_root = opt['dataroot_GT'], opt['dataroot_LQ'] self.data_type = self.opt['data_type'] self.LR_input = False if opt['GT_size'] == opt[ 'LQ_size'] else True # low resolution inputs #### directly load image keys if self.data_type == 'lmdb': self.paths_GT, _ = util.get_image_paths(self.data_type, opt['dataroot_GT']) train_list = [] if self.video_class == 'all': pass elif self.video_class == 'movie': with open('data/movie_list.txt', 'r') as f: for line in f.readlines(): line = line.strip() train_list.append(line) #print((train_list)) for item in self.paths_GT.copy(): if item.split('_')[0] not in train_list: self.paths_GT.remove(item) elif self.video_class == 'cartoon': with open('data/cartoon_list.txt', 'r') as f: for line in f.readlines(): line = line.strip() train_list.append(line) for item in self.paths_GT.copy(): if item.split('_')[0] not in train_list: self.paths_GT.remove(item) elif self.video_class == 'lego': with open('data/lego_list.txt', 'r') as f: for line in f.readlines(): line = line.strip() train_list.append(line) for item in self.paths_GT.copy(): if item.split('_')[0] not in train_list: self.paths_GT.remove(item) # clear bad data for item in self.paths_GT.copy(): if item.split('_')[0] == '15922480': self.paths_GT.remove(item) logger.info('Using lmdb meta info for cache keys.') elif self.data_type == 'img': self.paths_GT, _ = util.get_image_paths(self.data_type, opt['dataroot_GT']) elif opt['cache_keys']: logger.info('Using cache keys: {}'.format(opt['cache_keys'])) self.paths_GT = pickle.load(open(opt['cache_keys'], 'rb'))['keys'] #else: # raise ValueError( # 'Need to create cache keys (meta_info.pkl) by running [create_lmdb.py]') assert self.paths_GT, 'Error: GT path is empty.' #rint((self.paths_GT)) print(len(self.paths_GT)) if self.data_type == 'lmdb': self.GT_env, self.LQ_env = None, None elif self.data_type == 'mc': # memcached self.mclient = None elif self.data_type == 'img': pass else: raise ValueError('Wrong data type: {}'.format(self.data_type))
def __init__(self, opt): super(ImageLabelDatasetFiveLevels, self).__init__() self.opt = opt # read image list from subset list txt if opt['subset_file'] is not None and opt['phase'] == 'train': with open(opt['subset_file']) as f: self.paths_Images = sorted([ os.path.join(opt['dataroot_HR_Image'], line.rstrip('\n')) for line in f ]) if opt['dataroot_LR_Image'] is not None: raise NotImplementedError( 'Now subset only supports generating LR on-the-fly.') else: # read image list from lmdb or image files # self.HR_images_env, self.paths_HR_images = util.get_image_paths(opt['data_type'], opt['dataroot_HR_Image']) # self.LR_labels_env, self.paths_LR_labels = util.get_image_paths(opt['data_type'], opt['dataroot_HR_Label']) self.HR_images_env, self.paths_HR_images = util.get_image_paths( "lmdb", opt['dataroot_HR_Image']) self.LR_labels_env, self.paths_LR_labels = util.get_image_paths( "lmdb", opt['dataroot_HR_Label']) # self.LR_images_env, self.paths_LR_images = util.get_image_paths(opt['data_type'], opt['dataroot_LR_Image']) # self.LR_labels_env, self.paths_LR_labels = util.get_image_paths(opt['data_type'], opt['dataroot_LR_Label']) # # self.D1_images_env, self.paths_D1_images = util.get_image_paths(opt['data_type'], opt['dataroot_D1_Image']) # self.D1_labels_env, self.paths_D1_labels = util.get_image_paths(opt['data_type'], opt['dataroot_D1_Label']) # # self.D2_images_env, self.paths_D2_images = util.get_image_paths(opt['data_type'], opt['dataroot_D2_Image']) # self.D2_labels_env, self.paths_D2_labels = util.get_image_paths(opt['data_type'], opt['dataroot_D2_Label']) # # self.D3_images_env, self.paths_D3_images = util.get_image_paths(opt['data_type'], opt['dataroot_D3_Image']) # self.D3_labels_env, self.paths_D3_labels = util.get_image_paths(opt['data_type'], opt['dataroot_D3_Label']) # # self.D4_images_env, self.paths_D4_images = util.get_image_paths(opt['data_type'], opt['dataroot_D4_Image']) # self.D4_labels_env, self.paths_D4_labels = util.get_image_paths(opt['data_type'], opt['dataroot_D4_Label']) print("Images: {}".format(len(self.paths_HR_images))) print("Labels: {}".format(len(self.paths_LR_labels))) if opt['max_len'] is not None: print("Trimming the first", opt['max_len']) self.paths_HR_images = self.paths_HR_images[:opt['max_len']] # self.paths_LR_images = self.paths_LR_images[:opt['max_len']] self.paths_LR_labels = self.paths_LR_labels[:opt['max_len']] # self.paths_D1_images = self.paths_D1_images[:opt['max_len']] # self.paths_D1_labels = self.paths_D1_labels[:opt['max_len']] # self.paths_D2_images = self.paths_D2_images[:opt['max_len']] # self.paths_D2_labels = self.paths_D2_labels[:opt['max_len']] # self.paths_D3_images = self.paths_D3_images[:opt['max_len']] # self.paths_D3_labels = self.paths_D3_labels[:opt['max_len']] # self.paths_D4_images = self.paths_D4_images[:opt['max_len']] # self.paths_D4_labels = self.paths_D4_labels[:opt['max_len']] assert self.paths_HR_images, 'Error: Images path is empty.' # if self.paths_HR_images and self.paths_LR_images: # assert len(self.paths_HR_images) == len(self.paths_LR_images), \ # 'Images and Labels datasets have different number of images - {}, {}.'.format( \ # len(self.paths_HR_images), len(self.paths_LR_images)) self.random_scale_list = [1] # rarity maks self.rarity_masks = [] # if opt['phase'] == 'train': # # batch size for modifying the input index in the training time # # self.batch_size = opt["batch_size_per_month"] # if 'rarity_mask_1' in self.opt: # for i in range(0, int(math.log(opt['scale'], 2))): # self.rarity_masks.append(np.load(opt['rarity_mask_{}'.format(i + 1)], mmap_mode='r')) # # # rarity bins for dataset re-balancing # self.objectid = [0, 2, 10, 19, 1] # self.objectnum = 5 # self.rarity_bin = [] # for i in range(self.objectnum): # self.rarity_bin.append(np.load(opt['rarity_bins'] % self.objectid[i])) # color palette dataset self.dataset_color_map = util.PaletteDataset(opt['palette'])
def __init__(self, opt): super(OnlineVideoTestDataset, self).__init__() self.opt = opt self.data_name = opt['name'] self.cache_data = opt['cache_data'] self.half_N_frames = opt['N_frames'] // 2 self.GT_root = opt['dataroot_GT'] self.is_lmdb = self.GT_root.endswith('lmdb') self.data_info = self.data_info = { 'path_GT': [], 'idx': [], 'folder': [], 'border': [] } self.scale = opt['scale'] self.imgs_GT = {} if self.data_name.lower() in ['vid4', 'reds4']: if self.is_lmdb: raise TypeError("{} data should not lmdb".format( self.data_name)) subs = sorted(os.listdir(self.GT_root)) for sub in subs: sub_dir = osp.join(self.GT_root, sub) im_names = sorted(os.listdir(sub_dir)) im_paths = [osp.join(sub_dir, name) for name in im_names] max_idx = len(im_names) self.data_info['path_GT'].extend(im_paths) self.data_info['folder'].extend([sub] * max_idx) for i in range(max_idx): self.data_info['idx'].append('{}/{}'.format(i, max_idx)) border_l = [0] * max_idx for i in range(self.half_N_frames): border_l[i] = 1 border_l[max_idx - i - 1] = 1 self.data_info['border'].extend(border_l) if self.cache_data: self.imgs_GT[sub] = self.read_sub_images(im_paths) elif self.data_name.lower() in ['vimeo', 'vimeo90k', 'vimeo90k-test']: if self.is_lmdb: paths_GT, self.GT_size_tuple = util.get_image_paths( 'lmdb', self.GT_root) self.data_info['path_GT'] = [ x for x in paths_GT if x.endswith('_4') ] else: split_file = osp.join(self.GT_root, 'sep_trainlist.txt') paths_GT = [] with open(txt_file, 'r') as f: lines = f.readlines() img_list = [line.strip() for line in lines] for item in img_list: key = osp.join(*item.split('/'), 'im4.png') paths_GT.append(osp.join(self.GT_root, 'sequences', key)) self.data_info['path_GT'] = paths_GT else: raise ValueError( 'Not support video test dataset. Support Vid4, REDS4 and Vimeo90k-Test.' )
def __init__(self, opt): super(LRHRDataset, self).__init__() self.opt = opt self.paths_LR = None self.paths_HR = None self.LR_env = None # environment for lmdb self.HR_env = None self.output_sample_imgs = None # read image list from subset list txt if opt['subset_file'] is not None and opt['phase'] == 'train': with open(opt['subset_file']) as f: self.paths_HR = sorted([ os.path.join(opt['dataroot_HR'], line.rstrip('\n')) for line in f ]) if opt['dataroot_LR'] is not None: raise NotImplementedError( 'Now subset only supports generating LR on-the-fly.') else: # read image list from lmdb or image files # Check if dataroot_HR is a list of directories or a single directory. Note: lmdb will not currently work with a list HR_images_paths = opt['dataroot_HR'] if type(HR_images_paths) is list: self.HR_env = [] self.paths_HR = [] for path in HR_images_paths: HR_env, paths_HR = util.get_image_paths( opt['data_type'], path) if type(HR_env) is list: for imgs in HR_env: self.HR_env.append(imgs) for imgs in paths_HR: self.paths_HR.append(imgs) if self.HR_env.count(None) == len(self.HR_env): self.HR_env = None else: self.HR_env = sorted(self.HR_env) self.paths_HR = sorted(self.paths_HR) elif type(HR_images_paths) is str: self.HR_env, self.paths_HR = util.get_image_paths( opt['data_type'], HR_images_paths) # Check if dataroot_LR is a list of directories or a single directory. Note: lmdb will not currently work with a list LR_images_paths = opt['dataroot_LR'] if type(LR_images_paths) is list: self.LR_env = [] self.paths_LR = [] for path in LR_images_paths: LR_env, paths_LR = util.get_image_paths( opt['data_type'], path) if type(LR_env) is list: for imgs in LR_env: self.LR_env.append(imgs) for imgs in paths_LR: self.paths_LR.append(imgs) if self.LR_env.count(None) == len(self.LR_env): self.LR_env = None else: self.LR_env = sorted(self.LR_env) self.paths_LR = sorted(self.paths_LR) elif type(LR_images_paths) is str: self.LR_env, self.paths_LR = util.get_image_paths( opt['data_type'], LR_images_paths) assert self.paths_HR, 'Error: HR path is empty.' if self.paths_LR and self.paths_HR: # Modify to allow using HR and LR folders with different amount of images # - If an LR image pair is not found, downscale HR on the fly, else, use the LR # - If all LR are provided and 'lr_downscale' is enabled, randomize use of provided LR and OTF LR for augmentation """ assert len(self.paths_LR) == len(self.paths_HR), \ 'HR and LR datasets have different number of images - {}, {}.'.format(\ len(self.paths_LR), len(self.paths_HR)) """ # """ assert len(self.paths_HR) >= len(self.paths_LR), \ 'HR dataset contains less images than LR dataset - {}, {}.'.format( len(self.paths_LR), len(self.paths_HR)) # """ if len(self.paths_LR) < len(self.paths_HR): print( 'LR contains less images than HR dataset - {}, {}. Will generate missing images on the fly.' .format(len(self.paths_LR), len(self.paths_HR))) i = 0 tmp = [] for idx in range(0, len(self.paths_HR)): _, HRtail = os.path.split(self.paths_HR[idx]) if i < len(self.paths_LR): LRhead, LRtail = os.path.split(self.paths_LR[i]) if LRtail == HRtail: LRimg_path = os.path.join(LRhead, LRtail) tmp.append(LRimg_path) i += 1 else: LRimg_path = None tmp.append(LRimg_path) else: # if the last image is missing LRimg_path = None tmp.append(LRimg_path) self.paths_LR = tmp
def __init__(self, opt): self.opt = opt self.corruptor = ImageCorruptor(opt) if 'center_crop_hq_sz' in opt.keys(): self.center_crop = functools.partial(ndarray_center_crop, opt['center_crop_hq_sz']) self.target_hq_size = opt['target_size'] if 'target_size' in opt.keys( ) else None self.multiple = opt['force_multiple'] if 'force_multiple' in opt.keys( ) else 1 self.scale = opt['scale'] self.paths = opt['paths'] self.corrupt_before_downsize = opt[ 'corrupt_before_downsize'] if 'corrupt_before_downsize' in opt.keys( ) else False self.fetch_alt_image = opt[ 'fetch_alt_image'] # If specified, this dataset will attempt to find a second image # from the same video source. Search for 'fetch_alt_image' for more info. self.fetch_alt_tiled_image = opt[ 'fetch_alt_tiled_image'] # If specified, this dataset will attempt to find anoter tile from the same source image # Search for 'fetch_alt_tiled_image' for more info. assert not (self.fetch_alt_image and self.fetch_alt_tiled_image ) # These are mutually exclusive. self.skip_lq = opt_get(opt, ['skip_lq'], False) self.disable_flip = opt_get(opt, ['disable_flip'], False) self.rgb_n1_to_1 = opt_get(opt, ['rgb_n1_to_1'], False) self.force_square = opt_get(opt, ['force_square'], True) self.fixed_parameters = { k: torch.tensor(v) for k, v in opt_get(opt, ['fixed_parameters'], {}).items() } self.all_image_color_jitter = opt_get(opt, ['all_image_color_jitter'], 0) if 'normalize' in opt.keys(): if opt['normalize'] == 'stylegan2_norm': self.normalize = Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True) elif opt['normalize'] == 'imagenet': self.normalize = Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225), inplace=True) else: raise Exception('Unsupported normalize') else: self.normalize = None if self.target_hq_size is not None: assert ( self.target_hq_size // self.scale ) % self.multiple == 0 # If we dont throw here, we get some really obscure errors. if not isinstance(self.paths, list): self.paths = [self.paths] self.weights = [1] else: self.weights = opt['weights'] if 'labeler' in opt.keys(): if opt['labeler']['type'] == 'patch_labels': self.labeler = VsNetImageLabeler(opt['labeler']['label_file']) assert len( self.paths ) == 1 # Only a single base-path is supported for labeled images. self.image_paths = self.labeler.get_labeled_paths(self.paths[0]) else: self.labeler = None # Just scan the given directory for images of standard types. supported_types = ['jpg', 'jpeg', 'png', 'gif'] self.image_paths = [] for path, weight in zip(self.paths, self.weights): cache_path = os.path.join(path, 'cache.pth') if os.path.exists(cache_path): imgs = torch.load(cache_path) else: print( "Building image folder cache, this can take some time for large datasets.." ) imgs = util.get_image_paths('img', path)[0] torch.save(imgs, cache_path) for w in range(weight): self.image_paths.extend(imgs) self.len = len(self.image_paths)
def main(): device = torch.device("cuda") parser = argparse.ArgumentParser(description="Kernel extractor testing") parser.add_argument("--source_H", action="store", help="source image height", type=int, required=True) parser.add_argument("--source_W", action="store", help="source image width", type=int, required=True) parser.add_argument("--target_H", action="store", help="target image height", type=int, required=True) parser.add_argument("--target_W", action="store", help="target image width", type=int, required=True) parser.add_argument("--augmented_H", action="store", help="desired height of the augmented images", type=int, required=True) parser.add_argument("--augmented_W", action="store", help="desired width of the augmented images", type=int, required=True) parser.add_argument("--source_LQ_root", action="store", help="source low-quality dataroot", type=str, required=True) parser.add_argument("--source_HQ_root", action="store", help="source high-quality dataroot", type=str, required=True) parser.add_argument("--target_HQ_root", action="store", help="target high-quality dataroot", type=str, required=True) parser.add_argument("--save_path", action="store", help="save path", type=str, required=True) parser.add_argument("--yml_path", action="store", help="yml path", type=str, required=True) parser.add_argument("--num_images", action="store", help="number of desire augmented images", type=int, required=True) args = parser.parse_args() source_LQ_root = args.source_LQ_root source_HQ_root = args.source_HQ_root target_HQ_root = args.target_HQ_root save_path = args.save_path source_H, source_W = args.source_H, args.source_W target_H, target_W = args.target_H, args.target_W augmented_H, augmented_W = args.augmented_H, args.augmented_W yml_path = args.yml_path num_images = args.num_images # Initializing logger logger = logging.getLogger("base") os.makedirs(save_path, exist_ok=True) util.setup_logger("base", save_path, "test", level=logging.INFO, screen=True, tofile=True) logger.info("source LQ root: {}".format(source_LQ_root)) logger.info("source HQ root: {}".format(source_HQ_root)) logger.info("target HQ root: {}".format(target_HQ_root)) logger.info("augmented height: {}".format(augmented_H)) logger.info("augmented width: {}".format(augmented_W)) logger.info("Number of augmented images: {}".format(num_images)) # Initializing mode logger.info("Loading model...") with open(yml_path, "r") as f: print(yml_path) opt = yaml.load(f)["KernelWizard"] model_path = opt["pretrained"] model = KernelWizard(opt) model.eval() model.load_state_dict(torch.load(model_path)) model = model.to(device) logger.info("Done") # processing data source_HQ_env = lmdb.open(source_HQ_root, readonly=True, lock=False, readahead=False, meminit=False) source_LQ_env = lmdb.open(source_LQ_root, readonly=True, lock=False, readahead=False, meminit=False) target_HQ_env = lmdb.open(target_HQ_root, readonly=True, lock=False, readahead=False, meminit=False) paths_source_HQ, _ = data_util.get_image_paths("lmdb", source_HQ_root) paths_target_HQ, _ = data_util.get_image_paths("lmdb", target_HQ_root) psnr_avg = 0 for i in range(num_images): source_key = np.random.choice(paths_source_HQ) target_key = np.random.choice(paths_target_HQ) source_rnd_h = random.randint(0, max(0, source_H - augmented_H)) source_rnd_w = random.randint(0, max(0, source_W - augmented_W)) target_rnd_h = random.randint(0, max(0, target_H - augmented_H)) target_rnd_w = random.randint(0, max(0, target_W - augmented_W)) source_LQ = read_image(source_LQ_env, source_key, source_rnd_h, source_rnd_w, augmented_H, augmented_W) source_HQ = read_image(source_HQ_env, source_key, source_rnd_h, source_rnd_w, augmented_H, augmented_W) target_HQ = read_image(target_HQ_env, target_key, target_rnd_h, target_rnd_w, augmented_H, augmented_W) source_LQ = torch.Tensor(source_LQ).unsqueeze(0).to(device) source_HQ = torch.Tensor(source_HQ).unsqueeze(0).to(device) target_HQ = torch.Tensor(target_HQ).unsqueeze(0).to(device) with torch.no_grad(): kernel_mean, kernel_sigma = model(source_HQ, source_LQ) kernel = kernel_mean + kernel_sigma * torch.randn_like(kernel_mean) fake_source_LQ = model.adaptKernel(source_HQ, kernel) target_LQ = model.adaptKernel(target_HQ, kernel) LQ_img = util.tensor2img(source_LQ) fake_LQ_img = util.tensor2img(fake_source_LQ) target_LQ_img = util.tensor2img(target_LQ) target_HQ_img = util.tensor2img(target_HQ) target_HQ_dst = osp.join( save_path, "sharp/{:03d}/{:08d}.png".format(i // 100, i % 100)) target_LQ_dst = osp.join( save_path, "blur/{:03d}/{:08d}.png".format(i // 100, i % 100)) os.makedirs(osp.dirname(target_HQ_dst), exist_ok=True) os.makedirs(osp.dirname(target_LQ_dst), exist_ok=True) cv2.imwrite(target_HQ_dst, target_HQ_img) cv2.imwrite(target_LQ_dst, target_LQ_img) # torch.save(kernel, osp.join(osp.dirname(target_LQ_dst), f'kernel{i:03d}.pth')) psnr = util.calculate_psnr(LQ_img, fake_LQ_img) logger.info( "Reconstruction PSNR of image #{:03d}/{:03d}: {:.2f}db".format( i, num_images, psnr)) psnr_avg += psnr logger.info("Average reconstruction PSNR: {:.2f}db".format(psnr_avg / num_images))