def compute_current_psnr_ssim(self, save=False, name=None, save_path=None): """ compute ssim, psnr when validate the model """ num = self.get_info() visuals = self.get_current_visuals() psnr_interp_t_t = [] ssim_interp_t_t = [] for i in range(num): rlt_img = util.tensor2img(visuals['rlt'][i]) gt_img = util.tensor2img(visuals['GT'][i]) psnr = util.calculate_psnr(rlt_img, gt_img) ssim = util.calculate_ssim(rlt_img, gt_img) psnr_interp_t_t.append(psnr) ssim_interp_t_t.append(ssim) if save == True: import os.path as osp import cv2 cv2.imwrite(osp.join(save_path, 'rlt_{}_{}.png'.format(name, i)), rlt_img) cv2.imwrite(osp.join(save_path, 'gt_{}_{}.png'.format(name, i)), gt_img) return psnr_interp_t_t, ssim_interp_t_t
def _calculate_metrics(sr_vol, gt_vol, view='xy'): sum_psnr = 0. sum_ssim = 0. sum_pdist = 0. # [D,H,W] num_val = 0 # psnr could be inf at xz or yz (near edges), will not calculate for i, vol in enumerate(zip(sr_vol, gt_vol)): sr_img, gt_img = vol[0], vol[1] # calculate PSNR and SSIM # range is assume to be [0,255] so have to scale back from 1500 to 255 float64 crop_size = round(opt['scale']) cropped_sr_img = sr_img[crop_size:-crop_size, crop_size:-crop_size]\ .astype(np.float64) / 1500. * 255. cropped_gt_img = gt_img[crop_size:-crop_size, crop_size:-crop_size]\ .astype(np.float64) / 1500. * 255. psnr = util.calculate_psnr(cropped_sr_img, cropped_gt_img) ssim = util.calculate_ssim(cropped_sr_img, cropped_gt_img) if opt['datasets']['val']['get_pdist']: pdist = util.calculate_pdist(pdist_model, cropped_sr_img, cropped_gt_img) else: pdist = float('nan') if psnr != float('inf'): num_val += 1 sum_psnr += psnr sum_ssim += ssim sum_pdist += pdist logger.info('{:20s} - {:3d}- PSNR: {:.6f} dB; SSIM: {:.6f}; pdist: {:.6f}.'\ .format(patient_id, i+1, psnr, ssim, pdist)) pnsr_results[patient_id][view] = sum_psnr / num_val ssim_results[patient_id][view] = sum_ssim / num_val pdist_results[patient_id][view] = sum_pdist / num_val return pnsr_results, ssim_results, pdist_results
def cascade_test_main(opt, logger, model, test_loader): test_set_name = test_loader.dataset.opt['name'] logger.info('\nTesting [{:s}]...'.format(test_set_name)) results = [] try: total_nfe = model.netG.module.conv_trunk.nfe except AttributeError: total_nfe = None for data in test_loader: need_GT = False if test_loader.dataset.opt[ 'dataroot_GT'] is None else True model.feed_data(data, need_GT=need_GT) img_path = data['GT_path'][0] if need_GT else data['LQ_path'][0] img_name = osp.splitext(osp.basename(img_path))[0] model.test() visuals = model.get_current_visuals(need_GT=need_GT) sr_img = util.tensor2img(visuals['SR']) # uint8 # calculate PSNR and SSIM if need_GT: gt_img = util.tensor2img(visuals['GT']) gt_img = gt_img / 255. sr_img = sr_img / 255. crop_border = opt['crop_border'] if opt['crop_border'] else opt[ 'scale'] if crop_border == 0: cropped_sr_img = sr_img cropped_gt_img = gt_img else: cropped_sr_img = sr_img[crop_border:-crop_border, crop_border:-crop_border, :] cropped_gt_img = gt_img[crop_border:-crop_border, crop_border:-crop_border, :] psnr = util.calculate_psnr(cropped_sr_img * 255, cropped_gt_img * 255) ssim = util.calculate_ssim(cropped_sr_img * 255, cropped_gt_img * 255) # niqe = util.calculate_niqe(cropped_sr_img * 255) if total_nfe is not None: last_nfe = model.netG.module.conv_trunk.nfe - total_nfe total_nfe = model.netG.module.conv_trunk.nfe else: last_nfe = None if gt_img.shape[2] == 3: # RGB image sr_img_y = bgr2ycbcr(sr_img, only_y=True) gt_img_y = bgr2ycbcr(gt_img, only_y=True) if crop_border == 0: cropped_sr_img_y = sr_img_y cropped_gt_img_y = gt_img_y else: cropped_sr_img_y = sr_img_y[crop_border:-crop_border, crop_border:-crop_border] cropped_gt_img_y = gt_img_y[crop_border:-crop_border, crop_border:-crop_border] psnr_y = util.calculate_psnr(cropped_sr_img_y * 255, cropped_gt_img_y * 255) ssim_y = util.calculate_ssim(cropped_sr_img_y * 255, cropped_gt_img_y * 255) results.append((psnr, ssim, psnr_y, ssim_y)) logger.info( '{:20s} - PSNR: {:.6f} dB; SSIM: {:.6f}; PSNR_Y: {:.6f} dB; SSIM_Y: {:.6f}; NFE: {}' .format(img_name, psnr, ssim, psnr_y, ssim_y, last_nfe)) else: logger.info( '{:20s} - PSNR: {:.6f} dB; SSIM: {:.6f}; NFE: {}'.format( img_name, psnr, ssim, last_nfe)) else: logger.info(img_name) return results
def main(): #################### # arguments parser # #################### # [format] dataset(vid4, REDS4) N(number of frames) # data_mode = str(args.dataset) # N_in = int(args.n_frames) # metrics = str(args.metrics) # output_format = str(args.output_format) ################# # configurations ################# device = torch.device('cuda') os.environ['CUDA_VISIBLE_DEVICES'] = '0' #data_mode = 'Vid4' # Vid4 | sharp_bicubic | blur_bicubic | blur | blur_comp # Vid4: SR # REDS4: sharp_bicubic (SR-clean), blur_bicubic (SR-blur); # blur (deblur-clean), blur_comp (deblur-compression). # STAGE Vid4 # Collecting results for Vid4 model_path = '../experiments/pretrained_models/EDVR_Vimeo90K_SR_L.pth' stage = 1 # 1 or 2, use two stage strategy for REDS dataset. flip_test = False predeblur, HR_in = False, False back_RBs = 40 N_model_default = 7 data_mode = 'Vid4' # vid4_dir_map = {"calendar": 0, "city": 1, "foliage": 2, "walk": 3} vid4_results = {"calendar": {}, "city": {}, "foliage": {}, "walk": {}} #vid4_results = 4 * [[]] for N_in in range(1, N_model_default + 1): raw_model = EDVR_arch.EDVR(128, N_model_default, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in) model = EDVR_arch.EDVR(128, N_in, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in) test_dataset_folder = '../datasets/Vid4/BIx4' GT_dataset_folder = '../datasets/Vid4/GT' aposterior_GT_dataset_folder = '../datasets/Vid4/GT_7' crop_border = 0 border_frame = N_in // 2 # border frames when evaluate padding = 'new_info' save_imgs = False raw_model.load_state_dict(torch.load(model_path), strict=True) model.nf = raw_model.nf model.center = N_in // 2 # if center is None else center model.is_predeblur = raw_model.is_predeblur model.HR_in = raw_model.HR_in model.w_TSA = raw_model.w_TSA if model.is_predeblur: model.pre_deblur = raw_model.pre_deblur # Predeblur_ResNet_Pyramid(nf=nf, HR_in=self.HR_in) model.conv_1x1 = raw_model.conv_1x1 # nn.Conv2d(nf, nf, 1, 1, bias=True) else: if model.HR_in: model.conv_first_1 = raw_model.conv_first_1 # nn.Conv2d(3, nf, 3, 1, 1, bias=True) model.conv_first_2 = raw_model.conv_first_2 # nn.Conv2d(nf, nf, 3, 2, 1, bias=True) model.conv_first_3 = raw_model.conv_first_3 # nn.Conv2d(nf, nf, 3, 2, 1, bias=True) else: model.conv_first = raw_model.conv_first # nn.Conv2d(3, nf, 3, 1, 1, bias=True) model.feature_extraction = raw_model.feature_extraction # arch_util.make_layer(ResidualBlock_noBN_f, front_RBs) model.fea_L2_conv1 = raw_model.fea_L2_conv1 # nn.Conv2d(nf, nf, 3, 2, 1, bias=True) model.fea_L2_conv2 = raw_model.fea_L2_conv2 # nn.Conv2d(nf, nf, 3, 1, 1, bias=True) model.fea_L3_conv1 = raw_model.fea_L3_conv1 # nn.Conv2d(nf, nf, 3, 2, 1, bias=True) model.fea_L3_conv2 = raw_model.fea_L3_conv2 # nn.Conv2d(nf, nf, 3, 1, 1, bias=True) model.pcd_align = raw_model.pcd_align # PCD_Align(nf=nf, groups=groups) model.tsa_fusion.center = model.center model.tsa_fusion.tAtt_1 = raw_model.tsa_fusion.tAtt_1 model.tsa_fusion.tAtt_2 = raw_model.tsa_fusion.tAtt_2 model.tsa_fusion.fea_fusion = copy.deepcopy(raw_model.tsa_fusion.fea_fusion) model.tsa_fusion.fea_fusion.weight = copy.deepcopy(torch.nn.Parameter(raw_model.tsa_fusion.fea_fusion.weight[:, 0:N_in * 128, :, :])) model.tsa_fusion.sAtt_1 = copy.deepcopy(raw_model.tsa_fusion.sAtt_1) model.tsa_fusion.sAtt_1.weight = copy.deepcopy(torch.nn.Parameter(raw_model.tsa_fusion.sAtt_1.weight[:, 0:N_in * 128, :, :])) model.tsa_fusion.maxpool = raw_model.tsa_fusion.maxpool model.tsa_fusion.avgpool = raw_model.tsa_fusion.avgpool model.tsa_fusion.sAtt_2 = raw_model.tsa_fusion.sAtt_2 model.tsa_fusion.sAtt_3 = raw_model.tsa_fusion.sAtt_3 model.tsa_fusion.sAtt_4 = raw_model.tsa_fusion.sAtt_4 model.tsa_fusion.sAtt_5 = raw_model.tsa_fusion.sAtt_5 model.tsa_fusion.sAtt_L1 = raw_model.tsa_fusion.sAtt_L1 model.tsa_fusion.sAtt_L2 = raw_model.tsa_fusion.sAtt_L2 model.tsa_fusion.sAtt_L3 = raw_model.tsa_fusion.sAtt_L3 model.tsa_fusion.sAtt_add_1 = raw_model.tsa_fusion.sAtt_add_1 model.tsa_fusion.sAtt_add_2 = raw_model.tsa_fusion.sAtt_add_2 model.tsa_fusion.lrelu = raw_model.tsa_fusion.lrelu model.recon_trunk = raw_model.recon_trunk model.upconv1 = raw_model.upconv1 model.upconv2 = raw_model.upconv2 model.pixel_shuffle = raw_model.pixel_shuffle model.HRconv = raw_model.HRconv model.conv_last = raw_model.conv_last model.lrelu = raw_model.lrelu ##################################################### model.eval() model = model.to(device) #avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l = [], [], [] subfolder_name_l = [] subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*'))) subfolder_GT_l = sorted(glob.glob(osp.join(GT_dataset_folder, '*'))) subfolder_GT_a_l = sorted(glob.glob(osp.join(aposterior_GT_dataset_folder, "*"))) # for each subfolder for subfolder, subfolder_GT, subfolder_GT_a in zip(subfolder_l, subfolder_GT_l, subfolder_GT_a_l): subfolder_name = osp.basename(subfolder) subfolder_name_l.append(subfolder_name) img_path_l = sorted(glob.glob(osp.join(subfolder, '*'))) max_idx = len(img_path_l) print("MAX_IDX: ", max_idx) #### read LQ and GT images imgs_LQ = data_util.read_img_seq(subfolder) img_GT_l = [] for img_GT_path in sorted(glob.glob(osp.join(subfolder_GT, '*'))): img_GT_l.append(data_util.read_img(None, img_GT_path)) img_GT_a = [] for img_GT_a_path in sorted(glob.glob(osp.join(subfolder_GT_a, '*'))): img_GT_a.append(data_util.read_img(None, img_GT_a_path)) #avg_psnr, avg_psnr_border, avg_psnr_center, N_border, N_center = 0, 0, 0, 0, 0 # process each image for img_idx, img_path in enumerate(img_path_l): img_name = osp.splitext(osp.basename(img_path))[0] select_idx = data_util.index_generation(img_idx, max_idx, N_in, padding=padding) imgs_in = imgs_LQ.index_select(0, torch.LongTensor(select_idx)).unsqueeze(0).to(device) if flip_test: output = util.flipx4_forward(model, imgs_in) else: print("IMGS_IN SHAPE: ", imgs_in.shape) output = util.single_forward(model, imgs_in) output = util.tensor2img(output.squeeze(0)) # calculate PSNR output = output / 255. GT = np.copy(img_GT_l[img_idx]) # For REDS, evaluate on RGB channels; for Vid4, evaluate on the Y channel #if data_mode == 'Vid4': # bgr2y, [0, 1] GT = data_util.bgr2ycbcr(GT, only_y=True) output = data_util.bgr2ycbcr(output, only_y=True) GT_a = np.copy(img_GT_a[img_idx]) GT_a = data_util.bgr2ycbcr(GT_a, only_y=True) output_a = copy.deepcopy(output) output, GT = util.crop_border([output, GT], crop_border) crt_psnr = util.calculate_psnr(output * 255, GT * 255) crt_ssim = util.calculate_ssim(output * 255, GT * 255) output_a, GT_a = util.crop_border([output_a, GT_a], crop_border) crt_aposterior = util.calculate_ssim(output_a * 255, GT_a * 255) # CHANGE t = vid4_results[subfolder_name].get(str(img_name)) if t != None: vid4_results[subfolder_name][img_name].add_psnr(crt_psnr) vid4_results[subfolder_name][img_name].add_gt_ssim(crt_ssim) vid4_results[subfolder_name][img_name].add_aposterior_ssim(crt_aposterior) else: vid4_results[subfolder_name].update({img_name: metrics_file(img_name)}) vid4_results[subfolder_name][img_name].add_psnr(crt_psnr) vid4_results[subfolder_name][img_name].add_gt_ssim(crt_ssim) vid4_results[subfolder_name][img_name].add_aposterior_ssim(crt_aposterior) ############################################################################ #### model #### writing vid4 results util.mkdirs('../results/calendar') util.mkdirs('../results/city') util.mkdirs('../results/foliage') util.mkdirs('../results/walk') save_folder = '../results/' for i, dir_name in enumerate(["calendar", "city", "foliage", "walk"]): save_subfolder = osp.join(save_folder, dir_name) for j, value in vid4_results[dir_name].items(): # cur_result = json.dumps(_) with open(osp.join(save_subfolder, '{}.json'.format(value.name)), 'w') as outfile: json.dump(value.__dict__, outfile, ensure_ascii=False, indent=4) #json.dump(cur_result, outfile) #cv2.imwrite(osp.join(save_subfolder, '{}.png'.format(img_name)), output) ################################################################################### # STAGE REDS reds4_results = {"000": {}, "011": {}, "015": {}, "020": {}} data_mode = 'sharp_bicubic' N_model_default = 5 for N_in in range(1, N_model_default + 1): for stage in range(1,3): flip_test = False if data_mode == 'sharp_bicubic': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_REDS_SR_L.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_SR_Stage2.pth' elif data_mode == 'blur_bicubic': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_REDS_SRblur_L.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_SRblur_Stage2.pth' elif data_mode == 'blur': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_REDS_deblur_L.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_deblur_Stage2.pth' elif data_mode == 'blur_comp': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_L.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_Stage2.pth' else: raise NotImplementedError predeblur, HR_in = False, False back_RBs = 40 if data_mode == 'blur_bicubic': predeblur = True if data_mode == 'blur' or data_mode == 'blur_comp': predeblur, HR_in = True, True if stage == 2: HR_in = True back_RBs = 20 if stage == 1: test_dataset_folder = '../datasets/REDS4/{}'.format(data_mode) else: test_dataset_folder = '../results/REDS-EDVR_REDS_SR_L_flipx4' print('You should modify the test_dataset_folder path for stage 2') GT_dataset_folder = '../datasets/REDS4/GT' raw_model = EDVR_arch.EDVR(128, N_model_default, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in) model = EDVR_arch.EDVR(128, N_in, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in) crop_border = 0 border_frame = N_in // 2 # border frames when evaluate # temporal padding mode if data_mode == 'Vid4' or data_mode == 'sharp_bicubic': padding = 'new_info' else: padding = 'replicate' save_imgs = True data_mode_t = copy.deepcopy(data_mode) if stage == 1 and data_mode_t != 'Vid4': data_mode = 'REDS-EDVR_REDS_SR_L_flipx4' save_folder = '../results/{}'.format(data_mode) data_mode = copy.deepcopy(data_mode_t) util.mkdirs(save_folder) util.setup_logger('base', save_folder, 'test', level=logging.INFO, screen=True, tofile=True) aposterior_GT_dataset_folder = '../datasets/REDS4/GT_5' crop_border = 0 border_frame = N_in // 2 # border frames when evaluate raw_model.load_state_dict(torch.load(model_path), strict=True) model.nf = raw_model.nf model.center = N_in // 2 # if center is None else center model.is_predeblur = raw_model.is_predeblur model.HR_in = raw_model.HR_in model.w_TSA = raw_model.w_TSA if model.is_predeblur: model.pre_deblur = raw_model.pre_deblur # Predeblur_ResNet_Pyramid(nf=nf, HR_in=self.HR_in) model.conv_1x1 = raw_model.conv_1x1 # nn.Conv2d(nf, nf, 1, 1, bias=True) else: if model.HR_in: model.conv_first_1 = raw_model.conv_first_1 # nn.Conv2d(3, nf, 3, 1, 1, bias=True) model.conv_first_2 = raw_model.conv_first_2 # nn.Conv2d(nf, nf, 3, 2, 1, bias=True) model.conv_first_3 = raw_model.conv_first_3 # nn.Conv2d(nf, nf, 3, 2, 1, bias=True) else: model.conv_first = raw_model.conv_first # nn.Conv2d(3, nf, 3, 1, 1, bias=True) model.feature_extraction = raw_model.feature_extraction # arch_util.make_layer(ResidualBlock_noBN_f, front_RBs) model.fea_L2_conv1 = raw_model.fea_L2_conv1 # nn.Conv2d(nf, nf, 3, 2, 1, bias=True) model.fea_L2_conv2 = raw_model.fea_L2_conv2 # nn.Conv2d(nf, nf, 3, 1, 1, bias=True) model.fea_L3_conv1 = raw_model.fea_L3_conv1 # nn.Conv2d(nf, nf, 3, 2, 1, bias=True) model.fea_L3_conv2 = raw_model.fea_L3_conv2 # nn.Conv2d(nf, nf, 3, 1, 1, bias=True) model.pcd_align = raw_model.pcd_align # PCD_Align(nf=nf, groups=groups) model.tsa_fusion.center = model.center model.tsa_fusion.tAtt_1 = raw_model.tsa_fusion.tAtt_1 model.tsa_fusion.tAtt_2 = raw_model.tsa_fusion.tAtt_2 model.tsa_fusion.fea_fusion = copy.deepcopy(raw_model.tsa_fusion.fea_fusion) model.tsa_fusion.fea_fusion.weight = copy.deepcopy(torch.nn.Parameter(raw_model.tsa_fusion.fea_fusion.weight[:, 0:N_in * 128, :, :])) model.tsa_fusion.sAtt_1 = copy.deepcopy(raw_model.tsa_fusion.sAtt_1) model.tsa_fusion.sAtt_1.weight = copy.deepcopy(torch.nn.Parameter(raw_model.tsa_fusion.sAtt_1.weight[:, 0:N_in * 128, :, :])) model.tsa_fusion.maxpool = raw_model.tsa_fusion.maxpool model.tsa_fusion.avgpool = raw_model.tsa_fusion.avgpool model.tsa_fusion.sAtt_2 = raw_model.tsa_fusion.sAtt_2 model.tsa_fusion.sAtt_3 = raw_model.tsa_fusion.sAtt_3 model.tsa_fusion.sAtt_4 = raw_model.tsa_fusion.sAtt_4 model.tsa_fusion.sAtt_5 = raw_model.tsa_fusion.sAtt_5 model.tsa_fusion.sAtt_L1 = raw_model.tsa_fusion.sAtt_L1 model.tsa_fusion.sAtt_L2 = raw_model.tsa_fusion.sAtt_L2 model.tsa_fusion.sAtt_L3 = raw_model.tsa_fusion.sAtt_L3 model.tsa_fusion.sAtt_add_1 = raw_model.tsa_fusion.sAtt_add_1 model.tsa_fusion.sAtt_add_2 = raw_model.tsa_fusion.sAtt_add_2 model.tsa_fusion.lrelu = raw_model.tsa_fusion.lrelu model.recon_trunk = raw_model.recon_trunk model.upconv1 = raw_model.upconv1 model.upconv2 = raw_model.upconv2 model.pixel_shuffle = raw_model.pixel_shuffle model.HRconv = raw_model.HRconv model.conv_last = raw_model.conv_last model.lrelu = raw_model.lrelu ##################################################### model.eval() model = model.to(device) #avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l = [], [], [] subfolder_name_l = [] subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*'))) subfolder_GT_l = sorted(glob.glob(osp.join(GT_dataset_folder, '*'))) subfolder_GT_a_l = sorted(glob.glob(osp.join(aposterior_GT_dataset_folder, "*"))) # for each subfolder for subfolder, subfolder_GT, subfolder_GT_a in zip(subfolder_l, subfolder_GT_l, subfolder_GT_a_l): subfolder_name = osp.basename(subfolder) subfolder_name_l.append(subfolder_name) save_subfolder = osp.join(save_folder, subfolder_name) img_path_l = sorted(glob.glob(osp.join(subfolder, '*'))) max_idx = len(img_path_l) print("MAX_IDX: ", max_idx) print("SAVE FOLDER::::::", save_folder) if save_imgs: util.mkdirs(save_subfolder) #### read LQ and GT images imgs_LQ = data_util.read_img_seq(subfolder) img_GT_l = [] for img_GT_path in sorted(glob.glob(osp.join(subfolder_GT, '*'))): img_GT_l.append(data_util.read_img(None, img_GT_path)) img_GT_a = [] for img_GT_a_path in sorted(glob.glob(osp.join(subfolder_GT_a, '*'))): img_GT_a.append(data_util.read_img(None, img_GT_a_path)) #avg_psnr, avg_psnr_border, avg_psnr_center, N_border, N_center = 0, 0, 0, 0, 0 # process each image for img_idx, img_path in enumerate(img_path_l): img_name = osp.splitext(osp.basename(img_path))[0] select_idx = data_util.index_generation(img_idx, max_idx, N_in, padding=padding) imgs_in = imgs_LQ.index_select(0, torch.LongTensor(select_idx)).unsqueeze(0).to(device) if flip_test: output = util.flipx4_forward(model, imgs_in) else: print("IMGS_IN SHAPE: ", imgs_in.shape) output = util.single_forward(model, imgs_in) output = util.tensor2img(output.squeeze(0)) if save_imgs and stage == 1: cv2.imwrite(osp.join(save_subfolder, '{}.png'.format(img_name)), output) # calculate PSNR if stage == 2: output = output / 255. GT = np.copy(img_GT_l[img_idx]) # For REDS, evaluate on RGB channels; for Vid4, evaluate on the Y channel #if data_mode == 'Vid4': # bgr2y, [0, 1] GT_a = np.copy(img_GT_a[img_idx]) output_a = copy.deepcopy(output) output, GT = util.crop_border([output, GT], crop_border) crt_psnr = util.calculate_psnr(output * 255, GT * 255) crt_ssim = util.calculate_ssim(output * 255, GT * 255) output_a, GT_a = util.crop_border([output_a, GT_a], crop_border) crt_aposterior = util.calculate_ssim(output_a * 255, GT_a * 255) # CHANGE t = reds4_results[subfolder_name].get(str(img_name)) if t != None: reds4_results[subfolder_name][img_name].add_psnr(crt_psnr) reds4_results[subfolder_name][img_name].add_gt_ssim(crt_ssim) reds4_results[subfolder_name][img_name].add_aposterior_ssim(crt_aposterior) else: reds4_results[subfolder_name].update({img_name: metrics_file(img_name)}) reds4_results[subfolder_name][img_name].add_psnr(crt_psnr) reds4_results[subfolder_name][img_name].add_gt_ssim(crt_ssim) reds4_results[subfolder_name][img_name].add_aposterior_ssim(crt_aposterior) ############################################################################ #### model #### writing reds4 results util.mkdirs('../results/000') util.mkdirs('../results/011') util.mkdirs('../results/015') util.mkdirs('../results/020') save_folder = '../results/' for i, dir_name in enumerate(["000", "011", "015", "020"]): # + save_subfolder = osp.join(save_folder, dir_name) for j, value in reds4_results[dir_name].items(): # cur_result = json.dumps(value.__dict__) with open(osp.join(save_subfolder, '{}.json'.format(value.name)), 'w') as outfile: json.dump(value.__dict__, outfile, ensure_ascii=False, indent=4)
def main(): #################### # arguments parser # #################### # [format] dataset(vid4, REDS4) N(number of frames) parser = argparse.ArgumentParser() parser.add_argument('dataset') parser.add_argument('n_frames') parser.add_argument('stage') args = parser.parse_args() data_mode = str(args.dataset) N_in = int(args.n_frames) stage = int(args.stage) #if args.command == 'start': # start(int(args.params[0])) #elif args.command == 'stop': # stop(args.params[0], int(args.params[1])) #elif args.command == 'stop_all': # stop_all(args.params[0]) ################# # configurations ################# device = torch.device('cuda') os.environ['CUDA_VISIBLE_DEVICES'] = '0' #data_mode = 'Vid4' # Vid4 | sharp_bicubic | blur_bicubic | blur | blur_comp # Vid4: SR # REDS4: sharp_bicubic (SR-clean), blur_bicubic (SR-blur); # blur (deblur-clean), blur_comp (deblur-compression). #stage = 1 # 1 or 2, use two stage strategy for REDS dataset. flip_test = False ############################################################################ #### model if data_mode == 'Vid4': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_Vimeo90K_SR_L.pth' else: raise ValueError('Vid4 does not support stage 2.') elif data_mode == 'sharp_bicubic': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_REDS_SR_L.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_SR_Stage2.pth' elif data_mode == 'blur_bicubic': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_REDS_SRblur_L.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_SRblur_Stage2.pth' elif data_mode == 'blur': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_REDS_deblur_L.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_deblur_Stage2.pth' elif data_mode == 'blur_comp': if stage == 1: model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_L.pth' else: model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_Stage2.pth' else: raise NotImplementedError predeblur, HR_in = False, False back_RBs = 40 if data_mode == 'blur_bicubic': predeblur = True if data_mode == 'blur' or data_mode == 'blur_comp': predeblur, HR_in = True, True if stage == 2: HR_in = True back_RBs = 20 #### dataset if data_mode == 'Vid4': N_model_default = 7 test_dataset_folder = '../datasets/Vid4/BIx4' GT_dataset_folder = '../datasets/Vid4/GT' else: N_model_default = 5 if stage == 1: test_dataset_folder = '../datasets/REDS4/{}'.format(data_mode) else: test_dataset_folder = '../results/REDS-EDVR_REDS_SR_L_flipx4' print('You should modify the test_dataset_folder path for stage 2') GT_dataset_folder = '../datasets/REDS4/GT' raw_model = EDVR_arch.EDVR(128, N_model_default, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in) model = EDVR_arch.EDVR(128, N_in, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in) #### evaluation crop_border = 0 border_frame = N_in // 2 # border frames when evaluate # temporal padding mode if data_mode == 'Vid4' or data_mode == 'sharp_bicubic': padding = 'new_info' else: padding = 'replicate' save_imgs = True data_mode_t = copy.deepcopy(data_mode) if stage == 1 and data_mode_t != 'Vid4': data_mode = 'REDS-EDVR_REDS_SR_L_flipx4' save_folder = '../results/{}'.format(data_mode) data_mode = copy.deepcopy(data_mode_t) util.mkdirs(save_folder) util.setup_logger('base', save_folder, 'test', level=logging.INFO, screen=True, tofile=True) logger = logging.getLogger('base') #### log info logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder)) logger.info('Padding mode: {}'.format(padding)) logger.info('Model path: {}'.format(model_path)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Flip test: {}'.format(flip_test)) #### set up the models print([a for a in dir(model) if not callable(getattr(model, a))]) # not a.startswith('__') and #model.load_state_dict(torch.load(model_path), strict=True) raw_model.load_state_dict(torch.load(model_path), strict=True) # model.load_state_dict(torch.load(model_path), strict=True) #### change model so it can work with less input model.nf = raw_model.nf model.center = N_in // 2 # if center is None else center model.is_predeblur = raw_model.is_predeblur model.HR_in = raw_model.HR_in model.w_TSA = raw_model.w_TSA #ResidualBlock_noBN_f = functools.partial(arch_util.ResidualBlock_noBN, nf=nf) #### extract features (for each frame) if model.is_predeblur: model.pre_deblur = raw_model.pre_deblur #Predeblur_ResNet_Pyramid(nf=nf, HR_in=self.HR_in) model.conv_1x1 = raw_model.conv_1x1 #nn.Conv2d(nf, nf, 1, 1, bias=True) else: if model.HR_in: model.conv_first_1 = raw_model.conv_first_1 #nn.Conv2d(3, nf, 3, 1, 1, bias=True) model.conv_first_2 = raw_model.conv_first_2 #nn.Conv2d(nf, nf, 3, 2, 1, bias=True) model.conv_first_3 = raw_model.conv_first_3 #nn.Conv2d(nf, nf, 3, 2, 1, bias=True) else: model.conv_first = raw_model.conv_first # nn.Conv2d(3, nf, 3, 1, 1, bias=True) model.feature_extraction = raw_model.feature_extraction # arch_util.make_layer(ResidualBlock_noBN_f, front_RBs) model.fea_L2_conv1 = raw_model.fea_L2_conv1 #nn.Conv2d(nf, nf, 3, 2, 1, bias=True) model.fea_L2_conv2 = raw_model.fea_L2_conv2 #nn.Conv2d(nf, nf, 3, 1, 1, bias=True) model.fea_L3_conv1 = raw_model.fea_L3_conv1 #nn.Conv2d(nf, nf, 3, 2, 1, bias=True) model.fea_L3_conv2 = raw_model.fea_L3_conv2 #nn.Conv2d(nf, nf, 3, 1, 1, bias=True) model.pcd_align = raw_model.pcd_align #PCD_Align(nf=nf, groups=groups) ######## Resize TSA model.tsa_fusion.center = model.center # temporal attention (before fusion conv) model.tsa_fusion.tAtt_1 = raw_model.tsa_fusion.tAtt_1 model.tsa_fusion.tAtt_2 = raw_model.tsa_fusion.tAtt_2 # fusion conv: using 1x1 to save parameters and computation #print(raw_model.tsa_fusion.fea_fusion.weight.shape) #print(raw_model.tsa_fusion.fea_fusion.weight.shape) #print(raw_model.tsa_fusion.fea_fusion.weight[127][639].shape) #print("MAIN SHAPE(FEA): ", raw_model.tsa_fusion.fea_fusion.weight.shape) model.tsa_fusion.fea_fusion = copy.deepcopy( raw_model.tsa_fusion.fea_fusion) model.tsa_fusion.fea_fusion.weight = copy.deepcopy( torch.nn.Parameter(raw_model.tsa_fusion.fea_fusion.weight[:, 0:N_in * 128, :, :])) #[:][] #nn.Conv2d(nframes * nf, nf, 1, 1, bias=True) #model.tsa_fusion.fea_fusion.bias = raw_model.tsa_fusion.fea_fusion.bias # spatial attention (after fusion conv) model.tsa_fusion.sAtt_1 = copy.deepcopy(raw_model.tsa_fusion.sAtt_1) model.tsa_fusion.sAtt_1.weight = copy.deepcopy( torch.nn.Parameter(raw_model.tsa_fusion.sAtt_1.weight[:, 0:N_in * 128, :, :])) #[:][] #nn.Conv2d(nframes * nf, nf, 1, 1, bias=True) #model.tsa_fusion.sAtt_1.bias = raw_model.tsa_fusion.sAtt_1.bias #print(N_in * 128) #print(raw_model.tsa_fusion.fea_fusion.weight[:, 0:N_in * 128, :, :].shape) print("MODEL TSA SHAPE: ", model.tsa_fusion.fea_fusion.weight.shape) model.tsa_fusion.maxpool = raw_model.tsa_fusion.maxpool model.tsa_fusion.avgpool = raw_model.tsa_fusion.avgpool model.tsa_fusion.sAtt_2 = raw_model.tsa_fusion.sAtt_2 model.tsa_fusion.sAtt_3 = raw_model.tsa_fusion.sAtt_3 model.tsa_fusion.sAtt_4 = raw_model.tsa_fusion.sAtt_4 model.tsa_fusion.sAtt_5 = raw_model.tsa_fusion.sAtt_5 model.tsa_fusion.sAtt_L1 = raw_model.tsa_fusion.sAtt_L1 model.tsa_fusion.sAtt_L2 = raw_model.tsa_fusion.sAtt_L2 model.tsa_fusion.sAtt_L3 = raw_model.tsa_fusion.sAtt_L3 model.tsa_fusion.sAtt_add_1 = raw_model.tsa_fusion.sAtt_add_1 model.tsa_fusion.sAtt_add_2 = raw_model.tsa_fusion.sAtt_add_2 model.tsa_fusion.lrelu = raw_model.tsa_fusion.lrelu #if model.w_TSA: # model.tsa_fusion = raw_model.tsa_fusion[:][:128 * N_in][:][:] #TSA_Fusion(nf=nf, nframes=nframes, center=self.center) #else: # model.tsa_fusion = raw_model.tsa_fusion[:][:128 * N_in][:][:] #nn.Conv2d(nframes * nf, nf, 1, 1, bias=True) # print(self.tsa_fusion) #### reconstruction model.recon_trunk = raw_model.recon_trunk # arch_util.make_layer(ResidualBlock_noBN_f, back_RBs) #### upsampling model.upconv1 = raw_model.upconv1 #nn.Conv2d(nf, nf * 4, 3, 1, 1, bias=True) model.upconv2 = raw_model.upconv2 #nn.Conv2d(nf, 64 * 4, 3, 1, 1, bias=True) model.pixel_shuffle = raw_model.pixel_shuffle # nn.PixelShuffle(2) model.HRconv = raw_model.HRconv model.conv_last = raw_model.conv_last #### activation function model.lrelu = raw_model.lrelu ##################################################### model.eval() model = model.to(device) avg_ssim_l, avg_ssim_center_l, avg_ssim_border_l = [], [], [] subfolder_name_l = [] subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*'))) subfolder_GT_l = sorted(glob.glob(osp.join(GT_dataset_folder, '*'))) # for each subfolder for subfolder, subfolder_GT in zip(subfolder_l, subfolder_GT_l): subfolder_name = osp.basename(subfolder) subfolder_name_l.append(subfolder_name) save_subfolder = osp.join(save_folder, subfolder_name) img_path_l = sorted(glob.glob(osp.join(subfolder, '*'))) max_idx = len(img_path_l) print("MAX_IDX: ", max_idx) if save_imgs: util.mkdirs(save_subfolder) #### read LQ and GT images imgs_LQ = data_util.read_img_seq(subfolder) img_GT_l = [] for img_GT_path in sorted(glob.glob(osp.join(subfolder_GT, '*'))): img_GT_l.append(data_util.read_img(None, img_GT_path)) avg_ssim, avg_ssim_border, avg_ssim_center, N_border, N_center = 0, 0, 0, 0, 0 # process each image for img_idx, img_path in enumerate(img_path_l): img_name = osp.splitext(osp.basename(img_path))[0] if data_mode == "blur": select_idx = data_util.glarefree_index_generation( img_idx, max_idx, N_in, padding=padding) else: select_idx = data_util.index_generation( img_idx, max_idx, N_in, padding=padding) # HERE GOTCHA print("SELECT IDX: ", select_idx) imgs_in = imgs_LQ.index_select( 0, torch.LongTensor(select_idx)).unsqueeze(0).to(device) if flip_test: output = util.flipx4_forward(model, imgs_in) else: print("IMGS_IN SHAPE: ", imgs_in.shape) # check this output = util.single_forward(model, imgs_in) # error here 1 output = util.tensor2img(output.squeeze(0)) if save_imgs: cv2.imwrite( osp.join(save_subfolder, '{}.png'.format(img_name)), output) # calculate SSIM output = output / 255. GT = np.copy(img_GT_l[img_idx]) # For REDS, evaluate on RGB channels; for Vid4, evaluate on the Y channel if data_mode == 'Vid4': # bgr2y, [0, 1] GT = data_util.bgr2ycbcr(GT, only_y=True) output = data_util.bgr2ycbcr(output, only_y=True) output, GT = util.crop_border([output, GT], crop_border) crt_ssim = util.calculate_ssim(output * 255, GT * 255) logger.info('{:3d} - {:25} \tSSIM: {:.6f} dB'.format( img_idx + 1, img_name, crt_ssim)) if img_idx >= border_frame and img_idx < max_idx - border_frame: # center frames avg_ssim_center += crt_ssim N_center += 1 else: # border frames avg_ssim_border += crt_ssim N_border += 1 avg_ssim = (avg_ssim_center + avg_ssim_border) / (N_center + N_border) avg_ssim_center = avg_ssim_center / N_center avg_ssim_border = 0 if N_border == 0 else avg_ssim_border / N_border avg_ssim_l.append(avg_ssim) avg_ssim_center_l.append(avg_ssim_center) avg_ssim_border_l.append(avg_ssim_border) logger.info('Folder {} - Average SSIM: {:.6f} dB for {} frames; ' 'Center SSIM: {:.6f} dB for {} frames; ' 'Border SSIM: {:.6f} dB for {} frames.'.format( subfolder_name, avg_ssim, (N_center + N_border), avg_ssim_center, N_center, avg_ssim_border, N_border)) logger.info('################ Tidy Outputs ################') for subfolder_name, ssim, ssim_center, ssim_border in zip( subfolder_name_l, avg_ssim_l, avg_ssim_center_l, avg_ssim_border_l): logger.info('Folder {} - Average SSIM: {:.6f} dB. ' 'Center SSIM: {:.6f} dB. ' 'Border SSIM: {:.6f} dB.'.format(subfolder_name, ssim, ssim_center, ssim_border)) logger.info('################ Final Results ################') logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder)) logger.info('Padding mode: {}'.format(padding)) logger.info('Model path: {}'.format(model_path)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Flip test: {}'.format(flip_test)) logger.info('Total Average SSIM: {:.6f} dB for {} clips. ' 'Center SSIM: {:.6f} dB. Border SSIM: {:.6f} dB.'.format( sum(avg_ssim_l) / len(avg_ssim_l), len(subfolder_l), sum(avg_ssim_center_l) / len(avg_ssim_center_l), sum(avg_ssim_border_l) / len(avg_ssim_border_l)))
def main(): # options parser = argparse.ArgumentParser() parser.add_argument('-opt', type=str, required=True, help='Path to option JSON file.') opt = option.parse(parser.parse_args().opt, is_train=True) # Convert to NoneDict, which return None for missing key. opt = option.dict_to_nonedict(opt) pytorch_ver = get_pytorch_ver() # train from scratch OR resume training if opt['path']['resume_state']: if os.path.isdir(opt['path']['resume_state']): import glob resume_state_path = util.sorted_nicely( glob.glob( os.path.normpath(opt['path']['resume_state']) + '/*.state'))[-1] else: resume_state_path = opt['path']['resume_state'] resume_state = torch.load(resume_state_path) else: # training from scratch resume_state = None # rename old folder if exists util.mkdir_and_rename(opt['path']['experiments_root']) util.mkdirs((path for key, path in opt['path'].items() if not key == 'experiments_root' and 'pretrain_model' not in key and 'resume' not in key)) # config loggers. Before it, the log will not work util.setup_logger(None, opt['path']['log'], 'train', level=logging.INFO, screen=True) util.setup_logger('val', opt['path']['log'], 'val', level=logging.INFO) logger = logging.getLogger('base') if resume_state: logger.info('Set [resume_state] to ' + resume_state_path) logger.info('Resuming training from epoch: {}, iter: {}.'.format( resume_state['epoch'], resume_state['iter'])) option.check_resume(opt) # check resume options logger.info(option.dict2str(opt)) # tensorboard logger if opt['use_tb_logger'] and 'debug' not in opt['name']: from tensorboardX import SummaryWriter try: # for version tensorboardX >= 1.7 tb_logger = SummaryWriter(logdir='../tb_logger/' + opt['name']) except: # for version tensorboardX < 1.6 tb_logger = SummaryWriter(log_dir='../tb_logger/' + opt['name']) # random seed seed = opt['train']['manual_seed'] if seed is None: seed = random.randint(1, 10000) logger.info('Random seed: {}'.format(seed)) util.set_random_seed(seed) torch.backends.cudnn.benckmark = True # torch.backends.cudnn.deterministic = True # create train and val dataloader for phase, dataset_opt in opt['datasets'].items(): if phase == 'train': train_set = create_dataset(dataset_opt) train_size = int( math.ceil(len(train_set) / dataset_opt['batch_size'])) logger.info('Number of train images: {:,d}, iters: {:,d}'.format( len(train_set), train_size)) total_iters = int(opt['train']['niter']) total_epochs = int(math.ceil(total_iters / train_size)) logger.info('Total epochs needed: {:d} for iters {:,d}'.format( total_epochs, total_iters)) train_loader = create_dataloader(train_set, dataset_opt) elif phase == 'val': val_set = create_dataset(dataset_opt) val_loader = create_dataloader(val_set, dataset_opt) logger.info('Number of val images in [{:s}]: {:d}'.format( dataset_opt['name'], len(val_set))) else: raise NotImplementedError( 'Phase [{:s}] is not recognized.'.format(phase)) assert train_loader is not None # create model model = create_model(opt) # resume training if resume_state: start_epoch = resume_state['epoch'] current_step = resume_state['iter'] model.resume_training(resume_state) # handle optimizers and schedulers # updated schedulers in case JSON configuration has changed model.update_schedulers(opt['train']) else: current_step = 0 start_epoch = 0 # training logger.info('Start training from epoch: {:d}, iter: {:d}'.format( start_epoch, current_step)) for epoch in range(start_epoch, total_epochs): for n, train_data in enumerate(train_loader, start=1): current_step += 1 if current_step > total_iters: break if pytorch_ver == "pre": # Order for PyTorch ver < 1.1.0 # update learning rate model.update_learning_rate(current_step - 1) # training model.feed_data(train_data) model.optimize_parameters(current_step) elif pytorch_ver == "post": # Order for PyTorch ver > 1.1.0 # training model.feed_data(train_data) model.optimize_parameters(current_step) # update learning rate model.update_learning_rate(current_step - 1) else: print('Error identifying PyTorch version. ', torch.__version__) break # log if current_step % opt['logger']['print_freq'] == 0: logs = model.get_current_log() message = '<epoch:{:3d}, iter:{:8,d}, lr:{:.3e}> '.format( epoch, current_step, model.get_current_learning_rate()) for k, v in logs.items(): message += '{:s}: {:.4e} '.format(k, v) # tensorboard logger if opt['use_tb_logger'] and 'debug' not in opt['name']: tb_logger.add_scalar(k, v, current_step) logger.info(message) # save models and training states (changed to save models before validation) if current_step % opt['logger']['save_checkpoint_freq'] == 0: model.save(current_step) model.save_training_state(epoch + (n >= len(train_loader)), current_step) logger.info('Models and training states saved.') # validation if current_step % opt['train']['val_freq'] == 0: avg_psnr = 0.0 avg_ssim = 0.0 avg_lpips = 0.0 idx = 0 val_sr_imgs_list = [] val_gt_imgs_list = [] for val_data in val_loader: idx += 1 img_name = os.path.splitext( os.path.basename(val_data['LR_path'][0]))[0] img_dir = os.path.join(opt['path']['val_images'], img_name) util.mkdir(img_dir) model.feed_data(val_data) model.test() visuals = model.get_current_visuals() if opt['datasets']['train'][ 'znorm']: # If the image range is [-1,1] sr_img = util.tensor2img(visuals['SR'], min_max=(-1, 1)) # uint8 gt_img = util.tensor2img(visuals['HR'], min_max=(-1, 1)) # uint8 else: # Default: Image range is [0,1] sr_img = util.tensor2img(visuals['SR']) # uint8 gt_img = util.tensor2img(visuals['HR']) # uint8 # sr_img = util.tensor2img(visuals['SR']) # uint8 # gt_img = util.tensor2img(visuals['HR']) # uint8 # print("Min. SR value:",sr_img.min()) # Debug # print("Max. SR value:",sr_img.max()) # Debug # print("Min. GT value:",gt_img.min()) # Debug # print("Max. GT value:",gt_img.max()) # Debug # Save SR images for reference save_img_path = os.path.join( img_dir, '{:s}_{:d}.png'.format(img_name, current_step)) util.save_img(sr_img, save_img_path) # calculate PSNR, SSIM and LPIPS distance crop_size = opt['scale'] gt_img = gt_img / 255. sr_img = sr_img / 255. # For training models with only one channel ndim==2, if RGB ndim==3, etc. if gt_img.ndim == 2: cropped_gt_img = gt_img[crop_size:-crop_size, crop_size:-crop_size] else: cropped_gt_img = gt_img[crop_size:-crop_size, crop_size:-crop_size, :] if sr_img.ndim == 2: cropped_sr_img = sr_img[crop_size:-crop_size, crop_size:-crop_size] else: # Default: RGB images cropped_sr_img = sr_img[crop_size:-crop_size, crop_size:-crop_size, :] # If calculating only once for all images val_gt_imgs_list.append(cropped_gt_img) # If calculating only once for all images val_sr_imgs_list.append(cropped_sr_img) # LPIPS only works for RGB images avg_psnr += util.calculate_psnr(cropped_sr_img * 255, cropped_gt_img * 255) avg_ssim += util.calculate_ssim(cropped_sr_img * 255, cropped_gt_img * 255) # If calculating for each image # avg_lpips += lpips.calculate_lpips([cropped_sr_img], [cropped_gt_img]) avg_psnr = avg_psnr / idx avg_ssim = avg_ssim / idx # avg_lpips=avg_lpips / idx # If calculating for each image # If calculating only once for all images avg_lpips = lpips.calculate_lpips(val_sr_imgs_list, val_gt_imgs_list) # log # logger.info('# Validation # PSNR: {:.5g}, SSIM: {:.5g}'.format(avg_psnr, avg_ssim)) logger.info( '# Validation # PSNR: {:.5g}, SSIM: {:.5g}, LPIPS: {:.5g}'. format(avg_psnr, avg_ssim, avg_lpips)) logger_val = logging.getLogger('val') # validation logger # logger_val.info('<epoch:{:3d}, iter:{:8,d}> psnr: {:.5g}, ssim: {:.5g}'.format( # epoch, current_step, avg_psnr, avg_ssim)) logger_val.info( '<epoch:{:3d}, iter:{:8,d}> psnr: {:.5g}, ssim: {:.5g}, lpips: {:.5g}' .format(epoch, current_step, avg_psnr, avg_ssim, avg_lpips)) # tensorboard logger if opt['use_tb_logger'] and 'debug' not in opt['name']: tb_logger.add_scalar('psnr', avg_psnr, current_step) tb_logger.add_scalar('ssim', avg_ssim, current_step) tb_logger.add_scalar('lpips', avg_lpips, current_step) logger.info('Saving the final model.') model.save('latest') logger.info('End of training.')
def main(): # options parser = argparse.ArgumentParser() parser.add_argument('-opt', type=str, required=True, help='Path to options JSON file.') opt = option.parse(parser.parse_args().opt, is_train=False) util.mkdirs((path for key, path in opt['path'].items() if not key == 'pretrain_model_G')) opt = option.dict_to_nonedict(opt) util.setup_logger(None, opt['path']['log'], 'test.log', level=logging.INFO, screen=True) logger = logging.getLogger('base') logger.info(option.dict2str(opt)) scale = opt.get('scale', 4) # Create test dataset and dataloader test_loaders = [] znorm = False #TMP # znorm_list = [] ''' video_list = os.listdir(cfg.testset_dir) for idx_video in range(len(video_list)): video_name = video_list[idx_video] # dataloader test_set = TestsetLoader(cfg, video_name) test_loader = DataLoader(test_set, num_workers=1, batch_size=1, shuffle=False) ''' for phase, dataset_opt in sorted(opt['datasets'].items()): test_set = create_dataset(dataset_opt) test_loader = create_dataloader(test_set, dataset_opt) logger.info('Number of test images in [{:s}]: {:d}'.format( dataset_opt['name'], len(test_set))) test_loaders.append(test_loader) # Temporary, will turn znorm on for all the datasets. Will need to introduce a variable for each dataset and differentiate each one later in the loop. # if dataset_opt.get['znorm'] and znorm == False: # znorm = True znorm = dataset_opt.get('znorm', False) # znorm_list.apped(znorm) # Create model model = create_model(opt) for test_loader in test_loaders: test_set_name = test_loader.dataset.opt['name'] logger.info('\nTesting [{:s}]...'.format(test_set_name)) test_start_time = time.time() dataset_dir = os.path.join(opt['path']['results_root'], test_set_name) util.mkdir(dataset_dir) test_results = OrderedDict() test_results['psnr'] = [] test_results['ssim'] = [] test_results['psnr_y'] = [] test_results['ssim_y'] = [] for data in test_loader: need_HR = False if test_loader.dataset.opt[ 'dataroot_HR'] is None else True img_path = data['LR_path'][0] img_name = os.path.splitext(os.path.basename(img_path))[0] # tmp_vis(data['LR'][:,1,:,:,:], True) if opt.get('chop_forward', None): # data if len(data['LR'].size()) == 4: b, n_frames, h_lr, w_lr = data['LR'].size() LR_y_cube = data['LR'].view(b, -1, 1, h_lr, w_lr) # b, t, c, h, w elif len(data['LR'].size() ) == 5: #for networks that work with 3 channel images _, n_frames, _, _, _ = data['LR'].size() LR_y_cube = data['LR'] # b, t, c, h, w # print(LR_y_cube.shape) # print(data['LR_bicubic'].shape) # crop borders to ensure each patch can be divisible by 2 #TODO: this is modcrop, not sure if really needed, check (the dataloader already does modcrop) _, _, _, h, w = LR_y_cube.size() h = int(h // 16) * 16 w = int(w // 16) * 16 LR_y_cube = LR_y_cube[:, :, :, :h, :w] if isinstance(data['LR_bicubic'], torch.Tensor): # SR_cb = data['LR_bicubic'][:, 1, :, :][:, :, :h * scale, :w * scale] SR_cb = data['LR_bicubic'][:, 1, :h * scale, :w * scale] # SR_cr = data['LR_bicubic'][:, 2, :, :][:, :, :h * scale, :w * scale] SR_cr = data['LR_bicubic'][:, 2, :h * scale, :w * scale] SR_y = chop_forward(LR_y_cube, model, scale, need_HR=need_HR).squeeze(0) # SR_y = np.array(SR_y.data.cpu()) if test_loader.dataset.opt.get('srcolors', None): print(SR_y.shape, SR_cb.shape, SR_cr.shape) sr_img = ycbcr_to_rgb(torch.stack((SR_y, SR_cb, SR_cr), -3)) else: sr_img = SR_y else: # data model.feed_data(data, need_HR=need_HR) # SR_y = net(LR_y_cube).squeeze(0) model.test() # test visuals = model.get_current_visuals(need_HR=need_HR) # ds = torch.nn.AvgPool2d(2, stride=2, count_include_pad=False) # tmp_vis(ds(visuals['SR']), True) # tmp_vis(visuals['SR'], True) if test_loader.dataset.opt.get( 'y_only', None) and test_loader.dataset.opt.get( 'srcolors', None): SR_cb = data['LR_bicubic'][:, 1, :, :] SR_cr = data['LR_bicubic'][:, 2, :, :] # tmp_vis(ds(SR_cb), True) # tmp_vis(ds(SR_cr), True) sr_img = ycbcr_to_rgb( torch.stack((visuals['SR'], SR_cb, SR_cr), -3)) else: sr_img = visuals['SR'] #if znorm the image range is [-1,1], Default: Image range is [0,1] # testing, each "dataset" can have a different name (not train, val or other) sr_img = tensor2np(sr_img, denormalize=znorm) # uint8 # save images suffix = opt['suffix'] if suffix: save_img_path = os.path.join(dataset_dir, img_name + suffix + '.png') else: save_img_path = os.path.join(dataset_dir, img_name + '.png') util.save_img(sr_img, save_img_path) #TODO: update to use metrics functions # calculate PSNR and SSIM if need_HR: #if znorm the image range is [-1,1], Default: Image range is [0,1] # testing, each "dataset" can have a different name (not train, val or other) gt_img = tensor2img(visuals['HR'], denormalize=znorm) # uint8 gt_img = gt_img / 255. sr_img = sr_img / 255. crop_border = test_loader.dataset.opt['scale'] cropped_sr_img = sr_img[crop_border:-crop_border, crop_border:-crop_border, :] cropped_gt_img = gt_img[crop_border:-crop_border, crop_border:-crop_border, :] psnr = util.calculate_psnr(cropped_sr_img * 255, cropped_gt_img * 255) ssim = util.calculate_ssim(cropped_sr_img * 255, cropped_gt_img * 255) test_results['psnr'].append(psnr) test_results['ssim'].append(ssim) if gt_img.shape[2] == 3: # RGB image sr_img_y = bgr2ycbcr(sr_img, only_y=True) gt_img_y = bgr2ycbcr(gt_img, only_y=True) cropped_sr_img_y = sr_img_y[crop_border:-crop_border, crop_border:-crop_border] cropped_gt_img_y = gt_img_y[crop_border:-crop_border, crop_border:-crop_border] psnr_y = util.calculate_psnr(cropped_sr_img_y * 255, cropped_gt_img_y * 255) ssim_y = util.calculate_ssim(cropped_sr_img_y * 255, cropped_gt_img_y * 255) test_results['psnr_y'].append(psnr_y) test_results['ssim_y'].append(ssim_y) logger.info('{:20s} - PSNR: {:.6f} dB; SSIM: {:.6f}; PSNR_Y: {:.6f} dB; SSIM_Y: {:.6f}.'\ .format(img_name, psnr, ssim, psnr_y, ssim_y)) else: logger.info( '{:20s} - PSNR: {:.6f} dB; SSIM: {:.6f}.'.format( img_name, psnr, ssim)) else: logger.info(img_name) #TODO: update to use metrics functions if need_HR: # metrics # Average PSNR/SSIM results ave_psnr = sum(test_results['psnr']) / len(test_results['psnr']) ave_ssim = sum(test_results['ssim']) / len(test_results['ssim']) logger.info('----Average PSNR/SSIM results for {}----\n\tPSNR: {:.6f} dB; SSIM: {:.6f}\n'\ .format(test_set_name, ave_psnr, ave_ssim)) if test_results['psnr_y'] and test_results['ssim_y']: ave_psnr_y = sum(test_results['psnr_y']) / len( test_results['psnr_y']) ave_ssim_y = sum(test_results['ssim_y']) / len( test_results['ssim_y']) logger.info('----Y channel, average PSNR/SSIM----\n\tPSNR_Y: {:.6f} dB; SSIM_Y: {:.6f}\n'\ .format(ave_psnr_y, ave_ssim_y))
def main(opts): ################## configurations ################# device = torch.device('cuda') os.environ['CUDA_VISIBLE_DEVICES'] = opts.gpus cache_all_imgs = opts.cache > 0 n_gpus = len(opts.gpus.split(',')) flip_test, save_imgs = False, False scale = 4 N_in, nf = 5, 64 back_RBs = 10 w_TSA = False predeblur, HR_in = False, False crop_border = 0 border_frame = N_in // 2 padding = 'new_info' ################## model files #################### model_dir = opts.model_dir if osp.isfile(model_dir): model_names = [osp.basename(model_dir)] model_dir = osp.dirname(model_dir) elif osp.isdir(model_dir): model_names = [ x for x in os.listdir(model_dir) if str.isdigit(x.split('_')[0]) ] model_names = sorted(model_names, key=lambda x: int(x.split("_")[0])) else: raise IOError('Invalid model_dir: {}'.format(model_dir)) ################## dataset ######################## test_subs = sorted(os.listdir(opts.test_dir)) gt_subs = os.listdir(opts.gt_dir) valid_test_subs = [sub in gt_subs for sub in test_subs] assert (all(valid_test_subs)), 'Invalid sub folders exists in {}'.format( opts.test_dir) scale = float(os.path.basename(os.path.dirname(opts.test_dir))[1:]) if cache_all_imgs: print('Cacheing all testing images ...') all_imgs = {} for sub in test_subs: print('Reading sub-folder: {} ...'.format(sub)) test_sub_dir = osp.join(opts.test_dir, sub) gt_sub_dir = osp.join(opts.gt_dir, sub) all_imgs[sub] = {'test': [], 'gt': []} im_names = sorted(os.listdir(test_sub_dir)) for i, name in enumerate(im_names): test_im_path = osp.join(test_sub_dir, name) gt_im_path = osp.join(gt_sub_dir, name) test_im = cv2.imread(test_im_path, cv2.IMREAD_UNCHANGED)[:, :, (2, 1, 0)] test_im = test_im.astype(np.float32).transpose( (2, 0, 1)) / 255. all_imgs[sub]['test'].append(test_im) gt_im = cv2.imread(gt_im_path, cv2.IMREAD_UNCHANGED).astype(np.float32) all_imgs[sub]['gt'].append(gt_im) all_psnrs = [] for model_name in model_names: model_path = osp.join(model_dir, model_name) exp_name = model_name.split('_')[0] if 'meta' in opts.mode.lower(): model = EDVR_arch.MetaEDVR(nf=nf, nframes=N_in, groups=8, front_RBs=5, center=None, back_RBs=back_RBs, predeblur=predeblur, HR_in=HR_in, w_TSA=w_TSA) elif opts.mode.lower() == 'edvr': model = EDVR_arch.EDVR(nf=nf, nframes=N_in, groups=8, front_RBs=5, center=None, back_RBs=back_RBs, predeblur=predeblur, HR_in=HR_in, w_TSA=w_TSA) elif opts.mode.lower() == 'upedvr': model = EDVR_arch.UPEDVR(nf=nf, nframes=N_in, groups=8, front_RBs=5, center=None, back_RBs=10, w_TSA=w_TSA, down_scale=True, align_target=True, ret_valid=True) elif opts.mode.lower() == 'upcont1': model = EDVR_arch.UPControlEDVR(nf=nf, nframes=N_in, groups=8, front_RBs=5, center=None, back_RBs=10, w_TSA=w_TSA, down_scale=True, align_target=True, ret_valid=True, multi_scale_cont=False) elif opts.mode.lower() == 'upcont3': model = EDVR_arch.UPControlEDVR(nf=nf, nframes=N_in, groups=8, front_RBs=5, center=None, back_RBs=10, w_TSA=w_TSA, down_scale=True, align_target=True, ret_valid=True, multi_scale_cont=True) elif opts.mode.lower() == 'upcont2': model = EDVR_arch.UPControlEDVR(nf=nf, nframes=N_in, groups=8, front_RBs=5, center=None, back_RBs=10, w_TSA=w_TSA, down_scale=True, align_target=True, ret_valid=True, multi_scale_cont=True) else: raise TypeError('Unknown model mode: {}'.format(opts.mode)) save_folder = osp.join(opts.save_dir, exp_name) util.mkdirs(save_folder) util.setup_logger(exp_name, save_folder, 'test', level=logging.INFO, screen=True, tofile=True) logger = logging.getLogger(exp_name) #### log info logger.info('Data: {}'.format(opts.test_dir)) logger.info('Padding mode: {}'.format(padding)) logger.info('Model path: {}'.format(model_path)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Flip test: {}'.format(flip_test)) #### set up the models model.load_state_dict(torch.load(model_path), strict=True) model.eval() if n_gpus > 1: model = nn.DataParallel(model) model = model.to(device) avg_psnrs, avg_psnr_centers, avg_psnr_borders = [], [], [] avg_ssims, avg_ssim_centers, avg_ssim_borders = [], [], [] evaled_subs = [] # for each subfolder for sub in test_subs: evaled_subs.append(sub) test_sub_dir = osp.join(opts.test_dir, sub) gt_sub_dir = osp.join(opts.gt_dir, sub) img_names = sorted(os.listdir(test_sub_dir)) max_idx = len(img_names) if save_imgs: save_subfolder = osp.join(save_folder, sub) util.mkdirs(save_subfolder) #### get LQ and GT images if not cache_all_imgs: img_LQs, img_GTs = [], [] for i, name in enumerate(img_names): test_im_path = osp.join(test_sub_dir, name) gt_im_path = osp.join(gt_sub_dir, name) test_im = cv2.imread(test_im_path, cv2.IMREAD_UNCHANGED)[:, :, (2, 1, 0)] test_im = test_im.astype(np.float32).transpose( (2, 0, 1)) / 255. gt_im = cv2.imread(gt_im_path, cv2.IMREAD_UNCHANGED).astype(np.float32) img_LQs.append(test_im) img_GTs.append(gt_im) else: img_LQs = all_imgs[sub]['test'] img_GTs = all_imgs[sub]['gt'] avg_psnr, avg_psnr_border, avg_psnr_center, N_border, N_center = 0, 0, 0, 0, 0 avg_ssim, avg_ssim_border, avg_ssim_center = 0, 0, 0 # process each image for i in range(0, max_idx, n_gpus): end = min(i + n_gpus, max_idx) select_idxs = [ data_util.index_generation(j, max_idx, N_in, padding=padding) for j in range(i, end) ] imgs = [] for select_idx in select_idxs: im = torch.from_numpy( np.stack([img_LQs[k] for k in select_idx])) imgs.append(im) if (i + n_gpus) > max_idx: for _ in range(max_idx, i + n_gpus): imgs.append(torch.zeros_like(im)) imgs = torch.stack(imgs, 0).to(device) if flip_test: output = util.flipx4_forward(model, imgs) else: if 'meta' in opts.mode.lower(): output = util.meta_single_forward( model, imgs, scale, n_gpus) if 'up' in opts.mode.lower(): output = util.up_single_forward(model, imgs, scale) else: output = util.single_forward(model, imgs) output = [ util.tensor2img(x).astype(np.float32) for x in output ] if save_imgs: for ii in range(i, end): cv2.imwrite( osp.join(save_subfolder, '{}.png'.format(img_names[ii])), output[ii - i].astype(np.uint8)) # calculate PSNR GT = np.copy(img_GTs[i:end]) output = util.crop_border(output, crop_border) GT = util.crop_border(GT, crop_border) for m in range(i, end): crt_psnr = util.calculate_psnr(output[m - i], GT[m - i]) crt_ssim = util.calculate_ssim(output[m - i], GT[m - i]) logger.info( '{:3d} - {:25} \tPSNR: {:.6f} dB SSIM: {:.6}'. format(m + 1, img_names[m], crt_psnr, crt_ssim)) if m >= border_frame and m < max_idx - border_frame: # center frames avg_psnr_center += crt_psnr avg_ssim_center += crt_ssim N_center += 1 else: # border frames avg_psnr_border += crt_psnr avg_ssim_border += crt_ssim N_border += 1 avg_psnr = (avg_psnr_center + avg_psnr_border) / (N_center + N_border) avg_psnr_center = avg_psnr_center / N_center avg_psnr_border = 0 if N_border == 0 else avg_psnr_border / N_border avg_ssim = (avg_ssim_center + avg_ssim_border) / (N_center + N_border) avg_ssim_center = avg_ssim_center / N_center avg_ssim_border = 0 if N_border == 0 else avg_ssim_border / N_border avg_psnrs.append(avg_psnr) avg_psnr_centers.append(avg_psnr_center) avg_psnr_borders.append(avg_psnr_border) avg_ssims.append(avg_ssim) avg_ssim_centers.append(avg_ssim_center) avg_ssim_borders.append(avg_ssim_border) logger.info('Folder {} - Average PSNR: {:.6f} dB for {} frames; ' 'Center PSNR: {:.6f} dB for {} frames; ' 'Border PSNR: {:.6f} dB for {} frames.'.format( sub, avg_psnr, (N_center + N_border), avg_psnr_center, N_center, avg_psnr_border, N_border)) logger.info('Folder {} - Average SSIM: {:.6f} for {} frames; ' 'Center SSIM: {:.6f} for {} frames; ' 'Border SSIM: {:.6f} for {} frames.'.format( sub, avg_ssim, (N_center + N_border), avg_ssim_center, N_center, avg_ssim_border, N_border)) logger.info('################ Tidy Outputs ################') for sub_name, psnr, psnr_center, psnr_border, ssim, ssim_center, ssim_border in zip( evaled_subs, avg_psnrs, avg_psnr_centers, avg_psnr_borders, avg_ssims, avg_ssim_centers, avg_ssim_borders): logger.info( 'Folder {} - Average PSNR: {:.6f} dB. ' 'Center PSNR: {:.6f} dB. Border PSNR: {:.6f} dB.'.format( sub_name, psnr, psnr_center, psnr_border)) logger.info('Folder {} - Average SSIM: {:.6f} ' 'Center SSIM: {:.6f} Border SSIM: {:.6f} '.format( sub_name, ssim, ssim_center, ssim_border)) logger.info('################ Final Results ################') logger.info('Data: {}'.format(opts.test_dir)) logger.info('Padding mode: {}'.format(padding)) logger.info('Model path: {}'.format(model_path)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Flip test: {}'.format(flip_test)) logger.info('Total Average PSNR: {:.6f} dB for {} clips. ' 'Center PSNR: {:.6f} dB. Border PSNR: {:.6f} dB.'.format( sum(avg_psnrs) / len(avg_psnrs), len(test_subs), sum(avg_psnr_centers) / len(avg_psnr_centers), sum(avg_psnr_borders) / len(avg_psnr_borders))) logger.info('Total Average SSIM: {:.6f} for {} clips. ' 'Center SSIM: {:.6f} Border SSIM: {:.6f} '.format( sum(avg_ssims) / len(avg_ssims), len(test_subs), sum(avg_ssim_centers) / len(avg_ssim_centers), sum(avg_ssim_borders) / len(avg_ssim_borders)))
def cal_pnsr_ssim(sr_img, gt_img, lr_img, lrgt_img): # save images suffix = opt['suffix'] if suffix: save_img_path = osp.join(dataset_dir, folder, img_name + suffix + '.png') else: save_img_path = osp.join(dataset_dir, folder, img_name + '.png') util.save_img(sr_img, save_img_path) # # if suffix: # save_img_path = osp.join(dataset_dir, folder, img_name + suffix + '_GT.png') # else: # save_img_path = osp.join(dataset_dir, folder, img_name + '_GT.png') # util.save_img(gt_img, save_img_path) # if suffix: save_img_path = osp.join(dataset_dir, folder, img_name + suffix + '_LR.png') else: save_img_path = osp.join(dataset_dir, folder, img_name + '_LR.png') util.save_img(lr_img, save_img_path) # # if suffix: # save_img_path = osp.join(dataset_dir, folder, img_name + suffix + '_LR_ref.png') # else: # save_img_path = osp.join(dataset_dir, folder, img_name + '_LR_ref.png') # util.save_img(lrgt_img, save_img_path) # calculate PSNR and SSIM gt_img = gt_img / 255. sr_img = sr_img / 255. lr_img = lr_img / 255. lrgt_img = lrgt_img / 255. crop_border = opt['crop_border'] if opt['crop_border'] else opt['scale'] if crop_border == 0: cropped_sr_img = sr_img cropped_gt_img = gt_img else: cropped_sr_img = sr_img[crop_border:-crop_border, crop_border:-crop_border, :] cropped_gt_img = gt_img[crop_border:-crop_border, crop_border:-crop_border, :] psnr = util.calculate_psnr(cropped_sr_img * 255, cropped_gt_img * 255) ssim = util.calculate_ssim(cropped_sr_img * 255, cropped_gt_img * 255) test_results['psnr'].append(psnr) test_results['ssim'].append(ssim) # PSNR and SSIM for LR psnr_lr = util.calculate_psnr(lr_img * 255, lrgt_img * 255) ssim_lr = util.calculate_ssim(lr_img * 255, lrgt_img * 255) test_results['psnr_lr'].append(psnr_lr) test_results['ssim_lr'].append(ssim_lr) if gt_img.shape[2] == 3: # RGB image sr_img_y = bgr2ycbcr(sr_img, only_y=True) gt_img_y = bgr2ycbcr(gt_img, only_y=True) if crop_border == 0: cropped_sr_img_y = sr_img_y cropped_gt_img_y = gt_img_y else: cropped_sr_img_y = sr_img_y[crop_border:-crop_border, crop_border:-crop_border] cropped_gt_img_y = gt_img_y[crop_border:-crop_border, crop_border:-crop_border] psnr_y = util.calculate_psnr(cropped_sr_img_y * 255, cropped_gt_img_y * 255) ssim_y = util.calculate_ssim(cropped_sr_img_y * 255, cropped_gt_img_y * 255) test_results['psnr_y'].append(psnr_y) test_results['ssim_y'].append(ssim_y) lr_img_y = bgr2ycbcr(lr_img, only_y=True) lrgt_img_y = bgr2ycbcr(lrgt_img, only_y=True) psnr_y_lr = util.calculate_psnr(lr_img_y * 255, lrgt_img_y * 255) ssim_y_lr = util.calculate_ssim(lr_img_y * 255, lrgt_img_y * 255) test_results['psnr_y_lr'].append(psnr_y_lr) test_results['ssim_y_lr'].append(ssim_y_lr) writer.writerow([osp.join(folder, img_name), psnr_y, psnr_y_lr, ssim_y, ssim_y_lr]) logger.info( '{:20s} - PSNR: {:.6f} dB; SSIM: {:.6f}; PSNR_Y: {:.6f} dB; SSIM_Y: {:.6f}. LR PSNR: {:.6f} dB; SSIM: {:.6f}; PSNR_Y: {:.6f} dB; SSIM_Y: {:.6f}.'. format(osp.join(folder, img_name), psnr, ssim, psnr_y, ssim_y, psnr_lr, ssim_lr, psnr_y_lr, ssim_y_lr)) else: writer.writerow([osp.join(folder, img_name), psnr, psnr_lr]) logger.info('{:20s} - PSNR: {:.6f} dB; SSIM: {:.6f}. LR PSNR: {:.6f} dB; SSIM: {:.6f}.'.format( osp.join(folder, img_name), psnr, ssim, psnr_lr, ssim_lr)) return test_results
def main(): # options parser = argparse.ArgumentParser() parser.add_argument('-opt', type=str, required=True, help='Path to option JSON file.') opt = option.parse(parser.parse_args().opt, is_train=True) opt = option.dict_to_nonedict( opt) # Convert to NoneDict, which return None for missing key. # train from scratch OR resume training if opt['path']['resume_state']: # resuming training resume_state = torch.load(opt['path']['resume_state']) else: # training from scratch resume_state = None util.mkdir_and_rename( opt['path']['experiments_root']) # rename old folder if exists util.mkdirs((path for key, path in opt['path'].items() if not key == 'experiments_root' and 'pretrain_model' not in key and 'resume' not in key)) # config loggers. Before it, the log will not work util.setup_logger(None, opt['path']['log'], 'train', level=logging.INFO, screen=True) util.setup_logger('val', opt['path']['log'], 'val', level=logging.INFO) logger = logging.getLogger('base') if resume_state: logger.info('Resuming training from epoch: {}, iter: {}.'.format( resume_state['epoch'], resume_state['iter'])) option.check_resume(opt) # check resume options logger.info(option.dict2str(opt)) # tensorboard logger if opt['use_tb_logger'] and 'debug' not in opt['name']: from tensorboardX import SummaryWriter tb_logger = SummaryWriter(log_dir='../tb_logger/' + opt['name']) # random seed seed = opt['train']['manual_seed'] if seed is None: seed = random.randint(1, 10000) logger.info('Random seed: {}'.format(seed)) util.set_random_seed(seed) torch.backends.cudnn.benckmark = True # create train and val dataloader for phase, dataset_opt in opt['datasets'].items(): if phase == 'train': train_set = create_dataset(dataset_opt) train_size = int( math.ceil(len(train_set) / dataset_opt['batch_size'])) logger.info('Number of train images: {:,d}, iters: {:,d}'.format( len(train_set), train_size)) total_iters = int(opt['train']['niter']) total_epochs = int(math.ceil(total_iters / train_size)) logger.info('Total epochs needed: {:d} for iters {:,d}'.format( total_epochs, total_iters)) train_loader = create_dataloader(train_set, dataset_opt) elif phase == 'val': val_set = create_dataset(dataset_opt) val_loader = create_dataloader(val_set, dataset_opt) logger.info('Number of val images in [{:s}]: {:d}'.format( dataset_opt['name'], len(val_set))) else: raise NotImplementedError( 'Phase [{:s}] is not recognized.'.format(phase)) assert train_loader is not None # create model model = create_model(opt) # resume training if resume_state: start_epoch = resume_state['epoch'] current_step = resume_state['iter'] model.resume_training(resume_state) # handle optimizers and schedulers else: current_step = 0 start_epoch = 0 # training logger.info('Start training from epoch: {:d}, iter: {:d}'.format( start_epoch, current_step)) for epoch in range(start_epoch, total_epochs): for _, train_data in enumerate(train_loader): current_step += 1 if current_step > total_iters: break # update learning rate model.update_learning_rate() # training model.feed_data_diffuse(train_data) model.optimize_parameters(current_step) # log if current_step % opt['logger']['print_freq'] == 0: logs = model.get_current_log() message = '<epoch:{:3d}, iter:{:8,d}, lr:{:.3e}> '.format( epoch, current_step, model.get_current_learning_rate()) for k, v in logs.items(): message += '{:s}: {:.4e} '.format(k, v) # tensorboard logger if opt['use_tb_logger'] and 'debug' not in opt['name']: tb_logger.add_scalar(k, v, current_step) logger.info(message) # validation if current_step % opt['train']['val_freq'] == 0: avg_psnr = 0.0 avg_ssim = 0.0 avg_mrse = 0.0 idx = 0 for val_data in val_loader: idx += 1 img_name = os.path.splitext( os.path.basename(val_data['NOISY_path'][0]))[0] img_dir = opt['path'][ 'val_images'] #os.path.join(opt['path']['val_images'], img_name) util.mkdir(img_dir) model.feed_data_diffuse(val_data) model.test() if opt["image_type"] == "exr": y = val_data["x_offset"] x = val_data["y_offset"] visuals = model.get_current_visuals() avg_mrse += util.calculate_mrse( visuals["DENOISED"].numpy(), visuals["GT"].numpy()) lr_img = util.tensor2img(visuals['NOISY']) sr_img = util.tensor2img(visuals['DENOISED']) # uint8 gt_img = util.tensor2img(visuals['GT']) # uint8 ############################################################################################## # sr_img = util.tensor2img(visuals['DENOISED']) # uint8 # lr_img = util.tensor2img(visuals['NOISY']) # gt_img = util.tensor2img(visuals['GT']) # uint8 # if opt["image_type"] == "exr": # sr_img = sr_img[y:1280-y, x:1280-x, :] # lr_img = lr_img[y:1280-y, x:1280-x, :] # gt_img = gt_img[y:1280-y, x:1280-x, :] ############################################################################################## # Save DENOISED images for reference save_DENOISED_img_path = os.path.join( img_dir, '{:s}_{:d}_1denoised.png'.format( img_name, current_step)) save_NOISY_img_path = os.path.join( img_dir, '{:s}_{:d}_0noisy.png'.format(img_name, current_step)) save_GT_img_path = os.path.join( img_dir, '{:s}_{:d}_2gt.png'.format(img_name, current_step)) # if current_step % 10000 == 0 :#and idx%100 ==0: # util.save_img(sr_img, save_DENOISED_img_path) # util.save_img(lr_img, save_NOISY_img_path) # util.save_img(gt_img, save_GT_img_path) # calculate PSNR # crop_size = opt['scale'] gt_img = gt_img #/ 255. sr_img = sr_img #/ 255. # cropped_sr_img = sr_img[crop_size:-crop_size, crop_size:-crop_size, :] # cropped_gt_img = gt_img[crop_size:-crop_size, crop_size:-crop_size, :] # avg_psnr += util.calculate_psnr(sr_img * 255, gt_img * 255) avg_psnr += util.calculate_psnr(sr_img, gt_img) avg_ssim += util.calculate_ssim(sr_img, gt_img) ############################################################################################## if opt["image_type"] == "exr" and current_step % 10000 == 0: sr_exr = util.tensor2exr(visuals['DENOISED']) # uint8 lr_exr = util.tensor2exr(visuals['NOISY']) gt_exr = util.tensor2exr(visuals['GT']) # uint8 # sr_exr = sr_exr[y:1280-y, x:1280-x, :] # lr_exr = lr_exr[y:1280-y, x:1280-x, :] # gt_exr = gt_exr[y:1280-y, x:1280-x, :] save_DENOISED_img_path = os.path.join( img_dir, '{:s}_{:d}_1denoised.exr'.format( img_name, current_step)) save_NOISY_img_path = os.path.join( img_dir, '{:s}_{:d}_0noisy.exr'.format( img_name, current_step)) save_GT_img_path = os.path.join( img_dir, '{:s}_{:d}_2gt.exr'.format(img_name, current_step)) util.saveEXRfromMatrix(save_DENOISED_img_path, sr_exr, (x, y)) util.saveEXRfromMatrix(save_NOISY_img_path, lr_exr, (x, y)) util.saveEXRfromMatrix(save_GT_img_path, gt_exr, (x, y)) ############################################################################################## avg_psnr = avg_psnr / idx avg_ssim = avg_ssim / idx avg_mrse = avg_mrse / idx # log logger.info('# Validation # PSNR: {:.4e}'.format(avg_psnr)) logger.info('# Validation # SSIM: {:.4e}'.format(avg_ssim)) logger.info('# Validation # MRSE: {:.4e}'.format(avg_mrse)) logger_val = logging.getLogger('val') # validation logger logger_val.info( '<epoch:{:3d}, iter:{:8,d}> psnr: {:.4e} ssim: {:.4e} mrse: {:.4e}' .format(epoch, current_step, avg_psnr, avg_ssim, avg_mrse)) # tensorboard logger if opt['use_tb_logger'] and 'debug' not in opt['name']: tb_logger.add_scalar('psnr', avg_psnr, current_step) tb_logger.add_scalar('ssim', avg_ssim, current_step) tb_logger.add_scalar('mrse', avg_mrse, current_step) # save models and training states if current_step % opt['logger']['save_checkpoint_freq'] == 0: logger.info('Saving models and training states.') model.save(current_step) model.save_training_state(epoch, current_step) logger.info('Saving the final model.') model.save('latest') logger.info('End of training.')
def main(): ymlPath = './options/df2k/train_bicubic_noise.yml' parser = argparse.ArgumentParser() parser.add_argument('-opt', type=str, default=ymlPath, help='Path to option YMAL file.') parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none', help='job launcher') parser.add_argument('--local_rank', type=int, default=0) args = parser.parse_args() opt = option.parse(args.opt, is_train=True) if args.launcher == 'none': # disabled distributed training opt['dist'] = False rank = -1 print('Disabled distributed training.') else: opt['dist'] = True init_dist() world_size = torch.distributed.get_world_size() rank = torch.distributed.get_rank() if opt['path'].get('resume_state', None): device_id = torch.cuda.current_device() resume_state = torch.load(opt['path']['resume_state'], map_location=lambda storage, loc: storage.cuda(device_id)) option.check_resume(opt, resume_state['iter']) # check resume options print('resume_state=', opt['path']['resume_state']) else: resume_state = None if rank <= 0: if resume_state is None: util.mkdir_and_rename( opt['path']['experiments_root']) # rename experiment folder if exists util.mkdirs((path for key, path in opt['path'].items() if not key == 'experiments_root' and 'pretrain_model' not in key and 'resume' not in key)) util.setup_logger('base', opt['path']['log'], 'train_' + opt['name'], level=logging.INFO, screen=True, tofile=True) util.setup_logger('val', opt['path']['log'], 'val_' + opt['name'], level=logging.INFO, screen=True, tofile=True) logger = logging.getLogger('base') logger.info(option.dict2str(opt)) if opt['use_tb_logger'] and 'debug' not in opt['name']: version = float(torch.__version__[0:3]) if version >= 1.1: # PyTorch 1.1 from torch.utils.tensorboard import SummaryWriter else: logger.info( 'You are using PyTorch {}. Tensorboard will use [tensorboardX]'.format(version)) from tensorboardX import SummaryWriter tb_logger = SummaryWriter(log_dir='../tb_logger/' + opt['name']) else: util.setup_logger('base', opt['path']['log'], 'train', level=logging.INFO, screen=True) logger = logging.getLogger('base') opt = option.dict_to_nonedict(opt) seed = opt['train']['manual_seed'] if seed is None: seed = random.randint(1, 10000) if rank <= 0: logger.info('Random seed: {}'.format(seed)) util.set_random_seed(seed) torch.backends.cudnn.benckmark = True dataset_ratio = 1 # enlarge the size of each epoch train_loader = None val_loader = None for phase, dataset_opt in opt['datasets'].items(): if phase == 'train': train_set = create_dataset(dataset_opt) train_size = int(math.ceil(len(train_set) / dataset_opt['batch_size'])) total_iters = int(opt['train']['niter']) total_epochs = int(math.ceil(total_iters / train_size)) if opt['dist']: train_sampler = DistIterSampler(train_set, world_size, rank, dataset_ratio) total_epochs = int(math.ceil(total_iters / (train_size * dataset_ratio))) else: train_sampler = None train_loader = create_dataloader(train_set, dataset_opt, opt, train_sampler) if rank <= 0: logger.info('Number of train images: {:,d}, iters: {:,d}'.format( len(train_set), train_size)) logger.info('Total epochs needed: {:d} for iters {:,d}'.format( total_epochs, total_iters)) elif phase == 'val': val_set = create_dataset(dataset_opt) val_loader = create_dataloader(val_set, dataset_opt, opt, None) if rank <= 0: logger.info('Number of val images in [{:s}]: {:d}'.format( dataset_opt['name'], len(val_set))) else: raise NotImplementedError('Phase [{:s}] is not recognized.'.format(phase)) assert train_loader is not None model = create_model(opt) if resume_state: logger.info('Resuming training from epoch: {}, iter: {}.'.format( resume_state['epoch'], resume_state['iter'])) start_epoch = resume_state['epoch'] current_step = resume_state['iter'] model.resume_training(resume_state) # handle optimizers and schedulers else: current_step = 0 start_epoch = 0 logger.info('Start training from epoch: {:d}, iter: {:d}'.format(start_epoch, current_step)) for epoch in range(start_epoch, total_epochs + 1): if opt['dist']: train_sampler.set_epoch(epoch) for _, train_data in enumerate(train_loader): current_step += 1 if current_step > total_iters: break model.feed_data(train_data) model.optimize_parameters(current_step) if current_step % opt['logger']['print_freq'] == 0: logs = model.get_current_log() message = '<epoch:{:3d}, iter:{:8,d}, lr:{:.3e}> '.format( epoch, current_step, model.get_current_learning_rate()) for k, v in logs.items(): message += '{:s}: {:.4e} '.format(k, v) # tensorboard logger if opt['use_tb_logger'] and 'debug' not in opt['name']: if rank <= 0: tb_logger.add_scalar(k, v, current_step) if rank <= 0: logger.info(message) model.update_learning_rate(current_step, warmup_iter=opt['train']['warmup_iter']) loss_and_index = epoch-start_epoch if current_step % opt['train']['val_freq'] == 0 and rank <= 0 and val_loader is not None: avg_psnr = val_pix_err_f = val_pix_err_nf = val_mean_color_err = avg_ssim = 0.0 idx = 0 for val_data in val_loader: if idx > 2: break idx += 1 img_name = os.path.splitext(os.path.basename(val_data['LQ_path'][0]))[0] img_dir = os.path.join(opt['path']['val_images'], img_name) util.mkdir(img_dir) model.feed_data(val_data) model.test() visuals = model.get_current_visuals() sr_img = util.tensor2img(visuals['SR']) # uint8 gt_img = util.tensor2img(visuals['GT']) # uint8 print("!!!Now start to save Image!!!!") save_img_path = os.path.join(img_dir, '{:s}_{:d}.png'.format(img_name, current_step)) util.save_img(sr_img, save_img_path) crop_size = opt['scale'] gt_img = gt_img / 255. sr_img = sr_img / 255. cropped_sr_img = sr_img[crop_size:-crop_size, crop_size:-crop_size, :] cropped_gt_img = gt_img[crop_size:-crop_size, crop_size:-crop_size, :] avg_psnr += util.calculate_psnr(cropped_sr_img * 255, cropped_gt_img * 255) avg_ssim += util.calculate_ssim(cropped_sr_img * 255, cropped_gt_img * 255) avg_psnr = avg_psnr / idx avg_ssim = avg_ssim / idx val_pix_err_f /= idx val_pix_err_nf /= idx val_mean_color_err /= idx logger.info('# Validation # PSNR: {:.4e}'.format(avg_psnr)) logger.info('# Validation # SSIM: {:.4e}'.format(avg_ssim)) logger_val = logging.getLogger('val') # validation logger logger_val.info('<epoch:{:3d}, iter:{:8,d}> psnr: {:.4e}'.format( epoch, current_step, avg_psnr)) if opt['use_tb_logger'] and 'debug' not in opt['name']: tb_logger.add_scalar('psnr', avg_psnr, current_step) tb_logger.add_scalar('val_pix_err_f', val_pix_err_f, current_step) tb_logger.add_scalar('val_pix_err_nf', val_pix_err_nf, current_step) tb_logger.add_scalar('val_mean_color_err', val_mean_color_err, current_step) if current_step % opt['logger']['save_checkpoint_freq'] == 0: if rank <= 0: logger.info('Saving models and training states.') model.save(current_step) model.save_training_state(epoch, current_step) if rank <= 0: logger.info('Saving the final model.') model.save('latest') logger.info('End of training.')
def valid(opt, val_loader, current_step, epoch, model, logger, vis, vis_plots): """Validation -> get PSNR and SSIM or D_labmda D_s QNR""" range_max = val_loader.dataset.dynamic_range crop_size = opt['scale'] # calculate psnr and ssim if opt['dataset_type'] == 'reduced': avg_psnr = 0.0 avg_ssim = 0.0 elif opt['dataset_type'] == 'full': avg_D_lambda = 0.0 avg_D_s = 0.0 avg_qnr = 0.0 idx = 0 for val_data in val_loader: idx += 1 # forward model.feed_data(val_data) model.test() visuals = model.get_current_visuals() # Generated image sr_img = util.tensor2img(visuals['SR'], dynamic_range=range_max) # uint cropped_sr_img = sr_img[crop_size:-crop_size, crop_size:-crop_size, :] # "Reference" image if opt['dataset_type'] == 'reduced': gt_img = util.tensor2img(visuals['HRx'], dynamic_range=range_max) # uint cropped_gt_img = gt_img[crop_size:-crop_size, crop_size: -crop_size, :] avg_psnr += util.calculate_psnr(cropped_sr_img, cropped_gt_img, range_max) avg_ssim += util.calculate_ssim(cropped_sr_img, cropped_gt_img, range_max) elif opt['dataset_type'] == 'full': # lr_x lr_x_img = util.tensor2img(visuals['LRx'], dynamic_range=range_max) # lr_p lr_p_img = util.tensor2img(visuals['LRp'], dynamic_range=range_max) now_qnr, now_D_lambda, now_D_s = util.qnr(sr_img, lr_x_img, lr_p_img, satellite='QuickBird', scale=4, block_size=32, p=1, q=1, alpha=1, beta=1) avg_D_lambda += now_D_lambda avg_D_s += now_D_s avg_qnr += now_qnr # if idx == 10: # break if opt['dataset_type'] == 'reduced': avg_psnr = avg_psnr / idx avg_ssim = avg_ssim / idx # log logger.info('# Validation # PSNR: {:.4e} # SSIM: {:.4e}'.format( avg_psnr, avg_ssim)) # logger for validation only logger_val = logging.getLogger('val') # validation logger logger_val.info( '<epoch:{:3d}, iter:{:8,d}> psnr: {:.4e} ssim: {:.4e}'.format( epoch, current_step, avg_psnr, avg_ssim)) # visdom if opt['use_vis_logger'] and 'debug' not in opt['name']: util.update_vis(vis, vis_plots['val_psnr'], current_step, avg_psnr) util.update_vis(vis, vis_plots['val_ssim'], current_step, avg_ssim) return avg_psnr, avg_ssim if opt['dataset_type'] == 'full': avg_D_lambda = avg_D_lambda / idx avg_D_s = avg_D_s / idx avg_qnr = avg_qnr / idx # log logger.info( '# Validation # D_lambda: {:.4e} # D_s: {:.4e} # QNR: {:.4e}'. format(avg_D_lambda, avg_D_s, avg_qnr)) # logger for validation only logger_val = logging.getLogger('val') # validation logger logger_val.info( '<epoch:{:3d}, iter:{:8,d}> D_lambda: {:.4e} D_s: {:.4e} QNR: {:.4e}' .format(epoch, current_step, avg_D_lambda, avg_D_s, avg_qnr)) # visdom if opt['use_vis_logger'] and 'debug' not in opt['name']: util.update_vis(vis, vis_plots['val_no_ref'], current_step, *(avg_D_lambda, avg_D_s, avg_qnr)) return avg_D_lambda, avg_D_s, avg_qnr
if crop_border == 0: cropped_sr_img = sr_img cropped_gt_img = gt_img else: cropped_sr_img = sr_img[crop_border:-crop_border, crop_border:-crop_border, :] cropped_gt_img = gt_img[crop_border:-crop_border, crop_border:-crop_border, :] psnr = util.calculate_psnr(cropped_sr_img * 255, cropped_gt_img * 255) psnr_ori_rgb = util.calculate_psnr(lr_img * 255, cropped_gt_img * 255) dpsnr_rgb = psnr - psnr_ori_rgb ssim = util.calculate_ssim(cropped_sr_img * 255, cropped_gt_img * 255) ssim_ori_rgb = util.calculate_ssim(lr_img * 255, cropped_gt_img * 255) dssim_rgb = ssim - ssim_ori_rgb test_results['psnr'].append(psnr) test_results['dpsnr_rgb'].append(dpsnr_rgb) test_results['ssim'].append(ssim) test_results['dssim_rgb'].append(dssim_rgb) if gt_img.ndim > 2: if gt_img.shape[2] == 3: # RGB image sr_img_y = bgr2ycbcr(sr_img, only_y=True) gt_img_y = bgr2ycbcr(gt_img, only_y=True) lq_img_y = bgr2ycbcr(lr_img, only_y=True) if opt['enhance_uv']:
def main(): # options parser = argparse.ArgumentParser() parser.add_argument('-opt', type=str, required=True, help='Path to options file.') opt = option.parse(parser.parse_args().opt, is_train=False) util.mkdirs((path for key, path in opt['path'].items() if not key == 'pretrain_model_G')) opt = option.dict_to_nonedict(opt) chop2 = opt['chop'] chop_patch_size = opt['chop_patch_size'] multi_upscale = opt['multi_upscale'] scale = opt['scale'] util.setup_logger(None, opt['path']['log'], 'test.log', level=logging.INFO, screen=True) logger = logging.getLogger('base') logger.info(option.dict2str(opt)) # Create test dataset and dataloader test_loaders = [] # znorm = False for phase, dataset_opt in sorted(opt['datasets'].items()): test_set = create_dataset(dataset_opt) test_loader = create_dataloader(test_set, dataset_opt) logger.info('Number of test images in [{:s}]: {:d}'.format(dataset_opt['name'], len(test_set))) test_loaders.append(test_loader) #Temporary, will turn znorm on for all the datasets. Will need to introduce a variable for each dataset and differentiate each one later in the loop. # if dataset_opt['znorm'] and znorm == False: # znorm = True # Create model model = create_model(opt) for test_loader in test_loaders: test_set_name = test_loader.dataset.opt['name'] logger.info('\nTesting [{:s}]...'.format(test_set_name)) test_start_time = time.time() dataset_dir = os.path.join(opt['path']['results_root'], test_set_name) util.mkdir(dataset_dir) test_results = OrderedDict() test_results['psnr'] = [] test_results['ssim'] = [] test_results['psnr_y'] = [] test_results['ssim_y'] = [] for data in test_loader: need_HR = False if test_loader.dataset.opt['dataroot_HR'] is None else True img_path = data['LR_path'][0] # because there's only 1 image per "data" dataset loader? img_name = os.path.splitext(os.path.basename(img_path))[0] znorm = test_loader.dataset.opt['znorm'] if chop2==True: lowres_img = data['LR'] #.to('cuda') if multi_upscale: # Upscale 8 times in different rotations/flips and average the results in a single image LR_90 = lowres_img.transpose(2, 3).flip(2) #PyTorch > 0.4.1 LR_180 = LR_90.transpose(2, 3).flip(2) #PyTorch > 0.4.1 LR_270 = LR_180.transpose(2, 3).flip(2) #PyTorch > 0.4.1 LR_f = lowres_img.flip(3) # horizontal mirror (flip), dim=3 (B,C,H,W=0,1,2,3) #PyTorch > 0.4.1 LR_90f = LR_90.flip(3) # horizontal mirror (flip), dim=3 (B,C,H,W=0,1,2,3) #PyTorch > 0.4.1 LR_180f = LR_180.flip(3) # horizontal mirror (flip), dim=3 (B,C,H,W=0,1,2,3) #PyTorch > 0.4.1 LR_270f = LR_270.flip(3) # horizontal mirror (flip), dim=3 (B,C,H,W=0,1,2,3) #PyTorch > 0.4.1 pred = chop_forward2(lowres_img, model, scale=scale, patch_size=chop_patch_size) pred_90 = chop_forward2(LR_90, model, scale=scale, patch_size=chop_patch_size) pred_180 = chop_forward2(LR_180, model, scale=scale, patch_size=chop_patch_size) pred_270 = chop_forward2(LR_270, model, scale=scale, patch_size=chop_patch_size) pred_f = chop_forward2(LR_f, model, scale=scale, patch_size=chop_patch_size) pred_90f = chop_forward2(LR_90f, model, scale=scale, patch_size=chop_patch_size) pred_180f = chop_forward2(LR_180f, model, scale=scale, patch_size=chop_patch_size) pred_270f = chop_forward2(LR_270f, model, scale=scale, patch_size=chop_patch_size) #convert to numpy array if znorm: #opt['datasets']['train']['znorm']: # If the image range is [-1,1] # In testing, each "dataset" can have a different name (not train, val or other) pred = util.tensor2img(pred,min_max=(-1, 1)).clip(0, 255) # uint8 pred_90 = util.tensor2img(pred_90,min_max=(-1, 1)).clip(0, 255) # uint8 pred_180 = util.tensor2img(pred_180,min_max=(-1, 1)).clip(0, 255) # uint8 pred_270 = util.tensor2img(pred_270,min_max=(-1, 1)).clip(0, 255) # uint8 pred_f = util.tensor2img(pred_f,min_max=(-1, 1)).clip(0, 255) # uint8 pred_90f = util.tensor2img(pred_90f,min_max=(-1, 1)).clip(0, 255) # uint8 pred_180f = util.tensor2img(pred_180f,min_max=(-1, 1)).clip(0, 255) # uint8 pred_270f = util.tensor2img(pred_270f,min_max=(-1, 1)).clip(0, 255) # uint8 else: # Default: Image range is [0,1] pred = util.tensor2img(pred).clip(0, 255) # uint8 pred_90 = util.tensor2img(pred_90).clip(0, 255) # uint8 pred_180 = util.tensor2img(pred_180).clip(0, 255) # uint8 pred_270 = util.tensor2img(pred_270).clip(0, 255) # uint8 pred_f = util.tensor2img(pred_f).clip(0, 255) # uint8 pred_90f = util.tensor2img(pred_90f).clip(0, 255) # uint8 pred_180f = util.tensor2img(pred_180f).clip(0, 255) # uint8 pred_270f = util.tensor2img(pred_270f).clip(0, 255) # uint8 pred_90 = np.rot90(pred_90, 3) pred_180 = np.rot90(pred_180, 2) pred_270 = np.rot90(pred_270, 1) pred_f = np.fliplr(pred_f) pred_90f = np.rot90(np.fliplr(pred_90f), 3) pred_180f = np.rot90(np.fliplr(pred_180f), 2) pred_270f = np.rot90(np.fliplr(pred_270f), 1) #The reason for overflow is that your NumPy arrays (im1arr im2arr) are of the uint8 type (i.e. 8-bit). This means each element of the array can only hold values up to 255, so when your sum exceeds 255, it loops back around 0: #To avoid overflow, your arrays should be able to contain values beyond 255. You need to convert them to floats for instance, perform the blending operation and convert the result back to uint8: # sr_img = (pred + pred_90 + pred_180 + pred_270 + pred_f + pred_90f + pred_180f + pred_270f) / 8.0 sr_img = (pred.astype('float') + pred_90.astype('float') + pred_180.astype('float') + pred_270.astype('float') + pred_f.astype('float') + pred_90f.astype('float') + pred_180f.astype('float') + pred_270f.astype('float')) / 8.0 sr_img = sr_img.astype('uint8') else: highres_output = chop_forward2(lowres_img, model, scale=scale, patch_size=chop_patch_size) #convert to numpy array #highres_image = highres_output[0].permute(1, 2, 0).clamp(0.0, 1.0).cpu() if znorm: #opt['datasets']['train']['znorm']: # If the image range is [-1,1] # In testing, each "dataset" can have a different name (not train, val or other) sr_img = util.tensor2img(highres_output,min_max=(-1, 1)) # uint8 else: # Default: Image range is [0,1] sr_img = util.tensor2img(highres_output) # uint8 else: # will upscale each image in the batch without chopping model.feed_data(data, need_HR=need_HR) model.test() # test visuals = model.get_current_visuals(need_HR=need_HR) if znorm: #opt['datasets']['train']['znorm']: # If the image range is [-1,1] # In testing, each "dataset" can have a different name (not train, val or other) sr_img = util.tensor2img(visuals['SR'],min_max=(-1, 1)) # uint8 else: # Default: Image range is [0,1] sr_img = util.tensor2img(visuals['SR']) # uint8 # save images suffix = opt['suffix'] if suffix: save_img_path = os.path.join(dataset_dir, img_name + suffix + '.png') else: save_img_path = os.path.join(dataset_dir, img_name + '.png') util.save_img(sr_img, save_img_path) # calculate PSNR and SSIM if need_HR: if znorm: #opt['datasets']['train']['znorm']: # If the image range is [-1,1] # In testing, each "dataset" can have a different name (not train, val or other) gt_img = util.tensor2img(visuals['HR'],min_max=(-1, 1)) # uint8 else: # Default: Image range is [0,1] gt_img = util.tensor2img(visuals['HR']) # uint8 gt_img = gt_img / 255. sr_img = sr_img / 255. crop_border = test_loader.dataset.opt['scale'] cropped_sr_img = sr_img[crop_border:-crop_border, crop_border:-crop_border, :] cropped_gt_img = gt_img[crop_border:-crop_border, crop_border:-crop_border, :] psnr = util.calculate_psnr(cropped_sr_img * 255, cropped_gt_img * 255) ssim = util.calculate_ssim(cropped_sr_img * 255, cropped_gt_img * 255) test_results['psnr'].append(psnr) test_results['ssim'].append(ssim) if gt_img.shape[2] == 3: # RGB image sr_img_y = bgr2ycbcr(sr_img, only_y=True) gt_img_y = bgr2ycbcr(gt_img, only_y=True) cropped_sr_img_y = sr_img_y[crop_border:-crop_border, crop_border:-crop_border] cropped_gt_img_y = gt_img_y[crop_border:-crop_border, crop_border:-crop_border] psnr_y = util.calculate_psnr(cropped_sr_img_y * 255, cropped_gt_img_y * 255) ssim_y = util.calculate_ssim(cropped_sr_img_y * 255, cropped_gt_img_y * 255) test_results['psnr_y'].append(psnr_y) test_results['ssim_y'].append(ssim_y) logger.info('{:20s} - PSNR: {:.6f} dB; SSIM: {:.6f}; PSNR_Y: {:.6f} dB; SSIM_Y: {:.6f}.'\ .format(img_name, psnr, ssim, psnr_y, ssim_y)) else: logger.info('{:20s} - PSNR: {:.6f} dB; SSIM: {:.6f}.'.format(img_name, psnr, ssim)) else: logger.info(img_name) if need_HR: # metrics # Average PSNR/SSIM results ave_psnr = sum(test_results['psnr']) / len(test_results['psnr']) ave_ssim = sum(test_results['ssim']) / len(test_results['ssim']) logger.info('----Average PSNR/SSIM results for {}----\n\tPSNR: {:.6f} dB; SSIM: {:.6f}\n'\ .format(test_set_name, ave_psnr, ave_ssim)) if test_results['psnr_y'] and test_results['ssim_y']: ave_psnr_y = sum(test_results['psnr_y']) / len(test_results['psnr_y']) ave_ssim_y = sum(test_results['ssim_y']) / len(test_results['ssim_y']) logger.info('----Y channel, average PSNR/SSIM----\n\tPSNR_Y: {:.6f} dB; SSIM_Y: {:.6f}\n'\ .format(ave_psnr_y, ave_ssim_y))
def main(): #### options parser = argparse.ArgumentParser() parser.add_argument('-opt', type=str, help='Path to option YAML file.') parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none', help='job launcher') parser.add_argument('--local_rank', type=int, default=0) parser.add_argument('--exp_name', type=str, default='temp') parser.add_argument('--degradation_type', type=str, default=None) parser.add_argument('--sigma_x', type=float, default=None) parser.add_argument('--sigma_y', type=float, default=None) parser.add_argument('--theta', type=float, default=None) args = parser.parse_args() if args.exp_name == 'temp': opt = option.parse(args.opt, is_train=False) else: opt = option.parse(args.opt, is_train=False, exp_name=args.exp_name) # convert to NoneDict, which returns None for missing keys opt = option.dict_to_nonedict(opt) inner_loop_name = opt['train']['maml']['optimizer'][0] + str(opt['train']['maml']['adapt_iter']) + str(math.floor(math.log10(opt['train']['maml']['lr_alpha']))) meta_loop_name = opt['train']['optim'][0] + str(math.floor(math.log10(opt['train']['lr_G']))) if args.degradation_type is not None: if args.degradation_type == 'preset': opt['datasets']['val']['degradation_mode'] = args.degradation_type else: opt['datasets']['val']['degradation_type'] = args.degradation_type if args.sigma_x is not None: opt['datasets']['val']['sigma_x'] = args.sigma_x if args.sigma_y is not None: opt['datasets']['val']['sigma_y'] = args.sigma_y if args.theta is not None: opt['datasets']['val']['theta'] = args.theta if 'degradation_mode' not in opt['datasets']['val'].keys(): degradation_name = '' elif opt['datasets']['val']['degradation_mode'] == 'set': degradation_name = '_' + str(opt['datasets']['val']['degradation_type'])\ + '_' + str(opt['datasets']['val']['sigma_x']) \ + '_' + str(opt['datasets']['val']['sigma_y'])\ + '_' + str(opt['datasets']['val']['theta']) else: degradation_name = '_' + opt['datasets']['val']['degradation_mode'] folder_name = opt['name'] + '_' + degradation_name if args.exp_name != 'temp': folder_name = args.exp_name torch.backends.cudnn.benchmark = False torch.backends.cudnn.deterministic = True #### create train and val dataloader dataset_ratio = 200 # enlarge the size of each epoch for phase, dataset_opt in opt['datasets'].items(): if phase == 'train': pass elif phase == 'val': if '+' in opt['datasets']['val']['name']: raise NotImplementedError('Do not use + signs in test mode') else: val_set = create_dataset(dataset_opt, scale=opt['scale'], kernel_size=opt['datasets']['train']['kernel_size'], model_name=opt['network_E']['which_model_E']) # val_set = loader.get_dataset(opt, train=False) val_loader = create_dataloader(val_set, dataset_opt, opt, None) print('Number of val images in [{:s}]: {:d}'.format(dataset_opt['name'], len(val_set))) else: raise NotImplementedError('Phase [{:s}] is not recognized.'.format(phase)) #### create model models = create_model(opt) assert len(models) == 2 model, est_model = models[0], models[1] modelcp, est_modelcp = create_model(opt) _, est_model_fixed = create_model(opt) center_idx = (opt['datasets']['val']['N_frames']) // 2 lr_alpha = opt['train']['maml']['lr_alpha'] update_step = opt['train']['maml']['adapt_iter'] with_GT = False if opt['datasets']['val']['mode'] == 'demo' else True pd_log = pd.DataFrame(columns=['PSNR_Bicubic', 'PSNR_Ours', 'SSIM_Bicubic', 'SSIM_Ours']) def crop(LR_seq, HR, num_patches_for_batch=4, patch_size=44): """ Crop given patches. Args: LR_seq: (B=1) x T x C x H x W HR: (B=1) x C x H x W patch_size (int, optional): Return: B(=batch_size) x T x C x H x W """ # Find the lowest resolution cropped_lr = [] cropped_hr = [] assert HR.size(0) == 1 LR_seq_ = LR_seq[0] HR_ = HR[0] for _ in range(num_patches_for_batch): patch_lr, patch_hr = preprocessing.common_crop(LR_seq_, HR_, patch_size=patch_size // 2) cropped_lr.append(patch_lr) cropped_hr.append(patch_hr) cropped_lr = torch.stack(cropped_lr, dim=0) cropped_hr = torch.stack(cropped_hr, dim=0) return cropped_lr, cropped_hr # Single GPU # PSNR_rlt: psnr_init, psnr_before, psnr_after psnr_rlt = [{}, {}] # SSIM_rlt: ssim_init, ssim_after ssim_rlt = [{}, {}] pbar = util.ProgressBar(len(val_set)) for val_data in val_loader: folder = val_data['folder'][0] idx_d = int(val_data['idx'][0].split('/')[0]) if 'name' in val_data.keys(): name = val_data['name'][0][center_idx][0] else: name = folder train_folder = os.path.join('../test_results', folder_name, name) maml_train_folder = os.path.join(train_folder, 'DynaVSR') if not os.path.exists(train_folder): os.makedirs(train_folder, exist_ok=False) if not os.path.exists(maml_train_folder): os.mkdir(maml_train_folder) for i in range(len(psnr_rlt)): if psnr_rlt[i].get(folder, None) is None: psnr_rlt[i][folder] = [] for i in range(len(ssim_rlt)): if ssim_rlt[i].get(folder, None) is None: ssim_rlt[i][folder] = [] cropped_meta_train_data = {} meta_train_data = {} meta_test_data = {} # Make SuperLR seq using estimation model meta_train_data['GT'] = val_data['LQs'][:, center_idx] meta_test_data['LQs'] = val_data['LQs'][0:1] meta_test_data['GT'] = val_data['GT'][0:1, center_idx] if with_GT else None # Check whether the batch size of each validation data is 1 assert val_data['LQs'].size(0) == 1 if opt['network_G']['which_model_G'] == 'TOF': LQs = meta_test_data['LQs'] B, T, C, H, W = LQs.shape LQs = LQs.reshape(B*T, C, H, W) Bic_LQs = F.interpolate(LQs, scale_factor=opt['scale'], mode='bicubic', align_corners=True) meta_test_data['LQs'] = Bic_LQs.reshape(B, T, C, H*opt['scale'], W*opt['scale']) ## Before start testing # Bicubic Model Results modelcp.load_network(opt['path']['bicubic_G'], modelcp.netG) modelcp.feed_data(meta_test_data, need_GT=with_GT) modelcp.test() if with_GT: model_start_visuals = modelcp.get_current_visuals(need_GT=True) hr_image = util.tensor2img(model_start_visuals['GT'], mode='rgb') start_image = util.tensor2img(model_start_visuals['rlt'], mode='rgb') psnr_rlt[0][folder].append(util.calculate_psnr(start_image, hr_image)) ssim_rlt[0][folder].append(util.calculate_ssim(start_image, hr_image)) modelcp.netG, est_modelcp.netE = deepcopy(model.netG), deepcopy(est_model.netE) ########## SLR LOSS Preparation ############ est_model_fixed.load_network(opt['path']['fixed_E'], est_model_fixed.netE) optim_params = [] for k, v in modelcp.netG.named_parameters(): if v.requires_grad: optim_params.append(v) if not opt['train']['use_real']: for k, v in est_modelcp.netE.named_parameters(): if v.requires_grad: optim_params.append(v) if opt['train']['maml']['optimizer'] == 'Adam': inner_optimizer = torch.optim.Adam(optim_params, lr=lr_alpha, betas=( opt['train']['maml']['beta1'], opt['train']['maml']['beta2'])) elif opt['train']['maml']['optimizer'] == 'SGD': inner_optimizer = torch.optim.SGD(optim_params, lr=lr_alpha) else: raise NotImplementedError() # Inner Loop Update st = time.time() for i in range(update_step): # Make SuperLR seq using UPDATED estimation model if not opt['train']['use_real']: est_modelcp.feed_data(val_data) est_modelcp.forward_without_optim() superlr_seq = est_modelcp.fake_L meta_train_data['LQs'] = superlr_seq else: meta_train_data['LQs'] = val_data['SuperLQs'] if opt['network_G']['which_model_G'] == 'TOF': # Bicubic upsample to match the size LQs = meta_train_data['LQs'] B, T, C, H, W = LQs.shape LQs = LQs.reshape(B*T, C, H, W) Bic_LQs = F.interpolate(LQs, scale_factor=opt['scale'], mode='bicubic', align_corners=True) meta_train_data['LQs'] = Bic_LQs.reshape(B, T, C, H*opt['scale'], W*opt['scale']) # Update both modelcp + estmodelcp jointly inner_optimizer.zero_grad() if opt['train']['maml']['use_patch']: cropped_meta_train_data['LQs'], cropped_meta_train_data['GT'] = \ crop(meta_train_data['LQs'], meta_train_data['GT'], opt['train']['maml']['num_patch'], opt['train']['maml']['patch_size']) modelcp.feed_data(cropped_meta_train_data) else: modelcp.feed_data(meta_train_data) loss_train = modelcp.calculate_loss() ##################### SLR LOSS ################### est_model_fixed.feed_data(val_data) est_model_fixed.test() slr_initialized = est_model_fixed.fake_L slr_initialized = slr_initialized.to('cuda') if opt['network_G']['which_model_G'] == 'TOF': loss_train += 10 * F.l1_loss(LQs.to('cuda').squeeze(0), slr_initialized) else: loss_train += 10 * F.l1_loss(meta_train_data['LQs'].to('cuda'), slr_initialized) loss_train.backward() inner_optimizer.step() et = time.time() update_time = et - st modelcp.feed_data(meta_test_data, need_GT=with_GT) modelcp.test() model_update_visuals = modelcp.get_current_visuals(need_GT=False) update_image = util.tensor2img(model_update_visuals['rlt'], mode='rgb') # Save and calculate final image imageio.imwrite(os.path.join(maml_train_folder, '{:08d}.png'.format(idx_d)), update_image) if with_GT: psnr_rlt[1][folder].append(util.calculate_psnr(update_image, hr_image)) ssim_rlt[1][folder].append(util.calculate_ssim(update_image, hr_image)) name_df = '{}/{:08d}'.format(folder, idx_d) if name_df in pd_log.index: pd_log.at[name_df, 'PSNR_Bicubic'] = psnr_rlt[0][folder][-1] pd_log.at[name_df, 'PSNR_Ours'] = psnr_rlt[1][folder][-1] pd_log.at[name_df, 'SSIM_Bicubic'] = ssim_rlt[0][folder][-1] pd_log.at[name_df, 'SSIM_Ours'] = ssim_rlt[1][folder][-1] else: pd_log.loc[name_df] = [psnr_rlt[0][folder][-1], psnr_rlt[1][folder][-1], ssim_rlt[0][folder][-1], ssim_rlt[1][folder][-1]] pd_log.to_csv(os.path.join('../test_results', folder_name, 'psnr_update.csv')) pbar.update('Test {} - {}: I: {:.3f}/{:.4f} \tF+: {:.3f}/{:.4f} \tTime: {:.3f}s' .format(folder, idx_d, psnr_rlt[0][folder][-1], ssim_rlt[0][folder][-1], psnr_rlt[1][folder][-1], ssim_rlt[1][folder][-1], update_time )) else: pbar.update() if with_GT: psnr_rlt_avg = {} psnr_total_avg = 0. # Just calculate the final value of psnr_rlt(i.e. psnr_rlt[2]) for k, v in psnr_rlt[0].items(): psnr_rlt_avg[k] = sum(v) / len(v) psnr_total_avg += psnr_rlt_avg[k] psnr_total_avg /= len(psnr_rlt[0]) log_s = '# Validation # Bic PSNR: {:.4e}:'.format(psnr_total_avg) for k, v in psnr_rlt_avg.items(): log_s += ' {}: {:.4e}'.format(k, v) print(log_s) psnr_rlt_avg = {} psnr_total_avg = 0. # Just calculate the final value of psnr_rlt(i.e. psnr_rlt[2]) for k, v in psnr_rlt[1].items(): psnr_rlt_avg[k] = sum(v) / len(v) psnr_total_avg += psnr_rlt_avg[k] psnr_total_avg /= len(psnr_rlt[1]) log_s = '# Validation # PSNR: {:.4e}:'.format(psnr_total_avg) for k, v in psnr_rlt_avg.items(): log_s += ' {}: {:.4e}'.format(k, v) print(log_s) ssim_rlt_avg = {} ssim_total_avg = 0. # Just calculate the final value of ssim_rlt(i.e. ssim_rlt[1]) for k, v in ssim_rlt[0].items(): ssim_rlt_avg[k] = sum(v) / len(v) ssim_total_avg += ssim_rlt_avg[k] ssim_total_avg /= len(ssim_rlt[0]) log_s = '# Validation # Bicubic SSIM: {:.4e}:'.format(ssim_total_avg) for k, v in ssim_rlt_avg.items(): log_s += ' {}: {:.4e}'.format(k, v) print(log_s) ssim_rlt_avg = {} ssim_total_avg = 0. # Just calculate the final value of ssim_rlt(i.e. ssim_rlt[1]) for k, v in ssim_rlt[1].items(): ssim_rlt_avg[k] = sum(v) / len(v) ssim_total_avg += ssim_rlt_avg[k] ssim_total_avg /= len(ssim_rlt[1]) log_s = '# Validation # SSIM: {:.4e}:'.format(ssim_total_avg) for k, v in ssim_rlt_avg.items(): log_s += ' {}: {:.4e}'.format(k, v) print(log_s) print('End of evaluation.')
def main(jsonPath): # options opt = option.parse(jsonPath, is_train=False) util.mkdirs((path for key, path in opt["path"].items() if not key == "pretrain_model_G")) opt = option.dict_to_nonedict(opt) util.setup_logger(None, opt["path"]["log"], "test.log", level=logging.INFO, screen=True) logger = logging.getLogger("base") logger.info(option.dict2str(opt)) # Create test dataset and dataloader test_loaders = [] for phase, dataset_opt in sorted(opt["datasets"].items()): test_set = create_dataset(dataset_opt) test_loader = create_dataloader(test_set, dataset_opt) logger.info("Number of test images in [{:s}]: {:d}".format( dataset_opt["name"], len(test_set))) test_loaders.append(test_loader) # Create model model = create_model(opt) for test_loader in test_loaders: test_set_name = test_loader.dataset.opt["name"] logger.info("\nTesting [{:s}]...".format(test_set_name)) # test_start_time = time.time() dataset_dir = os.path.join(opt["path"]["results_root"], test_set_name) util.mkdir(dataset_dir) test_results = OrderedDict() test_results["psnr"] = [] test_results["ssim"] = [] test_results["psnr_y"] = [] test_results["ssim_y"] = [] for data in test_loader: need_HR = False if test_loader.dataset.opt[ "dataroot_HR"] is None else True model.feed_data(data, need_HR=need_HR) img_path = data["LR_path"][0] img_name = os.path.splitext(os.path.basename(img_path))[0] model.test() # test visuals = model.get_current_visuals(need_HR=need_HR) sr_img = util.tensor2img(visuals["SR"]) # uint8 # save images suffix = opt["suffix"] if suffix: save_img_path = os.path.join(dataset_dir, img_name + suffix + ".png") else: save_img_path = os.path.join(dataset_dir, img_name + ".png") util.save_img(sr_img, save_img_path) # calculate PSNR and SSIM if need_HR: gt_img = util.tensor2img(visuals["HR"]) gt_img = gt_img / 255.0 sr_img = sr_img / 255.0 crop_border = test_loader.dataset.opt["scale"] cropped_sr_img = sr_img[crop_border:-crop_border, crop_border:-crop_border, :] cropped_gt_img = gt_img[crop_border:-crop_border, crop_border:-crop_border, :] psnr = util.calculate_psnr(cropped_sr_img * 255, cropped_gt_img * 255) ssim = util.calculate_ssim(cropped_sr_img * 255, cropped_gt_img * 255) test_results["psnr"].append(psnr) test_results["ssim"].append(ssim) if gt_img.shape[2] == 3: # RGB image sr_img_y = bgr2ycbcr(sr_img, only_y=True) gt_img_y = bgr2ycbcr(gt_img, only_y=True) cropped_sr_img_y = sr_img_y[crop_border:-crop_border, crop_border:-crop_border] cropped_gt_img_y = gt_img_y[crop_border:-crop_border, crop_border:-crop_border] psnr_y = util.calculate_psnr(cropped_sr_img_y * 255, cropped_gt_img_y * 255) ssim_y = util.calculate_ssim(cropped_sr_img_y * 255, cropped_gt_img_y * 255) test_results["psnr_y"].append(psnr_y) test_results["ssim_y"].append(ssim_y) logger.info( "{:20s} - PSNR: {:.6f} dB; SSIM: {:.6f}; PSNR_Y: {:.6f} dB; SSIM_Y: {:.6f}." .format(img_name, psnr, ssim, psnr_y, ssim_y)) else: logger.info( "{:20s} - PSNR: {:.6f} dB; SSIM: {:.6f}.".format( img_name, psnr, ssim)) else: logger.info(img_name) if need_HR: # metrics # Average PSNR/SSIM results ave_psnr = sum(test_results["psnr"]) / len(test_results["psnr"]) ave_ssim = sum(test_results["ssim"]) / len(test_results["ssim"]) logger.info( "----Average PSNR/SSIM results for {}----\n\tPSNR: {:.6f} dB; SSIM: {:.6f}\n" .format(test_set_name, ave_psnr, ave_ssim)) if test_results["psnr_y"] and test_results["ssim_y"]: ave_psnr_y = sum(test_results["psnr_y"]) / len( test_results["psnr_y"]) ave_ssim_y = sum(test_results["ssim_y"]) / len( test_results["ssim_y"]) logger.info( "----Y channel, average PSNR/SSIM----\n\tPSNR_Y: {:.6f} dB; SSIM_Y: {:.6f}\n" .format(ave_psnr_y, ave_ssim_y))
def main(): #### options parser = argparse.ArgumentParser() parser.add_argument('-opt', type=str, help='Path to option YAML file.') parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none', help='job launcher') parser.add_argument('--local_rank', type=int, default=0) parser.add_argument('--exp_name', type=str, default='temp') parser.add_argument('--degradation_type', type=str, default=None) parser.add_argument('--sigma_x', type=float, default=None) parser.add_argument('--sigma_y', type=float, default=None) parser.add_argument('--theta', type=float, default=None) args = parser.parse_args() if args.exp_name == 'temp': opt = option.parse(args.opt, is_train=True) else: opt = option.parse(args.opt, is_train=True, exp_name=args.exp_name) # convert to NoneDict, which returns None for missing keys opt = option.dict_to_nonedict(opt) inner_loop_name = opt['train']['maml']['optimizer'][0] + str( opt['train']['maml']['adapt_iter']) + str( math.floor(math.log10(opt['train']['maml']['lr_alpha']))) meta_loop_name = opt['train']['optim'][0] + str( math.floor(math.log10(opt['train']['lr_G']))) if args.degradation_type is not None: if args.degradation_type == 'preset': opt['datasets']['val']['degradation_mode'] = args.degradation_type else: opt['datasets']['val']['degradation_type'] = args.degradation_type if args.sigma_x is not None: opt['datasets']['val']['sigma_x'] = args.sigma_x if args.sigma_y is not None: opt['datasets']['val']['sigma_y'] = args.sigma_y if args.theta is not None: opt['datasets']['val']['theta'] = args.theta if opt['datasets']['val']['degradation_mode'] == 'set': degradation_name = str(opt['datasets']['val']['degradation_type'])\ + '_' + str(opt['datasets']['val']['sigma_x']) \ + '_' + str(opt['datasets']['val']['sigma_y'])\ + '_' + str(opt['datasets']['val']['theta']) else: degradation_name = opt['datasets']['val']['degradation_mode'] patch_name = 'p{}x{}'.format( opt['train']['maml']['patch_size'], opt['train']['maml'] ['num_patch']) if opt['train']['maml']['use_patch'] else 'full' use_real_flag = '_ideal' if opt['train']['use_real'] else '' folder_name = opt[ 'name'] + '_' + degradation_name # + '_' + inner_loop_name + meta_loop_name + '_' + degradation_name + '_' + patch_name + use_real_flag if args.exp_name != 'temp': folder_name = args.exp_name #### distributed training settings if args.launcher == 'none': # disabled distributed training opt['dist'] = False rank = -1 print('Disabled distributed training.') else: opt['dist'] = True init_dist() world_size = torch.distributed.get_world_size() rank = torch.distributed.get_rank() #### loading resume state if exists if opt['path'].get('resume_state', None): # distributed resuming: all load into default GPU device_id = torch.cuda.current_device() resume_state = torch.load( opt['path']['resume_state'], map_location=lambda storage, loc: storage.cuda(device_id)) option.check_resume(opt, resume_state['iter']) # check resume options else: resume_state = None #### mkdir and loggers if rank <= 0: # normal training (rank -1) OR distributed training (rank 0) if resume_state is None: #util.mkdir_and_rename( # opt['path']['experiments_root']) # rename experiment folder if exists #util.mkdirs((path for key, path in opt['path'].items() if not key == 'experiments_root' # and 'pretrain_model' not in key and 'resume' not in key)) if not os.path.exists(opt['path']['experiments_root']): os.mkdir(opt['path']['experiments_root']) # raise ValueError('Path does not exists - check path') # config loggers. Before it, the log will not work util.setup_logger('base', opt['path']['log'], 'train_' + opt['name'], level=logging.INFO, screen=True, tofile=True) logger = logging.getLogger('base') #logger.info(option.dict2str(opt)) # tensorboard logger if opt['use_tb_logger'] and 'debug' not in opt['name']: version = float(torch.__version__[0:3]) if version >= 1.1: # PyTorch 1.1 from torch.utils.tensorboard import SummaryWriter else: logger.info( 'You are using PyTorch {}. Tensorboard will use [tensorboardX]' .format(version)) from tensorboardX import SummaryWriter tb_logger = SummaryWriter(log_dir='../tb_logger/' + folder_name) else: util.setup_logger('base', opt['path']['log'], 'train', level=logging.INFO, screen=True) logger = logging.getLogger('base') #### random seed seed = opt['train']['manual_seed'] if seed is None: seed = random.randint(1, 10000) if rank <= 0: logger.info('Random seed: {}'.format(seed)) util.set_random_seed(seed) torch.backends.cudnn.benchmark = False torch.backends.cudnn.deterministic = True #### create train and val dataloader dataset_ratio = 200 # enlarge the size of each epoch for phase, dataset_opt in opt['datasets'].items(): if phase == 'train': pass elif phase == 'val': if '+' in opt['datasets']['val']['name']: val_set, val_loader = [], [] valname_list = opt['datasets']['val']['name'].split('+') for i in range(len(valname_list)): val_set.append( create_dataset( dataset_opt, scale=opt['scale'], kernel_size=opt['datasets']['train'] ['kernel_size'], model_name=opt['network_E']['which_model_E'], idx=i)) val_loader.append( create_dataloader(val_set[-1], dataset_opt, opt, None)) else: val_set = create_dataset( dataset_opt, scale=opt['scale'], kernel_size=opt['datasets']['train']['kernel_size'], model_name=opt['network_E']['which_model_E']) # val_set = loader.get_dataset(opt, train=False) val_loader = create_dataloader(val_set, dataset_opt, opt, None) if rank <= 0: logger.info('Number of val images in [{:s}]: {:d}'.format( dataset_opt['name'], len(val_set))) else: raise NotImplementedError( 'Phase [{:s}] is not recognized.'.format(phase)) #### create model models = create_model(opt) assert len(models) == 2 model, est_model = models[0], models[1] modelcp, est_modelcp = create_model(opt) _, est_model_fixed = create_model(opt) center_idx = (opt['datasets']['val']['N_frames']) // 2 lr_alpha = opt['train']['maml']['lr_alpha'] update_step = opt['train']['maml']['adapt_iter'] pd_log = pd.DataFrame( columns=['PSNR_Bicubic', 'PSNR_Ours', 'SSIM_Bicubic', 'SSIM_Ours']) def crop(LR_seq, HR, num_patches_for_batch=4, patch_size=44): """ Crop given patches. Args: LR_seq: (B=1) x T x C x H x W HR: (B=1) x C x H x W patch_size (int, optional): Return: B(=batch_size) x T x C x H x W """ # Find the lowest resolution cropped_lr = [] cropped_hr = [] assert HR.size(0) == 1 LR_seq_ = LR_seq[0] HR_ = HR[0] for _ in range(num_patches_for_batch): patch_lr, patch_hr = preprocessing.common_crop( LR_seq_, HR_, patch_size=patch_size // 2) cropped_lr.append(patch_lr) cropped_hr.append(patch_hr) cropped_lr = torch.stack(cropped_lr, dim=0) cropped_hr = torch.stack(cropped_hr, dim=0) return cropped_lr, cropped_hr # Single GPU # PSNR_rlt: psnr_init, psnr_before, psnr_after psnr_rlt = [{}, {}] # SSIM_rlt: ssim_init, ssim_after ssim_rlt = [{}, {}] pbar = util.ProgressBar(len(val_set)) for val_data in val_loader: folder = val_data['folder'][0] idx_d = int(val_data['idx'][0].split('/')[0]) if 'name' in val_data.keys(): name = val_data['name'][0][center_idx][0] else: #name = '{}/{:08d}'.format(folder, idx_d) name = folder train_folder = os.path.join('../results_for_paper', folder_name, name) hr_train_folder = os.path.join(train_folder, 'hr') bic_train_folder = os.path.join(train_folder, 'bic') maml_train_folder = os.path.join(train_folder, 'maml') #slr_train_folder = os.path.join(train_folder, 'slr') # print(train_folder) if not os.path.exists(train_folder): os.makedirs(train_folder, exist_ok=False) if not os.path.exists(hr_train_folder): os.mkdir(hr_train_folder) if not os.path.exists(bic_train_folder): os.mkdir(bic_train_folder) if not os.path.exists(maml_train_folder): os.mkdir(maml_train_folder) #if not os.path.exists(slr_train_folder): # os.mkdir(slr_train_folder) for i in range(len(psnr_rlt)): if psnr_rlt[i].get(folder, None) is None: psnr_rlt[i][folder] = [] for i in range(len(ssim_rlt)): if ssim_rlt[i].get(folder, None) is None: ssim_rlt[i][folder] = [] if idx_d % 10 != 5: #continue pass cropped_meta_train_data = {} meta_train_data = {} meta_test_data = {} # Make SuperLR seq using estimation model meta_train_data['GT'] = val_data['LQs'][:, center_idx] meta_test_data['LQs'] = val_data['LQs'][0:1] meta_test_data['GT'] = val_data['GT'][0:1, center_idx] # Check whether the batch size of each validation data is 1 assert val_data['SuperLQs'].size(0) == 1 if opt['network_G']['which_model_G'] == 'TOF': LQs = meta_test_data['LQs'] B, T, C, H, W = LQs.shape LQs = LQs.reshape(B * T, C, H, W) Bic_LQs = F.interpolate(LQs, scale_factor=opt['scale'], mode='bicubic', align_corners=True) meta_test_data['LQs'] = Bic_LQs.reshape(B, T, C, H * opt['scale'], W * opt['scale']) ## Before start training, first save the bicubic, real outputs # Bicubic modelcp.load_network(opt['path']['bicubic_G'], modelcp.netG) modelcp.feed_data(meta_test_data) modelcp.test() model_start_visuals = modelcp.get_current_visuals(need_GT=True) hr_image = util.tensor2img(model_start_visuals['GT'], mode='rgb') start_image = util.tensor2img(model_start_visuals['rlt'], mode='rgb') #####imageio.imwrite(os.path.join(hr_train_folder, '{:08d}.png'.format(idx_d)), hr_image) #####imageio.imwrite(os.path.join(bic_train_folder, '{:08d}.png'.format(idx_d)), start_image) psnr_rlt[0][folder].append(util.calculate_psnr(start_image, hr_image)) ssim_rlt[0][folder].append(util.calculate_ssim(start_image, hr_image)) modelcp.netG, est_modelcp.netE = deepcopy(model.netG), deepcopy( est_model.netE) ########## SLR LOSS Preparation ############ est_model_fixed.load_network(opt['path']['fixed_E'], est_model_fixed.netE) optim_params = [] for k, v in modelcp.netG.named_parameters(): if v.requires_grad: optim_params.append(v) if not opt['train']['use_real']: for k, v in est_modelcp.netE.named_parameters(): if v.requires_grad: optim_params.append(v) if opt['train']['maml']['optimizer'] == 'Adam': inner_optimizer = torch.optim.Adam( optim_params, lr=lr_alpha, betas=(opt['train']['maml']['beta1'], opt['train']['maml']['beta2'])) elif opt['train']['maml']['optimizer'] == 'SGD': inner_optimizer = torch.optim.SGD(optim_params, lr=lr_alpha) else: raise NotImplementedError() # Inner Loop Update st = time.time() for i in range(update_step): # Make SuperLR seq using UPDATED estimation model if not opt['train']['use_real']: est_modelcp.feed_data(val_data) # est_model.test() est_modelcp.forward_without_optim() superlr_seq = est_modelcp.fake_L meta_train_data['LQs'] = superlr_seq else: meta_train_data['LQs'] = val_data['SuperLQs'] if opt['network_G']['which_model_G'] == 'TOF': # Bicubic upsample to match the size LQs = meta_train_data['LQs'] B, T, C, H, W = LQs.shape LQs = LQs.reshape(B * T, C, H, W) Bic_LQs = F.interpolate(LQs, scale_factor=opt['scale'], mode='bicubic', align_corners=True) meta_train_data['LQs'] = Bic_LQs.reshape( B, T, C, H * opt['scale'], W * opt['scale']) # Update both modelcp + estmodelcp jointly inner_optimizer.zero_grad() if opt['train']['maml']['use_patch']: cropped_meta_train_data['LQs'], cropped_meta_train_data['GT'] = \ crop(meta_train_data['LQs'], meta_train_data['GT'], opt['train']['maml']['num_patch'], opt['train']['maml']['patch_size']) modelcp.feed_data(cropped_meta_train_data) else: modelcp.feed_data(meta_train_data) loss_train = modelcp.calculate_loss() ##################### SLR LOSS ################### est_model_fixed.feed_data(val_data) est_model_fixed.test() slr_initialized = est_model_fixed.fake_L slr_initialized = slr_initialized.to('cuda') if opt['network_G']['which_model_G'] == 'TOF': loss_train += 10 * F.l1_loss( LQs.to('cuda').squeeze(0), slr_initialized) else: loss_train += 10 * F.l1_loss(meta_train_data['LQs'].to('cuda'), slr_initialized) loss_train.backward() inner_optimizer.step() et = time.time() update_time = et - st modelcp.feed_data(meta_test_data) modelcp.test() model_update_visuals = modelcp.get_current_visuals(need_GT=False) update_image = util.tensor2img(model_update_visuals['rlt'], mode='rgb') # Save and calculate final image imageio.imwrite( os.path.join(maml_train_folder, '{:08d}.png'.format(idx_d)), update_image) psnr_rlt[1][folder].append(util.calculate_psnr(update_image, hr_image)) ssim_rlt[1][folder].append(util.calculate_ssim(update_image, hr_image)) name_df = '{}/{:08d}'.format(folder, idx_d) if name_df in pd_log.index: pd_log.at[name_df, 'PSNR_Bicubic'] = psnr_rlt[0][folder][-1] pd_log.at[name_df, 'PSNR_Ours'] = psnr_rlt[1][folder][-1] pd_log.at[name_df, 'SSIM_Bicubic'] = ssim_rlt[0][folder][-1] pd_log.at[name_df, 'SSIM_Ours'] = ssim_rlt[1][folder][-1] else: pd_log.loc[name_df] = [ psnr_rlt[0][folder][-1], psnr_rlt[1][folder][-1], ssim_rlt[0][folder][-1], ssim_rlt[1][folder][-1] ] pd_log.to_csv( os.path.join('../results_for_paper', folder_name, 'psnr_update.csv')) pbar.update( 'Test {} - {}: I: {:.3f}/{:.4f} \tF+: {:.3f}/{:.4f} \tTime: {:.3f}s' .format(folder, idx_d, psnr_rlt[0][folder][-1], ssim_rlt[0][folder][-1], psnr_rlt[1][folder][-1], ssim_rlt[1][folder][-1], update_time)) psnr_rlt_avg = {} psnr_total_avg = 0. # Just calculate the final value of psnr_rlt(i.e. psnr_rlt[2]) for k, v in psnr_rlt[0].items(): psnr_rlt_avg[k] = sum(v) / len(v) psnr_total_avg += psnr_rlt_avg[k] psnr_total_avg /= len(psnr_rlt[0]) log_s = '# Validation # Bic PSNR: {:.4e}:'.format(psnr_total_avg) for k, v in psnr_rlt_avg.items(): log_s += ' {}: {:.4e}'.format(k, v) logger.info(log_s) psnr_rlt_avg = {} psnr_total_avg = 0. # Just calculate the final value of psnr_rlt(i.e. psnr_rlt[2]) for k, v in psnr_rlt[1].items(): psnr_rlt_avg[k] = sum(v) / len(v) psnr_total_avg += psnr_rlt_avg[k] psnr_total_avg /= len(psnr_rlt[1]) log_s = '# Validation # PSNR: {:.4e}:'.format(psnr_total_avg) for k, v in psnr_rlt_avg.items(): log_s += ' {}: {:.4e}'.format(k, v) logger.info(log_s) ssim_rlt_avg = {} ssim_total_avg = 0. # Just calculate the final value of ssim_rlt(i.e. ssim_rlt[1]) for k, v in ssim_rlt[0].items(): ssim_rlt_avg[k] = sum(v) / len(v) ssim_total_avg += ssim_rlt_avg[k] ssim_total_avg /= len(ssim_rlt[0]) log_s = '# Validation # Bicubic SSIM: {:.4e}:'.format(ssim_total_avg) for k, v in ssim_rlt_avg.items(): log_s += ' {}: {:.4e}'.format(k, v) logger.info(log_s) ssim_rlt_avg = {} ssim_total_avg = 0. # Just calculate the final value of ssim_rlt(i.e. ssim_rlt[1]) for k, v in ssim_rlt[1].items(): ssim_rlt_avg[k] = sum(v) / len(v) ssim_total_avg += ssim_rlt_avg[k] ssim_total_avg /= len(ssim_rlt[1]) log_s = '# Validation # SSIM: {:.4e}:'.format(ssim_total_avg) for k, v in ssim_rlt_avg.items(): log_s += ' {}: {:.4e}'.format(k, v) logger.info(log_s) logger.info('End of evaluation.')
def main(): #### options parser = argparse.ArgumentParser() parser.add_argument('-opt', type=str, required=True, help='Path to options YMAL file.') opt = option.parse(parser.parse_args().opt, is_train=False) opt = option.dict_to_nonedict(opt) util.mkdirs((path for key, path in opt['path'].items() if not key == 'experiments_root' and 'pretrain_model' not in key and 'resume' not in key)) util.setup_logger('base', opt['path']['log'], 'test_' + opt['name'], level=logging.INFO, screen=True, tofile=True) logger = logging.getLogger('base') logger.info(option.dict2str(opt)) #### Create test dataset and dataloader test_loaders = [] for phase, dataset_opt in sorted(opt['datasets'].items()): test_set = create_dataset(dataset_opt) test_loader = create_dataloader(test_set, dataset_opt) logger.info('Number of test images in [{:s}]: {:d}'.format( dataset_opt['name'], len(test_set))) test_loaders.append(test_loader) model = create_model(opt) for test_loader in test_loaders: test_set_name = test_loader.dataset.opt['name'] logger.info('\nTesting [{:s}]...'.format(test_set_name)) test_start_time = time.time() dataset_dir = osp.join(opt['path']['results_root'], test_set_name) util.mkdir(dataset_dir) test_results = OrderedDict() test_results['psnr'] = [] test_results['ssim'] = [] test_results['psnr_y'] = [] test_results['ssim_y'] = [] for data in test_loader: need_GT = False if test_loader.dataset.opt[ 'dataroot_GT'] is None else True model.feed_data(data, need_GT=need_GT) img_path = data['GT_path'][0] if need_GT else data['LQ_path'][0] img_name = osp.splitext(osp.basename(img_path))[0] model.test() visuals = model.get_current_visuals(need_GT=need_GT) sr_img = util.tensor2img(visuals['rlt']) # uint8 # save images suffix = opt['suffix'] if suffix: save_img_path = osp.join(dataset_dir, img_name + suffix + '.png') else: save_img_path = osp.join(dataset_dir, img_name + '.png') util.save_img(sr_img, save_img_path) # calculate PSNR and SSIM if need_GT: gt_img = util.tensor2img(visuals['GT']) sr_img, gt_img = util.crop_border([sr_img, gt_img], opt['scale']) psnr = util.calculate_psnr(sr_img, gt_img) ssim = util.calculate_ssim(sr_img, gt_img) test_results['psnr'].append(psnr) test_results['ssim'].append(ssim) if gt_img.shape[2] == 3: # RGB image sr_img_y = bgr2ycbcr(sr_img / 255., only_y=True) gt_img_y = bgr2ycbcr(gt_img / 255., only_y=True) psnr_y = util.calculate_psnr(sr_img_y * 255, gt_img_y * 255) ssim_y = util.calculate_ssim(sr_img_y * 255, gt_img_y * 255) test_results['psnr_y'].append(psnr_y) test_results['ssim_y'].append(ssim_y) logger.info( '{:20s} - PSNR: {:.6f} dB; SSIM: {:.6f}; PSNR_Y: {:.6f} dB; SSIM_Y: {:.6f}.' .format(img_name, psnr, ssim, psnr_y, ssim_y)) else: logger.info( '{:20s} - PSNR: {:.6f} dB; SSIM: {:.6f}.'.format( img_name, psnr, ssim)) else: logger.info(img_name) if need_GT: # metrics # Average PSNR/SSIM results ave_psnr = sum(test_results['psnr']) / len(test_results['psnr']) ave_ssim = sum(test_results['ssim']) / len(test_results['ssim']) logger.info( '----Average PSNR/SSIM results for {}----\n\tPSNR: {:.6f} dB; SSIM: {:.6f}\n' .format(test_set_name, ave_psnr, ave_ssim)) if test_results['psnr_y'] and test_results['ssim_y']: ave_psnr_y = sum(test_results['psnr_y']) / len( test_results['psnr_y']) ave_ssim_y = sum(test_results['ssim_y']) / len( test_results['ssim_y']) logger.info( '----Y channel, average PSNR/SSIM----\n\tPSNR_Y: {:.6f} dB; SSIM_Y: {:.6f}\n' .format(ave_psnr_y, ave_ssim_y))
img_LR = torch.from_numpy( np.ascontiguousarray(np.transpose( img_LR, (2, 0, 1)))).float().unsqueeze(0).cuda() with torch.no_grad(): begin_time = time.time() frame, sr_base, output = model(img_LR) end_time = time.time() stat_time += (end_time - begin_time) #print(end_time-begin_time) output = util.tensor2img(output.squeeze(0)) frame = util.tensor2img(frame.squeeze(0)) sr_base = util.tensor2img(sr_base.squeeze(0)) # save images save_path_name = osp.join( save_path, '{}_exp{}/{}.png'.format(dataset, exp_name, base_name)) merge = np.concatenate((frame, sr_base, output), axis=1) util.save_img(merge, save_path_name) # calculate PSNR sr_img, gt_img = util.crop_border([output, img_GT], scale) PSNR_avg += util.calculate_psnr(sr_img, gt_img) SSIM_avg += util.calculate_ssim(sr_img, gt_img) print('average PSNR: ', PSNR_avg / len(img_list)) print('average SSIM: ', SSIM_avg / len(img_list)) print('time: ', stat_time / len(img_list))
opt = option.dict_to_nonedict( opt) # Convert to NoneDict, which return None for missing key. dataset = DstlDataset(opt['datasets']['val']) model = create_model(opt).netG.module model.load_state_dict(torch.load(model_path), strict=True) model.eval() model = model.cuda() print('sftgan testing...') for i in range(len(dataset)): data = dataset[i] img_HR = data['HR'] img_LR = data['LR'] output = model(img_LR.cuda()).data img_HR = img_HR.numpy() output = output.squeeze().cpu().numpy() img_HR = denormalize(img_HR) output = denormalize(output) psnr = util.calculate_psnr(img_HR, output) ssim = util.calculate_ssim(img_HR, output) print(psnr, ssim) util.save_img(img_HR, os.path.join('../results', '{}_hr.png'.format(i))) util.save_img(output, os.path.join('../results', '{}_fake.png'.format(i)))
def main(): #### options parser = argparse.ArgumentParser() parser.add_argument('-opt', type=str, help='Path to option YAML file.') parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none', help='job launcher') parser.add_argument('--local_rank', type=int, default=0) args = parser.parse_args() opt = option.parse(args.opt, is_train=True) #### distributed training settings if args.launcher == 'none': # disabled distributed training opt['dist'] = False rank = -1 print('Disabled distributed training.') else: opt['dist'] = True init_dist() world_size = torch.distributed.get_world_size() rank = torch.distributed.get_rank() #### loading resume state if exists if opt['path'].get('resume_state', None): # distributed resuming: all load into default GPU device_id = torch.cuda.current_device() resume_state = torch.load( opt['path']['resume_state'], map_location=lambda storage, loc: storage.cuda(device_id)) option.check_resume(opt, resume_state['iter']) # check resume options else: resume_state = None #### mkdir and loggers if rank <= 0: # normal training (rank -1) OR distributed training (rank 0) if resume_state is None: util.mkdir_and_rename( opt['path'] ['experiments_root']) # rename experiment folder if exists util.mkdirs( (path for key, path in opt['path'].items() if not key == 'experiments_root' and 'pretrain_model' not in key and 'resume' not in key)) # config loggers. Before it, the log will not work util.setup_logger('base', opt['path']['log'], 'train_' + opt['name'], level=logging.INFO, screen=True, tofile=True) logger = logging.getLogger('base') logger.info(option.dict2str(opt)) # tensorboard logger if opt['use_tb_logger'] and 'debug' not in opt['name']: version = float(torch.__version__[0:3]) if version >= 1.1: # PyTorch 1.1 from torch.utils.tensorboard import SummaryWriter else: logger.info( 'You are using PyTorch {}. Tensorboard will use [tensorboardX]' .format(version)) from tensorboardX import SummaryWriter tb_logger = SummaryWriter(log_dir='../tb_logger/' + opt['name']) else: util.setup_logger('base', opt['path']['log'], 'train', level=logging.INFO, screen=True) logger = logging.getLogger('base') # convert to NoneDict, which returns None for missing keys opt = option.dict_to_nonedict(opt) #### random seed seed = opt['train']['manual_seed'] if seed is None: seed = random.randint(1, 10000) if rank <= 0: logger.info('Random seed: {}'.format(seed)) util.set_random_seed(seed) torch.backends.cudnn.benchmark = True # torch.backends.cudnn.deterministic = True #### create train and val dataloader dataset_ratio = 200 # enlarge the size of each epoch for phase, dataset_opt in opt['datasets'].items(): if phase == 'train': train_set = create_dataset(dataset_opt) train_size = int( math.ceil(len(train_set) / dataset_opt['batch_size'])) total_iters = int(opt['train']['niter']) total_epochs = int(math.ceil(total_iters / train_size)) if opt['dist']: train_sampler = DistIterSampler(train_set, world_size, rank, dataset_ratio) total_epochs = int( math.ceil(total_iters / (train_size * dataset_ratio))) else: train_sampler = None train_loader = create_dataloader(train_set, dataset_opt, opt, train_sampler) if rank <= 0: logger.info( 'Number of train images: {:,d}, iters: {:,d}'.format( len(train_set), train_size)) logger.info('Total epochs needed: {:d} for iters {:,d}'.format( total_epochs, total_iters)) elif phase == 'val': val_set = create_dataset(dataset_opt) val_loader = create_dataloader(val_set, dataset_opt, opt, None) if rank <= 0: logger.info('Number of val images in [{:s}]: {:d}'.format( dataset_opt['name'], len(val_set))) else: raise NotImplementedError( 'Phase [{:s}] is not recognized.'.format(phase)) assert train_loader is not None #### create model model = create_model(opt) #### resume training if resume_state: logger.info('Resuming training from epoch: {}, iter: {}.'.format( resume_state['epoch'], resume_state['iter'])) start_epoch = resume_state['epoch'] current_step = resume_state['iter'] model.resume_training(resume_state) # handle optimizers and schedulers else: current_step = 0 start_epoch = 0 #### training logger.info('Start training from epoch: {:d}, iter: {:d}'.format( start_epoch, current_step)) for epoch in range(start_epoch, total_epochs + 1): if opt['dist']: train_sampler.set_epoch(epoch) for _, train_data in enumerate(train_loader): current_step += 1 if current_step > total_iters: break #### update learning rate model.update_learning_rate(current_step, warmup_iter=opt['train']['warmup_iter']) #### training model.feed_data(train_data) model.optimize_parameters(current_step) #### log if current_step % opt['logger']['print_freq'] == 0: logs = model.get_current_log() message = '[epoch:{:3d}, iter:{:8,d}, lr:('.format( epoch, current_step) for v in model.get_current_learning_rate(): message += '{:.3e},'.format(v) message += ')] ' for k, v in logs.items(): message += '{:s}: {:.4e} '.format(k, v) # tensorboard logger if opt['use_tb_logger'] and 'debug' not in opt['name']: if rank <= 0: tb_logger.add_scalar(k, v, current_step) if rank <= 0: logger.info(message) #### validation if opt['datasets'].get( 'val', None) and current_step % opt['train']['val_freq'] == 0: if opt['model'] in [ 'sr', 'srgan' ] and rank <= 0: # image restoration validation # does not support multi-GPU validation pbar = util.ProgressBar(len(val_loader)) avg_psnr = 0. avg_ssim = 0. idx = 0 for val_data in val_loader: idx += 1 img_name = os.path.splitext( os.path.basename(val_data['LQ_path'][0]))[0] img_dir = os.path.join(opt['path']['val_images'], img_name) util.mkdir(img_dir) # we don't need GT data to be run with the model gt_img = util.tensor2img(val_data['GT']) # uint8 val_data['GT'] = torch.tensor(0) import sys #print(val_data['LQ'].shape) model.feed_data(val_data, need_GT=False) model.test() visuals = model.get_current_visuals() sr_img = util.tensor2img(visuals['rlt']) # uint8 if "mask" in visuals: mask = util.tensor2img(visuals['mask']) # print(mask) # Save SR mask for reference save_img_path = os.path.join( img_dir, '{:s}_{:d}_mask.png'.format( img_name, current_step)) util.save_img(mask, save_img_path) #gt_mask_dir = os.path.join(opt['path']['gt_mask'], img_name) #gt_mask = util.tensor2img(val_data['GT_mask']) #print(gt_mask) #import cv2 #cv2.imshow("window", gt_mask) #cv2.waitKey() #util.mkdir(gt_mask_dir) #save_img_path = os.path.join(gt_mask_dir, # '{:s}_{:d}_GT_mask.png'.format(img_name, current_step)) #util.save_img(gt_mask, save_img_path) # Save SR images for reference save_img_path = os.path.join( img_dir, '{:s}_{:d}.png'.format(img_name, current_step)) util.save_img(sr_img, save_img_path) # calculate PSNR sr_img, gt_img = util.crop_border([sr_img, gt_img], opt['scale']) avg_psnr += util.calculate_psnr(sr_img, gt_img) avg_ssim += util.calculate_ssim(sr_img, gt_img) pbar.update('Test {}'.format(img_name)) avg_psnr = avg_psnr / idx avg_ssim = avg_ssim / idx # log logger.info('# Validation # PSNR: {:.4e}'.format(avg_psnr)) logger.info('# Validation # SSIM: {:.4e}'.format(avg_ssim)) # tensorboard logger if opt['use_tb_logger'] and 'debug' not in opt['name']: tb_logger.add_scalar('psnr', avg_psnr, current_step) tb_logger.add_scalar('ssim', avg_ssim, current_step) else: # video restoration validation if opt['dist']: # multi-GPU testing psnr_rlt = {} # with border and center frames if rank == 0: pbar = util.ProgressBar(len(val_set)) for idx in range(rank, len(val_set), world_size): val_data = val_set[idx] val_data['LQs'].unsqueeze_(0) val_data['GT'].unsqueeze_(0) folder = val_data['folder'] idx_d, max_idx = val_data['idx'].split('/') idx_d, max_idx = int(idx_d), int(max_idx) if psnr_rlt.get(folder, None) is None: psnr_rlt[folder] = torch.zeros( max_idx, dtype=torch.float32, device='cuda') # tmp = torch.zeros(max_idx, dtype=torch.float32, device='cuda') model.feed_data(val_data) model.test() visuals = model.get_current_visuals() rlt_img = util.tensor2img(visuals['rlt']) # uint8 gt_img = util.tensor2img(visuals['GT']) # uint8 # calculate PSNR psnr_rlt[folder][idx_d] = util.calculate_psnr( rlt_img, gt_img) if rank == 0: for _ in range(world_size): pbar.update('Test {} - {}/{}'.format( folder, idx_d, max_idx)) # # collect data for _, v in psnr_rlt.items(): dist.reduce(v, 0) dist.barrier() if rank == 0: psnr_rlt_avg = {} psnr_total_avg = 0. for k, v in psnr_rlt.items(): psnr_rlt_avg[k] = torch.mean(v).cpu().item() psnr_total_avg += psnr_rlt_avg[k] psnr_total_avg /= len(psnr_rlt) log_s = '# Validation # PSNR: {:.4e}:'.format( psnr_total_avg) for k, v in psnr_rlt_avg.items(): log_s += ' {}: {:.4e}'.format(k, v) logger.info(log_s) if opt['use_tb_logger'] and 'debug' not in opt[ 'name']: tb_logger.add_scalar('psnr_avg', psnr_total_avg, current_step) for k, v in psnr_rlt_avg.items(): tb_logger.add_scalar(k, v, current_step) else: pbar = util.ProgressBar(len(val_loader)) psnr_rlt = {} # with border and center frames psnr_rlt_avg = {} psnr_total_avg = 0. for val_data in val_loader: folder = val_data['folder'][0] idx_d = val_data['idx'].item() # border = val_data['border'].item() if psnr_rlt.get(folder, None) is None: psnr_rlt[folder] = [] model.feed_data(val_data) model.test() visuals = model.get_current_visuals() rlt_img = util.tensor2img(visuals['rlt']) # uint8 gt_img = util.tensor2img(visuals['GT']) # uint8 # calculate PSNR psnr = util.calculate_psnr(rlt_img, gt_img) psnr_rlt[folder].append(psnr) pbar.update('Test {} - {}'.format(folder, idx_d)) for k, v in psnr_rlt.items(): psnr_rlt_avg[k] = sum(v) / len(v) psnr_total_avg += psnr_rlt_avg[k] psnr_total_avg /= len(psnr_rlt) log_s = '# Validation # PSNR: {:.4e}:'.format( psnr_total_avg) for k, v in psnr_rlt_avg.items(): log_s += ' {}: {:.4e}'.format(k, v) logger.info(log_s) if opt['use_tb_logger'] and 'debug' not in opt['name']: tb_logger.add_scalar('psnr_avg', psnr_total_avg, current_step) for k, v in psnr_rlt_avg.items(): tb_logger.add_scalar(k, v, current_step) #### save models and training states if current_step % opt['logger']['save_checkpoint_freq'] == 0: if rank <= 0: logger.info('Saving models and training states.') model.save(current_step) model.save_training_state(epoch, current_step) if rank <= 0: logger.info('Saving the final model.') model.save('latest') logger.info('End of training.') tb_logger.close()
def main(): # options parser = argparse.ArgumentParser() parser.add_argument('-opt', type=str, default='options/test/test_ppon.json', help='Path to options JSON file.') opt = option.parse(parser.parse_args().opt, is_train=False) util.mkdirs((path for key, path in opt['path'].items() if not key == 'pretrain_model_G')) opt = option.dict_to_nonedict(opt) util.setup_logger(None, opt['path']['log'], 'test.log', level=logging.INFO, screen=True) logger = logging.getLogger('base') logger.info(option.dict2str(opt)) # Create test dataset and dataloader test_loaders = [] for phase, dataset_opt in sorted(opt['datasets'].items()): test_set = create_dataset(dataset_opt) test_loader = create_dataloader(test_set, dataset_opt) logger.info('Number of test images in [{:s}]: {:d}'.format(dataset_opt['name'], len(test_set))) test_loaders.append(test_loader) # Create model model = create_model(opt) for test_loader in test_loaders: test_set_name = test_loader.dataset.opt['name'] logger.info('\nTesting [{:s}]...'.format(test_set_name)) test_start_time = time.time() dataset_dir = os.path.join(opt['path']['results_root'], test_set_name) util.mkdir(dataset_dir) test_results = OrderedDict() test_results['psnr'] = [] test_results['ssim'] = [] test_results['psnr_y'] = [] test_results['ssim_y'] = [] for data in test_loader: need_HR = False if test_loader.dataset.opt['dataroot_HR'] is None else True model.feed_data(data, need_HR=need_HR) img_path = data['LR_path'][0] img_name = os.path.splitext(os.path.basename(img_path))[0] model.test() # test visuals = model.get_current_visuals(need_HR=need_HR) img_c = util.tensor2img(visuals['img_c']) # uint8 img_s = util.tensor2img(visuals['img_s']) # uint8 img_p = util.tensor2img(visuals['img_p']) # uint8 # save images suffix = opt['suffix'] if suffix: save_c_img_path = os.path.join(dataset_dir, img_name + suffix + '_c.png') save_s_img_path = os.path.join(dataset_dir, img_name + suffix + '_s.png') save_p_img_path = os.path.join(dataset_dir, img_name + suffix + '_p.png') else: save_c_img_path = os.path.join(dataset_dir, img_name + '_c.png') save_s_img_path = os.path.join(dataset_dir, img_name + '_s.png') save_p_img_path = os.path.join(dataset_dir, img_name + '_p.png') util.save_img(img_c, save_c_img_path) util.save_img(img_s, save_s_img_path) util.save_img(img_p, save_p_img_path) # calculate PSNR and SSIM if need_HR: gt_img = util.tensor2img(visuals['HR']) gt_img = gt_img / 255. sr_img = img_c / 255. crop_border = test_loader.dataset.opt['scale'] cropped_sr_img = sr_img[crop_border:-crop_border, crop_border:-crop_border, :] cropped_gt_img = gt_img[crop_border:-crop_border, crop_border:-crop_border, :] psnr = util.calculate_psnr(cropped_sr_img * 255, cropped_gt_img * 255) ssim = util.calculate_ssim(cropped_sr_img * 255, cropped_gt_img * 255) test_results['psnr'].append(psnr) test_results['ssim'].append(ssim) if gt_img.shape[2] == 3: # RGB image sr_img_y = bgr2ycbcr(sr_img, only_y=True) gt_img_y = bgr2ycbcr(gt_img, only_y=True) cropped_sr_img_y = sr_img_y[crop_border:-crop_border, crop_border:-crop_border] cropped_gt_img_y = gt_img_y[crop_border:-crop_border, crop_border:-crop_border] psnr_y = util.calculate_psnr(cropped_sr_img_y * 255, cropped_gt_img_y * 255) ssim_y = util.calculate_ssim(cropped_sr_img_y * 255, cropped_gt_img_y * 255) test_results['psnr_y'].append(psnr_y) test_results['ssim_y'].append(ssim_y) logger.info('{:20s} - PSNR: {:.6f} dB; SSIM: {:.6f}; PSNR_Y: {:.6f} dB; SSIM_Y: {:.6f}.'\ .format(img_name, psnr, ssim, psnr_y, ssim_y)) else: logger.info('{:20s} - PSNR: {:.6f} dB; SSIM: {:.6f}.'.format(img_name, psnr, ssim)) else: logger.info(img_name) if need_HR: # metrics # Average PSNR/SSIM results ave_psnr = sum(test_results['psnr']) / len(test_results['psnr']) ave_ssim = sum(test_results['ssim']) / len(test_results['ssim']) logger.info('----Average PSNR/SSIM results for {}----\n\tPSNR: {:.6f} dB; SSIM: {:.6f}\n'\ .format(test_set_name, ave_psnr, ave_ssim)) if test_results['psnr_y'] and test_results['ssim_y']: ave_psnr_y = sum(test_results['psnr_y']) / len(test_results['psnr_y']) ave_ssim_y = sum(test_results['ssim_y']) / len(test_results['ssim_y']) logger.info('----Y channel, average PSNR/SSIM----\n\tPSNR_Y: {:.6f} dB; SSIM_Y: {:.6f}\n'\ .format(ave_psnr_y, ave_ssim_y))
else: # normalized_pixel_STD = 0 pixel_STD = 0 # Save GT image for reference: util.save_img( (255 * util.tensor2img( visuals['HR'], out_type=np.float32)).astype( np.uint8), os.path.join( dataset_dir, img_name + '_HR_STD%.3f_SR_STD%.3f.png' % (HR_STD, pixel_STD))) sr_img *= 255. if LATENT_DISTRIBUTION in NON_ARBITRARY_Z_INPUTS or cur_channel_cur_Z == 0: psnr = util.calculate_psnr(sr_img, gt_img) ssim = util.calculate_ssim(sr_img, gt_img) test_results['psnr'].append(psnr) test_results['ssim'].append(ssim) if SAVE_IMAGE_COLLAGE: if len(test_set) > 1: margins2crop = ((np.array(sr_img.shape[:2]) - per_image_saved_patch) / 2).astype( np.int32) else: margins2crop = [0, 0] image_collage[-1].append( np.clip(util.crop_center(sr_img, margins2crop), 0, 255).astype(np.uint8)) if LATENT_DISTRIBUTION in NON_ARBITRARY_Z_INPUTS or cur_channel_cur_Z == 0: # Save GT HR images: GT_image_collage[-1].append(
def main(): opt = option.parse("options/train/train_degnet.json", is_train=True) opt = option.dict_to_nonedict(opt) # Convert to NoneDict, which return None for missing key. # train from scratch OR resume training if opt['path']['resume_state']: # resuming training resume_state = torch.load(opt['path']['resume_state']) else: # training from scratch resume_state = None util.mkdir_and_rename(opt['path']['experiments_root']) # rename old folder if exists util.mkdirs((path for key, path in opt['path'].items() if not key == 'experiments_root' and 'pretrain_model' not in key and 'resume' not in key)) # config loggers. Before it, the log will not work util.setup_logger(None, opt['path']['log'], 'train', level=logging.INFO, screen=True) util.setup_logger('val', opt['path']['log'], 'val', level=logging.INFO) logger = logging.getLogger('base') if resume_state: logger.info('Resuming training from epoch: {}, iter: {}.'.format( resume_state['epoch'], resume_state['iter'])) option.check_resume(opt) # check resume options logger.info(option.dict2str(opt)) # tensorboard logger if opt['use_tb_logger'] and 'debug' not in opt['name']: from tensorboardX import SummaryWriter tb_logger = SummaryWriter(log_dir='./tb_logger/' + opt['name']) # random seed seed = opt['train']['manual_seed'] if seed is None: seed = random.randint(1, 10000) logger.info('Random seed: {}'.format(seed)) util.set_random_seed(seed) torch.backends.cudnn.benckmark = True # torch.backends.cudnn.deterministic = True # create train and val dataloader for phase, dataset_opt in opt['datasets'].items(): if phase == 'train': train_set = create_dataset(dataset_opt) train_size = int(math.ceil(len(train_set) / dataset_opt['batch_size'])) logger.info('Number of train images: {:,d}, iters: {:,d}'.format( len(train_set), train_size)) total_iters = int(opt['train']['niter']) total_epochs = int(math.ceil(total_iters / train_size)) logger.info('Total epochs needed: {:d} for iters {:,d}'.format( total_epochs, total_iters)) train_loader = create_dataloader(train_set, dataset_opt) elif phase == 'val': val_set = create_dataset(dataset_opt) val_loader = create_dataloader(val_set, dataset_opt) logger.info('Number of val images in [{:s}]: {:d}'.format(dataset_opt['name'], len(val_set))) else: raise NotImplementedError('Phase [{:s}] is not recognized.'.format(phase)) assert train_loader is not None # create model model = create_model(opt) # resume training if resume_state: start_epoch = resume_state['epoch'] current_step = resume_state['iter'] model.resume_training(resume_state) # handle optimizers and schedulers else: current_step = 0 start_epoch = 0 # training logger.info('Start training from epoch: {:d}, iter: {:d}'.format(start_epoch, current_step)) for epoch in range(start_epoch, total_epochs): for _, train_data in enumerate(train_loader): current_step += 1 if current_step > total_iters: break # training model.feed_data(train_data) model.optimize_parameters(current_step) # update learning rate model.update_learning_rate() # log if current_step % opt['logger']['print_freq'] == 0: logs = model.get_current_log() message = '<epoch:{:3d}, iter:{:8,d}, lr:{:.3e}> '.format( epoch, current_step, model.get_current_learning_rate()) for k, v in logs.items(): message += '{:s}: {:.4e} '.format(k, v) if opt['use_tb_logger'] and 'debug' not in opt['name']: tb_logger.add_scalar(k, v, current_step) logger.info(message) # validation if current_step % opt['train']['val_freq'] == 0: avg_psnr = 0.0 avg_ssim = 0.0 idx = 0 for val_data in val_loader: idx += 1 img_name = os.path.splitext(os.path.basename(val_data['LR_path'][0]))[0] img_dir = os.path.join(opt['path']['val_images'], img_name) util.mkdir(img_dir) model.feed_data(val_data) model.test() visuals = model.get_current_visuals() lr_img = util.tensor2img(visuals['FLR']) # uint8 gt_img = util.tensor2img(visuals['LR']) # uint8 # Save degradation map for reference #save_DM_path = os.path.join(img_dir, '{:s}_{:d}.png'.format(img_name, current_step)) #degradation_img = util.tensor2img(visuals['DM']) # uint8 #util.save_img(degradation_img, save_DM_path[:-4] + '_degradation.png') # Save fake LR save_LR_path = opt['path']['root'] + 'validations/' + opt['name'] + '/{:d}/'.format(current_step) util.mkdir(save_LR_path) save_LR_img_path = os.path.join(save_LR_path, '{:s}.png'.format( \ img_name)) util.save_img(lr_img, save_LR_img_path) # calculate PSNR crop_size = opt['scale'] gt_img = gt_img / 255. lr_img = lr_img / 255. cropped_lr_img = lr_img[crop_size:-crop_size, crop_size:-crop_size, :] cropped_gt_img = gt_img[crop_size:-crop_size, crop_size:-crop_size, :] avg_psnr += util.calculate_psnr(cropped_lr_img * 255, cropped_gt_img * 255) avg_ssim += util.calculate_ssim(cropped_lr_img * 255, cropped_gt_img * 255) avg_psnr = avg_psnr / idx avg_ssim = avg_ssim / idx # log logger.info('# Validation # PSNR: {:.4e} SSIM: {:.4e}'.format(avg_psnr, avg_ssim)) logger_val = logging.getLogger('val') # validation logger logger_val.info('<epoch:{:3d}, iter:{:8,d}> psnr: {:.4e} ssim: {:.4e}'.format( epoch, current_step, avg_psnr, avg_ssim)) if opt['use_tb_logger'] and 'debug' not in opt['name']: tb_logger.add_scalar('psnr', avg_psnr, current_step) tb_logger.add_scalar('ssim', avg_ssim, current_step) # save models and training states if current_step % opt['logger']['save_checkpoint_freq'] == 0: logger.info('Saving models and training states.') model.save(current_step) model.save_training_state(epoch, current_step) logger.info('Saving the final model.') model.save('latest') logger.info('End of training.')
if need_HR: gt_img = util.tensor2img(visuals['HR']) gt_img = gt_img / 255. sr_img = sr_img / 255. if opt['val_lpips']: lpips = visuals['LPIPS'].numpy() crop_border = test_loader.dataset.opt['scale'] cropped_sr_img = sr_img[crop_border:-crop_border, crop_border:-crop_border, :] cropped_gt_img = gt_img[crop_border:-crop_border, crop_border:-crop_border, :] psnr = util.calculate_psnr(cropped_sr_img * 255, cropped_gt_img * 255) ssim = util.calculate_ssim(cropped_sr_img * 255, cropped_gt_img * 255) test_results['psnr'].append(psnr) test_results['ssim'].append(ssim) if opt['val_lpips']: test_results['lpips'].append(lpips) if gt_img.shape[2] == 3: # RGB image sr_img_y = bgr2ycbcr(sr_img, only_y=True) gt_img_y = bgr2ycbcr(gt_img, only_y=True) cropped_sr_img_y = sr_img_y[crop_border:-crop_border, crop_border:-crop_border] cropped_gt_img_y = gt_img_y[crop_border:-crop_border, crop_border:-crop_border] psnr_y = util.calculate_psnr(cropped_sr_img_y * 255, cropped_gt_img_y * 255) ssim_y = util.calculate_ssim(cropped_sr_img_y * 255,
quantized_image = util.tensor2img( model.Return_Compressed( gt_im_YCbCr.to(model.device)), out_type=np.uint8, min_max=[0, 255], chroma_mode=chroma_mode) # quantized_image = util.tensor2img(model.jpeg_extractor(model.jpeg_compressor(data['Uncomp'])), out_type=np.uint8,min_max=[0, 255],chroma_mode=chroma_mode) if SAVE_QUANTIZED: util.save_img( quantized_image, os.path.join(dataset_dir + '_Quant', img_name + suffix + '.png')) test_results['psnr_quantized'].append( util.calculate_psnr(quantized_image, gt_img)) test_results['ssim_quantized'].append( util.calculate_ssim(quantized_image, gt_img)) psnr = util.calculate_psnr(sr_img, gt_img) ssim = util.calculate_ssim(sr_img, gt_img) test_results['psnr'].append(psnr) test_results['ssim'].append(ssim) if SAVE_IMAGE_COLLAGE: if len(test_set) > 1: margins2crop = ((np.array(sr_img.shape[:2]) - per_image_saved_patch) / 2).astype( np.int32) else: margins2crop = [0, 0] image_collage[-1].append( np.clip(util.crop_center(sr_img, margins2crop), 0, 255).astype(np.uint8)) if LATENT_DISTRIBUTION in NON_ARBITRARY_Z_INPUTS or cur_channel_cur_Z == 0:
sr_img = util.tensor2img(visuals['rlt']) # uint8 # save images suffix = opt['suffix'] if suffix: save_img_path = osp.join(dataset_dir, img_name + suffix + '.png') else: save_img_path = osp.join(dataset_dir, img_name + '.png') util.save_img(sr_img, save_img_path) # calculate PSNR and SSIM if need_GT: gt_img = util.tensor2img(visuals['GT']) sr_img, gt_img = util.crop_border([sr_img, gt_img], opt['scale']) psnr = util.calculate_psnr(sr_img, gt_img) ssim = util.calculate_ssim(sr_img, gt_img) test_results['psnr'].append(psnr) test_results['ssim'].append(ssim) if gt_img.shape[2] == 3: # RGB image sr_img_y = bgr2ycbcr(sr_img / 255., only_y=True) gt_img_y = bgr2ycbcr(gt_img / 255., only_y=True) psnr_y = util.calculate_psnr(sr_img_y * 255, gt_img_y * 255) ssim_y = util.calculate_ssim(sr_img_y * 255, gt_img_y * 255) test_results['psnr_y'].append(psnr_y) test_results['ssim_y'].append(ssim_y) logger.info( '{:20s} - PSNR: {:.6f} dB; SSIM: {:.6f}; PSNR_Y: {:.6f} dB; SSIM_Y: {:.6f}.' .format(img_name, psnr, ssim, psnr_y, ssim_y)) else:
def main(): ############################################ # # set options # ############################################ parser = argparse.ArgumentParser() parser.add_argument('--opt', type=str, help='Path to option YAML file.') parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none', help='job launcher') parser.add_argument('--local_rank', type=int, default=0) args = parser.parse_args() opt = option.parse(args.opt, is_train=True) ############################################ # # distributed training settings # ############################################ if args.launcher == 'none': # disabled distributed training opt['dist'] = False rank = -1 print('Disabled distributed training.') else: opt['dist'] = True init_dist() world_size = torch.distributed.get_world_size() rank = torch.distributed.get_rank() print("Rank:", rank) print("------------------DIST-------------------------") ############################################ # # loading resume state if exists # ############################################ if opt['path'].get('resume_state', None): # distributed resuming: all load into default GPU device_id = torch.cuda.current_device() resume_state = torch.load( opt['path']['resume_state'], map_location=lambda storage, loc: storage.cuda(device_id)) option.check_resume(opt, resume_state['iter']) # check resume options else: resume_state = None ############################################ # # mkdir and loggers # ############################################ if rank <= 0: # normal training (rank -1) OR distributed training (rank 0) if resume_state is None: util.mkdir_and_rename( opt['path'] ['experiments_root']) # rename experiment folder if exists util.mkdirs( (path for key, path in opt['path'].items() if not key == 'experiments_root' and 'pretrain_model' not in key and 'resume' not in key)) # config loggers. Before it, the log will not work util.setup_logger('base', opt['path']['log'], 'train_' + opt['name'], level=logging.INFO, screen=True, tofile=True) util.setup_logger('base_val', opt['path']['log'], 'val_' + opt['name'], level=logging.INFO, screen=True, tofile=True) logger = logging.getLogger('base') logger_val = logging.getLogger('base_val') logger.info(option.dict2str(opt)) # tensorboard logger if opt['use_tb_logger'] and 'debug' not in opt['name']: version = float(torch.__version__[0:3]) if version >= 1.1: # PyTorch 1.1 from torch.utils.tensorboard import SummaryWriter else: logger.info( 'You are using PyTorch {}. Tensorboard will use [tensorboardX]' .format(version)) from tensorboardX import SummaryWriter tb_logger = SummaryWriter(log_dir='../tb_logger/' + opt['name']) else: # config loggers. Before it, the log will not work util.setup_logger('base', opt['path']['log'], 'train_', level=logging.INFO, screen=True) print("set train log") util.setup_logger('base_val', opt['path']['log'], 'val_', level=logging.INFO, screen=True) print("set val log") logger = logging.getLogger('base') logger_val = logging.getLogger('base_val') # convert to NoneDict, which returns None for missing keys opt = option.dict_to_nonedict(opt) #### random seed seed = opt['train']['manual_seed'] if seed is None: seed = random.randint(1, 10000) if rank <= 0: logger.info('Random seed: {}'.format(seed)) util.set_random_seed(seed) torch.backends.cudnn.benchmark = True # torch.backends.cudnn.deterministic = True ############################################ # # create train and val dataloader # ############################################ #### # dataset_ratio = 200 # enlarge the size of each epoch, todo: what it is dataset_ratio = 1 # enlarge the size of each epoch, todo: what it is for phase, dataset_opt in opt['datasets'].items(): if phase == 'train': train_set = create_dataset(dataset_opt) train_size = int( math.ceil(len(train_set) / dataset_opt['batch_size'])) # total_iters = int(opt['train']['niter']) # total_epochs = int(math.ceil(total_iters / train_size)) total_iters = train_size total_epochs = int(opt['train']['epoch']) if opt['dist']: train_sampler = DistIterSampler(train_set, world_size, rank, dataset_ratio) # total_epochs = int(math.ceil(total_iters / (train_size * dataset_ratio))) total_epochs = int(opt['train']['epoch']) if opt['train']['enable'] == False: total_epochs = 1 else: train_sampler = None train_loader = create_dataloader(train_set, dataset_opt, opt, train_sampler) if rank <= 0: logger.info( 'Number of train images: {:,d}, iters: {:,d}'.format( len(train_set), train_size)) logger.info('Total epochs needed: {:d} for iters {:,d}'.format( total_epochs, total_iters)) elif phase == 'val': val_set = create_dataset(dataset_opt) val_loader = create_dataloader(val_set, dataset_opt, opt, None) if rank <= 0: logger.info('Number of val images in [{:s}]: {:d}'.format( dataset_opt['name'], len(val_set))) else: raise NotImplementedError( 'Phase [{:s}] is not recognized.'.format(phase)) assert train_loader is not None ############################################ # # create model # ############################################ #### model = create_model(opt) print("Model Created! ") #### resume training if resume_state: logger.info('Resuming training from epoch: {}, iter: {}.'.format( resume_state['epoch'], resume_state['iter'])) start_epoch = resume_state['epoch'] current_step = resume_state['iter'] model.resume_training(resume_state) # handle optimizers and schedulers else: current_step = 0 start_epoch = 0 print("Not Resume Training") ############################################ # # training # ############################################ #### #### logger.info('Start training from epoch: {:d}, iter: {:d}'.format( start_epoch, current_step)) Avg_train_loss = AverageMeter() # total if (opt['train']['pixel_criterion'] == 'cb+ssim'): Avg_train_loss_pix = AverageMeter() Avg_train_loss_ssim = AverageMeter() elif (opt['train']['pixel_criterion'] == 'cb+ssim+vmaf'): Avg_train_loss_pix = AverageMeter() Avg_train_loss_ssim = AverageMeter() Avg_train_loss_vmaf = AverageMeter() elif (opt['train']['pixel_criterion'] == 'ssim'): Avg_train_loss_ssim = AverageMeter() elif (opt['train']['pixel_criterion'] == 'msssim'): Avg_train_loss_msssim = AverageMeter() elif (opt['train']['pixel_criterion'] == 'cb+msssim'): Avg_train_loss_pix = AverageMeter() Avg_train_loss_msssim = AverageMeter() saved_total_loss = 10e10 saved_total_PSNR = -1 for epoch in range(start_epoch, total_epochs): ############################################ # # Start a new epoch # ############################################ # Turn into training mode #model = model.train() # reset total loss Avg_train_loss.reset() current_step = 0 if (opt['train']['pixel_criterion'] == 'cb+ssim'): Avg_train_loss_pix.reset() Avg_train_loss_ssim.reset() elif (opt['train']['pixel_criterion'] == 'cb+ssim+vmaf'): Avg_train_loss_pix.reset() Avg_train_loss_ssim.reset() Avg_train_loss_vmaf.reset() elif (opt['train']['pixel_criterion'] == 'ssim'): Avg_train_loss_ssim = AverageMeter() elif (opt['train']['pixel_criterion'] == 'msssim'): Avg_train_loss_msssim = AverageMeter() elif (opt['train']['pixel_criterion'] == 'cb+msssim'): Avg_train_loss_pix = AverageMeter() Avg_train_loss_msssim = AverageMeter() if opt['dist']: train_sampler.set_epoch(epoch) for train_idx, train_data in enumerate(train_loader): if 'debug' in opt['name']: img_dir = os.path.join(opt['path']['train_images']) util.mkdir(img_dir) LQ = train_data['LQs'] GT = train_data['GT'] GT_img = util.tensor2img(GT) # uint8 save_img_path = os.path.join( img_dir, '{:4d}_{:s}.png'.format(train_idx, 'debug_GT')) util.save_img(GT_img, save_img_path) for i in range(5): LQ_img = util.tensor2img(LQ[0, i, ...]) # uint8 save_img_path = os.path.join( img_dir, '{:4d}_{:s}_{:1d}.png'.format(train_idx, 'debug_LQ', i)) util.save_img(LQ_img, save_img_path) if (train_idx >= 3): break if opt['train']['enable'] == False: message_train_loss = 'None' break current_step += 1 if current_step > total_iters: print("Total Iteration Reached !") break #### update learning rate if opt['train']['lr_scheme'] == 'ReduceLROnPlateau': pass else: model.update_learning_rate( current_step, warmup_iter=opt['train']['warmup_iter']) #### training model.feed_data(train_data) # if opt['train']['lr_scheme'] == 'ReduceLROnPlateau': # model.optimize_parameters_without_schudlue(current_step) # else: model.optimize_parameters(current_step) if (opt['train']['pixel_criterion'] == 'cb+ssim'): Avg_train_loss.update(model.log_dict['total_loss'], 1) Avg_train_loss_pix.update(model.log_dict['l_pix'], 1) Avg_train_loss_ssim.update(model.log_dict['ssim_loss'], 1) elif (opt['train']['pixel_criterion'] == 'cb+ssim+vmaf'): Avg_train_loss.update(model.log_dict['total_loss'], 1) Avg_train_loss_pix.update(model.log_dict['l_pix'], 1) Avg_train_loss_ssim.update(model.log_dict['ssim_loss'], 1) Avg_train_loss_vmaf.update(model.log_dict['vmaf_loss'], 1) elif (opt['train']['pixel_criterion'] == 'ssim'): Avg_train_loss.update(model.log_dict['total_loss'], 1) Avg_train_loss_ssim.update(model.log_dict['ssim_loss'], 1) elif (opt['train']['pixel_criterion'] == 'msssim'): Avg_train_loss.update(model.log_dict['total_loss'], 1) Avg_train_loss_msssim.update(model.log_dict['msssim_loss'], 1) elif (opt['train']['pixel_criterion'] == 'cb+msssim'): Avg_train_loss.update(model.log_dict['total_loss'], 1) Avg_train_loss_pix.update(model.log_dict['l_pix'], 1) Avg_train_loss_msssim.update(model.log_dict['msssim_loss'], 1) else: Avg_train_loss.update(model.log_dict['l_pix'], 1) # add total train loss if (opt['train']['pixel_criterion'] == 'cb+ssim'): message_train_loss = ' pix_avg_loss: {:.4e}'.format( Avg_train_loss_pix.avg) message_train_loss += ' ssim_avg_loss: {:.4e}'.format( Avg_train_loss_ssim.avg) message_train_loss += ' total_avg_loss: {:.4e}'.format( Avg_train_loss.avg) elif (opt['train']['pixel_criterion'] == 'cb+ssim+vmaf'): message_train_loss = ' pix_avg_loss: {:.4e}'.format( Avg_train_loss_pix.avg) message_train_loss += ' ssim_avg_loss: {:.4e}'.format( Avg_train_loss_ssim.avg) message_train_loss += ' vmaf_avg_loss: {:.4e}'.format( Avg_train_loss_vmaf.avg) message_train_loss += ' total_avg_loss: {:.4e}'.format( Avg_train_loss.avg) elif (opt['train']['pixel_criterion'] == 'ssim'): message_train_loss = ' ssim_avg_loss: {:.4e}'.format( Avg_train_loss_ssim.avg) message_train_loss += ' total_avg_loss: {:.4e}'.format( Avg_train_loss.avg) elif (opt['train']['pixel_criterion'] == 'msssim'): message_train_loss = ' msssim_avg_loss: {:.4e}'.format( Avg_train_loss_msssim.avg) message_train_loss += ' total_avg_loss: {:.4e}'.format( Avg_train_loss.avg) elif (opt['train']['pixel_criterion'] == 'cb+msssim'): message_train_loss = ' pix_avg_loss: {:.4e}'.format( Avg_train_loss_pix.avg) message_train_loss += ' msssim_avg_loss: {:.4e}'.format( Avg_train_loss_msssim.avg) message_train_loss += ' total_avg_loss: {:.4e}'.format( Avg_train_loss.avg) else: message_train_loss = ' train_avg_loss: {:.4e}'.format( Avg_train_loss.avg) #### log if current_step % opt['logger']['print_freq'] == 0: logs = model.get_current_log() message = '[epoch:{:3d}, iter:{:8,d}, lr:('.format( epoch, current_step) for v in model.get_current_learning_rate(): message += '{:.3e},'.format(v) message += ')] ' for k, v in logs.items(): message += '{:s}: {:.4e} '.format(k, v) # tensorboard logger if opt['use_tb_logger'] and 'debug' not in opt['name']: if rank <= 0: tb_logger.add_scalar(k, v, current_step) message += message_train_loss if rank <= 0: logger.info(message) ############################################ # # end of one epoch, save epoch model # ############################################ #### save models and training states # if current_step % opt['logger']['save_checkpoint_freq'] == 0: # if rank <= 0: # logger.info('Saving models and training states.') # model.save(current_step) # model.save('latest') # # model.save_training_state(epoch, current_step) # # todo delete previous weights # previous_step = current_step - opt['logger']['save_checkpoint_freq'] # save_filename = '{}_{}.pth'.format(previous_step, 'G') # save_path = os.path.join(opt['path']['models'], save_filename) # if os.path.exists(save_path): # os.remove(save_path) if epoch == 1: save_filename = '{:04d}_{}.pth'.format(0, 'G') save_path = os.path.join(opt['path']['models'], save_filename) if os.path.exists(save_path): os.remove(save_path) save_filename = '{:04d}_{}.pth'.format(epoch - 1, 'G') save_path = os.path.join(opt['path']['models'], save_filename) if os.path.exists(save_path): os.remove(save_path) if rank <= 0: logger.info('Saving models and training states.') save_filename = '{:04d}'.format(epoch) model.save(save_filename) # model.save('latest') # model.save_training_state(epoch, current_step) ############################################ # # end of one epoch, do validation # ############################################ #### validation #if opt['datasets'].get('val', None) and current_step % opt['train']['val_freq'] == 0: if opt['datasets'].get('val', None): if opt['model'] in [ 'sr', 'srgan' ] and rank <= 0: # image restoration validation # does not support multi-GPU validation pbar = util.ProgressBar(len(val_loader)) avg_psnr = 0. idx = 0 for val_data in val_loader: idx += 1 img_name = os.path.splitext( os.path.basename(val_data['LQ_path'][0]))[0] img_dir = os.path.join(opt['path']['val_images'], img_name) util.mkdir(img_dir) model.feed_data(val_data) model.test() visuals = model.get_current_visuals() sr_img = util.tensor2img(visuals['rlt']) # uint8 gt_img = util.tensor2img(visuals['GT']) # uint8 # Save SR images for reference save_img_path = os.path.join( img_dir, '{:s}_{:d}.png'.format(img_name, current_step)) #util.save_img(sr_img, save_img_path) # calculate PSNR sr_img, gt_img = util.crop_border([sr_img, gt_img], opt['scale']) avg_psnr += util.calculate_psnr(sr_img, gt_img) pbar.update('Test {}'.format(img_name)) avg_psnr = avg_psnr / idx # log logger.info('# Validation # PSNR: {:.4e}'.format(avg_psnr)) # tensorboard logger if opt['use_tb_logger'] and 'debug' not in opt['name']: tb_logger.add_scalar('psnr', avg_psnr, current_step) else: # video restoration validation if opt['dist']: # todo : multi-GPU testing psnr_rlt = {} # with border and center frames psnr_rlt_avg = {} psnr_total_avg = 0. ssim_rlt = {} # with border and center frames ssim_rlt_avg = {} ssim_total_avg = 0. val_loss_rlt = {} val_loss_rlt_avg = {} val_loss_total_avg = 0. if rank == 0: pbar = util.ProgressBar(len(val_set)) for idx in range(rank, len(val_set), world_size): print('idx', idx) if 'debug' in opt['name']: if (idx >= 3): break val_data = val_set[idx] val_data['LQs'].unsqueeze_(0) val_data['GT'].unsqueeze_(0) folder = val_data['folder'] idx_d, max_idx = val_data['idx'].split('/') idx_d, max_idx = int(idx_d), int(max_idx) if psnr_rlt.get(folder, None) is None: psnr_rlt[folder] = torch.zeros(max_idx, dtype=torch.float32, device='cuda') if ssim_rlt.get(folder, None) is None: ssim_rlt[folder] = torch.zeros(max_idx, dtype=torch.float32, device='cuda') if val_loss_rlt.get(folder, None) is None: val_loss_rlt[folder] = torch.zeros( max_idx, dtype=torch.float32, device='cuda') # tmp = torch.zeros(max_idx, dtype=torch.float32, device='cuda') model.feed_data(val_data) # model.test() # model.test_stitch() if opt['stitch'] == True: model.test_stitch() else: model.test() # large GPU memory # visuals = model.get_current_visuals() visuals = model.get_current_visuals( save=True, name='{}_{}'.format(folder, idx), save_path=opt['path']['val_images']) rlt_img = util.tensor2img(visuals['rlt']) # uint8 gt_img = util.tensor2img(visuals['GT']) # uint8 # calculate PSNR psnr = util.calculate_psnr(rlt_img, gt_img) psnr_rlt[folder][idx_d] = psnr # calculate SSIM ssim = util.calculate_ssim(rlt_img, gt_img) ssim_rlt[folder][idx_d] = ssim # calculate Val loss val_loss = model.get_loss() val_loss_rlt[folder][idx_d] = val_loss logger.info( '{}_{:02d} PSNR: {:.4f}, SSIM: {:.4f}'.format( folder, idx, psnr, ssim)) if rank == 0: for _ in range(world_size): pbar.update('Test {} - {}/{}'.format( folder, idx_d, max_idx)) # # collect data for _, v in psnr_rlt.items(): dist.reduce(v, 0) for _, v in ssim_rlt.items(): dist.reduce(v, 0) for _, v in val_loss_rlt.items(): dist.reduce(v, 0) dist.barrier() if rank == 0: psnr_rlt_avg = {} psnr_total_avg = 0. for k, v in psnr_rlt.items(): psnr_rlt_avg[k] = torch.mean(v).cpu().item() psnr_total_avg += psnr_rlt_avg[k] psnr_total_avg /= len(psnr_rlt) log_s = '# Validation # PSNR: {:.4e}:'.format( psnr_total_avg) for k, v in psnr_rlt_avg.items(): log_s += ' {}: {:.4e}'.format(k, v) logger.info(log_s) # ssim ssim_rlt_avg = {} ssim_total_avg = 0. for k, v in ssim_rlt.items(): ssim_rlt_avg[k] = torch.mean(v).cpu().item() ssim_total_avg += ssim_rlt_avg[k] ssim_total_avg /= len(ssim_rlt) log_s = '# Validation # PSNR: {:.4e}:'.format( ssim_total_avg) for k, v in ssim_rlt_avg.items(): log_s += ' {}: {:.4e}'.format(k, v) logger.info(log_s) # added val_loss_rlt_avg = {} val_loss_total_avg = 0. for k, v in val_loss_rlt.items(): val_loss_rlt_avg[k] = torch.mean(v).cpu().item() val_loss_total_avg += val_loss_rlt_avg[k] val_loss_total_avg /= len(val_loss_rlt) log_l = '# Validation # Loss: {:.4e}:'.format( val_loss_total_avg) for k, v in val_loss_rlt_avg.items(): log_l += ' {}: {:.4e}'.format(k, v) logger.info(log_l) message = '' for v in model.get_current_learning_rate(): message += '{:.5e}'.format(v) logger_val.info( 'Epoch {:02d}, LR {:s}, PSNR {:.4f}, SSIM {:.4f} Train {:s}, Val Total Loss {:.4e}' .format(epoch, message, psnr_total_avg, ssim_total_avg, message_train_loss, val_loss_total_avg)) if opt['use_tb_logger'] and 'debug' not in opt['name']: tb_logger.add_scalar('psnr_avg', psnr_total_avg, current_step) for k, v in psnr_rlt_avg.items(): tb_logger.add_scalar(k, v, current_step) # add val loss tb_logger.add_scalar('val_loss_avg', val_loss_total_avg, current_step) for k, v in val_loss_rlt_avg.items(): tb_logger.add_scalar(k, v, current_step) else: # Todo: our function One GPU pbar = util.ProgressBar(len(val_loader)) psnr_rlt = {} # with border and center frames psnr_rlt_avg = {} psnr_total_avg = 0. ssim_rlt = {} # with border and center frames ssim_rlt_avg = {} ssim_total_avg = 0. val_loss_rlt = {} val_loss_rlt_avg = {} val_loss_total_avg = 0. for val_inx, val_data in enumerate(val_loader): if 'debug' in opt['name']: if (val_inx >= 5): break folder = val_data['folder'][0] # idx_d = val_data['idx'].item() idx_d = val_data['idx'] # border = val_data['border'].item() if psnr_rlt.get(folder, None) is None: psnr_rlt[folder] = [] if ssim_rlt.get(folder, None) is None: ssim_rlt[folder] = [] if val_loss_rlt.get(folder, None) is None: val_loss_rlt[folder] = [] # process the black blank [B N C H W] print(val_data['LQs'].size()) H_S = val_data['LQs'].size(3) # 540 W_S = val_data['LQs'].size(4) # 960 print(H_S) print(W_S) blank_1_S = 0 blank_2_S = 0 print(val_data['LQs'][0, 2, 0, :, :].size()) for i in range(H_S): if not sum(val_data['LQs'][0, 2, 0, i, :]) == 0: blank_1_S = i - 1 # assert not sum(data_S[:, :, 0][i+1]) == 0 break for i in range(H_S): if not sum(val_data['LQs'][0, 2, 0, :, H_S - i - 1]) == 0: blank_2_S = (H_S - 1) - i - 1 # assert not sum(data_S[:, :, 0][blank_2_S-1]) == 0 break print('LQ :', blank_1_S, blank_2_S) if blank_1_S == -1: print('LQ has no blank') blank_1_S = 0 blank_2_S = H_S # val_data['LQs'] = val_data['LQs'][:,:,:,blank_1_S:blank_2_S,:] print("LQ", val_data['LQs'].size()) # end of process the black blank model.feed_data(val_data) if opt['stitch'] == True: model.test_stitch() else: model.test() # large GPU memory # process blank blank_1_L = blank_1_S << 2 blank_2_L = blank_2_S << 2 print(blank_1_L, blank_2_L) print(model.fake_H.size()) if not blank_1_S == 0: # model.fake_H = model.fake_H[:,:,blank_1_L:blank_2_L,:] model.fake_H[:, :, 0:blank_1_L, :] = 0 model.fake_H[:, :, blank_2_L:H_S, :] = 0 # end of # process blank visuals = model.get_current_visuals( save=True, name='{}_{:02d}'.format(folder, val_inx), save_path=opt['path']['val_images']) rlt_img = util.tensor2img(visuals['rlt']) # uint8 gt_img = util.tensor2img(visuals['GT']) # uint8 # calculate PSNR psnr = util.calculate_psnr(rlt_img, gt_img) psnr_rlt[folder].append(psnr) # calculate SSIM ssim = util.calculate_ssim(rlt_img, gt_img) ssim_rlt[folder].append(ssim) # val loss val_loss = model.get_loss() val_loss_rlt[folder].append(val_loss.item()) logger.info( '{}_{:02d} PSNR: {:.4f}, SSIM: {:.4f}'.format( folder, val_inx, psnr, ssim)) pbar.update('Test {} - {}'.format(folder, idx_d)) # average PSNR for k, v in psnr_rlt.items(): psnr_rlt_avg[k] = sum(v) / len(v) psnr_total_avg += psnr_rlt_avg[k] psnr_total_avg /= len(psnr_rlt) log_s = '# Validation # PSNR: {:.4e}:'.format( psnr_total_avg) for k, v in psnr_rlt_avg.items(): log_s += ' {}: {:.4e}'.format(k, v) logger.info(log_s) # average SSIM for k, v in ssim_rlt.items(): ssim_rlt_avg[k] = sum(v) / len(v) ssim_total_avg += ssim_rlt_avg[k] ssim_total_avg /= len(ssim_rlt) log_s = '# Validation # SSIM: {:.4e}:'.format( ssim_total_avg) for k, v in ssim_rlt_avg.items(): log_s += ' {}: {:.4e}'.format(k, v) logger.info(log_s) # average VMAF # average Val LOSS for k, v in val_loss_rlt.items(): val_loss_rlt_avg[k] = sum(v) / len(v) val_loss_total_avg += val_loss_rlt_avg[k] val_loss_total_avg /= len(val_loss_rlt) log_l = '# Validation # Loss: {:.4e}:'.format( val_loss_total_avg) for k, v in val_loss_rlt_avg.items(): log_l += ' {}: {:.4e}'.format(k, v) logger.info(log_l) # toal validation log message = '' for v in model.get_current_learning_rate(): message += '{:.5e}'.format(v) logger_val.info( 'Epoch {:02d}, LR {:s}, PSNR {:.4f}, SSIM {:.4f} Train {:s}, Val Total Loss {:.4e}' .format(epoch, message, psnr_total_avg, ssim_total_avg, message_train_loss, val_loss_total_avg)) # end add if opt['use_tb_logger'] and 'debug' not in opt['name']: tb_logger.add_scalar('psnr_avg', psnr_total_avg, current_step) for k, v in psnr_rlt_avg.items(): tb_logger.add_scalar(k, v, current_step) # tb_logger.add_scalar('ssim_avg', ssim_total_avg, current_step) # for k, v in ssim_rlt_avg.items(): # tb_logger.add_scalar(k, v, current_step) # add val loss tb_logger.add_scalar('val_loss_avg', val_loss_total_avg, current_step) for k, v in val_loss_rlt_avg.items(): tb_logger.add_scalar(k, v, current_step) ############################################ # # end of validation, save model # ############################################ # logger.info("Finished an epoch, Check and Save the model weights") # we check the validation loss instead of training loss. OK~ if saved_total_loss >= val_loss_total_avg: saved_total_loss = val_loss_total_avg #torch.save(model.state_dict(), args.save_path + "/best" + ".pth") model.save('best') logger.info( "Best Weights updated for decreased validation loss") else: logger.info( "Weights Not updated for undecreased validation loss") if saved_total_PSNR <= psnr_total_avg: saved_total_PSNR = psnr_total_avg model.save('bestPSNR') logger.info( "Best Weights updated for increased validation PSNR") else: logger.info( "Weights Not updated for unincreased validation PSNR") ############################################ # # end of one epoch, schedule LR # ############################################ # add scheduler todo if opt['train']['lr_scheme'] == 'ReduceLROnPlateau': for scheduler in model.schedulers: # scheduler.step(val_loss_total_avg) scheduler.step(val_loss_total_avg) if rank <= 0: logger.info('Saving the final model.') model.save('last') logger.info('End of training.') tb_logger.close()
def main(): # options parser = argparse.ArgumentParser() parser.add_argument('-opt', type=str, required=True, help='Path to options file.') opt = option.parse(parser.parse_args().opt, is_train=False) util.mkdirs((path for key, path in opt['path'].items() if not key == 'pretrain_model_G')) opt = option.dict_to_nonedict(opt) util.setup_logger(None, opt['path']['log'], 'test.log', level=logging.INFO, screen=True) logger = logging.getLogger('base') logger.info(option.dict2str(opt)) # Create test dataset and dataloader test_loaders = [] znorm = False for phase, dataset_opt in sorted(opt['datasets'].items()): test_set = create_dataset(dataset_opt) test_loader = create_dataloader(test_set, dataset_opt) logger.info('Number of test images in [{:s}]: {:d}'.format( dataset_opt['name'], len(test_set))) test_loaders.append(test_loader) # Temporary, will turn znorm on for all the datasets. Will need to introduce a variable for each dataset and differentiate each one later in the loop. if dataset_opt['znorm'] and znorm == False: znorm = True # Create model model = create_model(opt) for test_loader in test_loaders: test_set_name = test_loader.dataset.opt['name'] logger.info('\nTesting [{:s}]...'.format(test_set_name)) test_start_time = time.time() dataset_dir = os.path.join(opt['path']['results_root'], test_set_name) util.mkdir(dataset_dir) test_results = OrderedDict() test_results['psnr'] = [] test_results['ssim'] = [] test_results['psnr_y'] = [] test_results['ssim_y'] = [] for data in test_loader: need_HR = False if test_loader.dataset.opt[ 'dataroot_HR'] is None else True model.feed_data(data, need_HR=need_HR) img_path = data['LR_path'][0] img_name = os.path.splitext(os.path.basename(img_path))[0] model.test() # test visuals = model.get_current_visuals(need_HR=need_HR) if znorm: #opt['datasets']['train']['znorm']: # If the image range is [-1,1] # In testing, each "dataset" can have a different name (not train, val or other) sr_img = util.tensor2img(visuals['SR'], min_max=(-1, 1)) # uint8 else: # Default: Image range is [0,1] sr_img = util.tensor2img(visuals['SR']) # uint8 # save images suffix = opt['suffix'] if suffix: save_img_path = os.path.join(dataset_dir, img_name + suffix + '.png') else: save_img_path = os.path.join(dataset_dir, img_name + '.png') util.save_img(sr_img, save_img_path) # calculate PSNR and SSIM if need_HR: if znorm: #opt['datasets']['train']['znorm']: # If the image range is [-1,1] # In testing, each "dataset" can have a different name (not train, val or other) gt_img = util.tensor2img(visuals['HR'], min_max=(-1, 1)) # uint8 else: # Default: Image range is [0,1] gt_img = util.tensor2img(visuals['HR']) # uint8 gt_img = gt_img / 255. sr_img = sr_img / 255. crop_border = test_loader.dataset.opt['scale'] cropped_sr_img = sr_img[crop_border:-crop_border, crop_border:-crop_border, :] cropped_gt_img = gt_img[crop_border:-crop_border, crop_border:-crop_border, :] psnr = util.calculate_psnr(cropped_sr_img * 255, cropped_gt_img * 255) ssim = util.calculate_ssim(cropped_sr_img * 255, cropped_gt_img * 255) test_results['psnr'].append(psnr) test_results['ssim'].append(ssim) if gt_img.shape[2] == 3: # RGB image sr_img_y = bgr2ycbcr(sr_img, only_y=True) gt_img_y = bgr2ycbcr(gt_img, only_y=True) cropped_sr_img_y = sr_img_y[crop_border:-crop_border, crop_border:-crop_border] cropped_gt_img_y = gt_img_y[crop_border:-crop_border, crop_border:-crop_border] psnr_y = util.calculate_psnr(cropped_sr_img_y * 255, cropped_gt_img_y * 255) ssim_y = util.calculate_ssim(cropped_sr_img_y * 255, cropped_gt_img_y * 255) test_results['psnr_y'].append(psnr_y) test_results['ssim_y'].append(ssim_y) logger.info('{:20s} - PSNR: {:.6f} dB; SSIM: {:.6f}; PSNR_Y: {:.6f} dB; SSIM_Y: {:.6f}.'\ .format(img_name, psnr, ssim, psnr_y, ssim_y)) else: logger.info( '{:20s} - PSNR: {:.6f} dB; SSIM: {:.6f}.'.format( img_name, psnr, ssim)) else: logger.info(img_name) if need_HR: # metrics # Average PSNR/SSIM results ave_psnr = sum(test_results['psnr']) / len(test_results['psnr']) ave_ssim = sum(test_results['ssim']) / len(test_results['ssim']) logger.info('----Average PSNR/SSIM results for {}----\n\tPSNR: {:.6f} dB; SSIM: {:.6f}\n'\ .format(test_set_name, ave_psnr, ave_ssim)) if test_results['psnr_y'] and test_results['ssim_y']: ave_psnr_y = sum(test_results['psnr_y']) / len( test_results['psnr_y']) ave_ssim_y = sum(test_results['ssim_y']) / len( test_results['ssim_y']) logger.info('----Y channel, average PSNR/SSIM----\n\tPSNR_Y: {:.6f} dB; SSIM_Y: {:.6f}\n'\ .format(ave_psnr_y, ave_ssim_y))
def main(): ################# # configurations ################# flip_test = False scale = 4 N_in = 7 predeblur, HR_in = False, False n_feats = 128 back_RBs = 40 save_imgs = False prog = argparse.ArgumentParser() prog.add_argument('--train_mode', '-t', type=str, default='REDS', help='train mode') prog.add_argument('--data_mode', '-m', type=str, default=None, help='data_mode') prog.add_argument('--degradation_mode', '-d', type=str, default='impulse', choices=('impulse', 'bicubic', 'preset'), help='path to image output directory.') prog.add_argument('--sigma_x', '-sx', type=float, default=1, help='sigma_x') prog.add_argument('--sigma_y', '-sy', type=float, default=0, help='sigma_y') prog.add_argument('--theta', '-th', type=float, default=0, help='theta') prog.add_argument('--model', type=str, default=None, help='name for subdirectory') args = prog.parse_args() train_data_mode = args.train_mode data_mode = args.data_mode if data_mode is None: if train_data_mode == 'Vimeo': data_mode = 'Vid4' elif train_data_mode == 'REDS': data_mode = 'REDS' degradation_mode = args.degradation_mode # impulse | bicubic | preset sig_x, sig_y, the = args.sigma_x, args.sigma_y, args.theta if sig_y == 0: sig_y = sig_x folder_subname = 'preset' if degradation_mode == 'preset' else degradation_mode + '_' + str( '{:.1f}'.format(sig_x)) + '_' + str('{:.1f}'.format(sig_y)) + '_' + str('{:.1f}'.format(the)) #### dataset if data_mode == 'Vid4': test_dataset_folder = '../dataset/Vid4/LR_{}/X1_KG_ZSSR'.format(folder_subname) #test_dataset_folder = '../dataset/Vid4/LR_{}/X1_CF_DBPN'.format(folder_subname) GT_dataset_folder = '../dataset/Vid4/HR' elif data_mode == 'MM522': test_dataset_folder = '../dataset/MM522val/LR_bicubic/X1_KG_ZSSR' GT_dataset_folder = '../dataset/MM522val/HR' else: # test_dataset_folder = '../dataset/REDS4/LR_bicubic/X{}'.format(scale) test_dataset_folder = '../dataset/REDS/train/LR_{}/X1_KG_ZSSR'.format(folder_subname) #test_dataset_folder = '../dataset/REDS/train/LR_{}/X1_CF_DBPN'.format(folder_subname) GT_dataset_folder = '../dataset/REDS/train/HR' #### evaluation crop_border = 0 border_frame = 0 # N_in // 2 # border frames when evaluate # temporal padding mode padding = 'new_info' save_folder = '../results/{}'.format(data_mode) util.mkdirs(save_folder) util.setup_logger('base', save_folder, 'test', level=logging.INFO, screen=True, tofile=True) logger = logging.getLogger('base') #### log info logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder)) logger.info('Padding mode: {}'.format(padding)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Flip test: {}'.format(flip_test)) avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l = [], [], [] avg_ssim_l, avg_ssim_center_l, avg_ssim_border_l = [], [], [] subfolder_name_l = [] subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*'))) subfolder_GT_l = sorted(glob.glob(osp.join(GT_dataset_folder, '*'))) if data_mode == 'REDS': subfolder_GT_l = [k for k in subfolder_GT_l if k.find('000') >= 0 or k.find('011') >= 0 or k.find('015') >= 0 or k.find('020') >= 0] # for each subfolder for subfolder, subfolder_GT in zip(subfolder_l, subfolder_GT_l): subfolder_name = osp.basename(subfolder) subfolder_name_l.append(subfolder_name) save_subfolder = osp.join(save_folder, subfolder_name) img_path_l = sorted(glob.glob(osp.join(subfolder_GT, '*5.png'))) ## *5.png max_idx = len(img_path_l) if save_imgs: util.mkdirs(save_subfolder) #### read LQ and GT images img_LQ_l = [] img_GT_l = [] for img_LQ_path in sorted(glob.glob(osp.join(subfolder, '*5.png.png'))): img_LQ_l.append(data_util.read_img(None, img_LQ_path)) for img_GT_path in sorted(glob.glob(osp.join(subfolder_GT, '*5.png'))): ### *5.png img_GT_l.append(data_util.read_img(None, img_GT_path)) avg_psnr, avg_psnr_border, avg_psnr_center, N_border, N_center = 0, 0, 0, 0, 0 avg_ssim, avg_ssim_border, avg_ssim_center = 0, 0, 0 # process each image for img_idx, img_path in enumerate(img_path_l): img_name = osp.splitext(osp.basename(img_path))[0] if save_imgs: cv2.imwrite(osp.join(save_subfolder, '{}.png'.format(img_name)), output) # calculate PSNR output = np.copy(img_LQ_l[img_idx]) GT = np.copy(img_GT_l[img_idx]) ''' output_tensor = torch.from_numpy(np.copy(output[:,:,::-1])).permute(2,0,1) GT_tensor = torch.from_numpy(np.copy(GT[:,:,::-1])).permute(2,0,1).type_as(output_tensor) torch.save(output_tensor.cpu(), '../results/sr_test.pt') torch.save(GT_tensor.cpu(), '../results/hr_test.pt') my_psnr = utility.calc_psnr(output_tensor, GT_tensor) GT_tensor = GT_tensor.cpu().numpy().transpose(1,2,0) imageio.imwrite('../results/hr_test.png', GT_tensor) print('saved', my_psnr) ''' ''' # For REDS, evaluate on RGB channels; for Vid4, evaluate on the Y channel if data_mode == 'Vid4' or 'sharp_bicubic' or 'MM522': # bgr2y, [0, 1] GT = data_util.bgr2ycbcr(GT, only_y=True) output = data_util.bgr2ycbcr(output, only_y=True) ''' output = (output * 255).round().astype('uint8') GT = (GT * 255).round().astype('uint8') output, GT = util.crop_border([output, GT], crop_border) crt_psnr = util.calculate_psnr(output, GT) crt_ssim = util.calculate_ssim(output, GT) # logger.info('{:3d} - {:16} \tPSNR: {:.6f} dB \tSSIM: {:.6f}'.format(img_idx + 1, img_name, crt_psnr, crt_ssim)) if img_idx >= border_frame and img_idx < max_idx - border_frame: # center frames avg_psnr_center += crt_psnr avg_ssim_center += crt_ssim N_center += 1 else: # border frames avg_psnr_border += crt_psnr avg_ssim_border += crt_ssim N_border += 1 avg_psnr = (avg_psnr_center + avg_psnr_border) / (N_center + N_border) avg_psnr_center = avg_psnr_center / N_center avg_psnr_border = 0 if N_border == 0 else avg_psnr_border / N_border avg_psnr_l.append(avg_psnr) avg_psnr_center_l.append(avg_psnr_center) avg_psnr_border_l.append(avg_psnr_border) avg_ssim = (avg_ssim_center + avg_ssim_border) / (N_center + N_border) avg_ssim_center = avg_ssim_center / N_center avg_ssim_border = 0 if N_border == 0 else avg_ssim_border / N_border avg_ssim_l.append(avg_ssim) avg_ssim_center_l.append(avg_ssim_center) avg_ssim_border_l.append(avg_ssim_border) logger.info('Folder {} - Average PSNR: {:.6f} dB for {} frames; ' 'Center PSNR: {:.6f} dB for {} frames; ' 'Border PSNR: {:.6f} dB for {} frames.'.format(subfolder_name, avg_psnr, (N_center + N_border), avg_psnr_center, N_center, avg_psnr_border, N_border)) logger.info('Folder {} - Average SSIM: {:.6f} for {} frames; ' 'Center SSIM: {:.6f} for {} frames; ' 'Border SSIM: {:.6f} for {} frames.'.format(subfolder_name, avg_ssim, (N_center + N_border), avg_ssim_center, N_center, avg_ssim_border, N_border)) ''' logger.info('################ Tidy Outputs ################') for subfolder_name, psnr, psnr_center, psnr_border in zip(subfolder_name_l, avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l): logger.info('Folder {} - Average PSNR: {:.6f} dB. ' 'Center PSNR: {:.6f} dB. ' 'Border PSNR: {:.6f} dB.'.format(subfolder_name, psnr, psnr_center, psnr_border)) for subfolder_name, ssim, ssim_center, ssim_border in zip(subfolder_name_l, avg_ssim_l, avg_ssim_center_l, avg_ssim_border_l): logger.info('Folder {} - Average SSIM: {:.6f}. ' 'Center SSIM: {:.6f}. ' 'Border SSIM: {:.6f}.'.format(subfolder_name, ssim, ssim_center, ssim_border)) ''' logger.info('################ Final Results ################') logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder)) logger.info('Padding mode: {}'.format(padding)) logger.info('Save images: {}'.format(save_imgs)) logger.info('Flip test: {}'.format(flip_test)) logger.info('Total Average PSNR: {:.6f} dB for {} clips. ' 'Center PSNR: {:.6f} dB. Border PSNR: {:.6f} dB.'.format( sum(avg_psnr_l) / len(avg_psnr_l), len(subfolder_l), sum(avg_psnr_center_l) / len(avg_psnr_center_l), sum(avg_psnr_border_l) / len(avg_psnr_border_l))) logger.info('Total Average SSIM: {:.6f} for {} clips. ' 'Center SSIM: {:.6f}. Border PSNR: {:.6f}.'.format( sum(avg_ssim_l) / len(avg_ssim_l), len(subfolder_l), sum(avg_ssim_center_l) / len(avg_ssim_center_l), sum(avg_ssim_border_l) / len(avg_ssim_border_l))) print('\n\n\n')