def main(): parser = argparse.ArgumentParser() parser.add_argument('-opt', type=str, 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) opt = option.dict_to_nonedict(opt) dataset_opt = opt['datasets']['test'] test_set = create_dataset(dataset_opt) test_loader = create_dataloader(test_set, dataset_opt, opt, None) model = Model(opt) if test_loader is not None: calc_lpips = PerceptualLossLPIPS() if True: avg_ssim = avg_psnr = avg_lpips = 0 print("Testing Starts!") idx = 0 test_bar = tqdm(test_loader) for test_data in test_bar: idx += 1 img_name = os.path.splitext(os.path.basename(test_data['LQ_path'][0]))[0] img_dir = '../test_results_' + opt['name'] util.mkdir(img_dir) model.feed_data(test_data) model.test() visuals = model.get_current_visuals() sr_img = util.tensor2img(visuals['SR']) # uint8 gt_img = util.tensor2img(visuals['GT']) # uint8 lq_img = util.tensor2img(visuals['LQ']) # uint8 # Save SR images for reference save_sr_img_path = os.path.join(img_dir, '{:s}_sr.png'.format(img_name)) util.save_img(sr_img, save_sr_img_path) gt_img = gt_img / 255. sr_img = sr_img / 255. lq_img = lq_img / 255. avg_psnr += util.calculate_psnr(sr_img * 255, gt_img * 255) avg_ssim += util.calculate_ssim(sr_img * 255, gt_img * 255) avg_lpips += calc_lpips(visuals['SR'], visuals['GT']) avg_psnr = avg_psnr / idx avg_ssim = avg_ssim / idx avg_lpips = avg_lpips / idx print('Test_Result_{:s} psnr: {:.4e} ssim: {:.4e} lpips: {:.4e}'.format( opt['name'], avg_psnr, avg_ssim, avg_lpips))
def main(): from options import NoneDict, dict_to_nonedict opt = {} opt['network_G'] = 'ESRGAN' opt['network_D'] = 'patchgan' opt['datasets'] = {} opt['datasets']['train'] = {} opt['datasets']['train']['crop_size'] = 128 opt = dict_to_nonedict(opt) opt = get_network_defaults(opt) print(opt)
import os from PIL import Image os.environ['KMP_DUPLICATE_LIB_OK'] = 'True' #### options parser = argparse.ArgumentParser() parser.add_argument("-opt", type=str, default="test_setting2.yml", help="Path to options YMAL file.") parser.add_argument("-input_dir", type=str, default="../../../data_samples/") parser.add_argument("-output_dir", type=str, default="../../../data_samples/") args = parser.parse_args() opt = option.parse(args.opt, is_train=False) opt = option.dict_to_nonedict(opt) model = create_model(opt) if not osp.exists(args.output_dir): os.makedirs(args.output_dir) def predict(img): # img = cv2.imread(args.input_dir)[:, :, [2, 1, 0]] img = img.transpose(2, 0, 1)[None] / 255 img_t = torch.as_tensor(np.ascontiguousarray(img)).float() model.feed_data(img_t) model.test() sr = model.fake_SR.detach().float().cpu()[0] sr_im = tensor2img(sr) sr_im = sr_im[:, :, ::-1]
def main(): # options parser = argparse.ArgumentParser() parser.add_argument('-opt', type=str, required=True, help='Path to options file.') opt = options.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 = options.dict_to_nonedict(opt) util.setup_logger(None, opt['path']['log'], 'test.log', level=logging.INFO, screen=True) logger = logging.getLogger('base') logger.info(options.dict2str(opt)) # Create test dataset and dataloader test_loaders = [] znorm = False # TMP 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['in_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 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) top_img = tensor2np(visuals['top_fake']) # uint8 bot_img = tensor2np(visuals['bottom_fake']) # uint8 # save images suffix = opt['suffix'] if suffix: save_img_path = os.path.join( dataset_dir, img_name + suffix) else: save_img_path = os.path.join(dataset_dir, img_name) util.save_img(top_img, save_img_path + '_top.png') util.save_img(bot_img, save_img_path + '_bot.png') #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(): #### setup options of three networks parser = argparse.ArgumentParser() parser.add_argument("-opt", type=str, help="Path to option YMAL file of Predictor.") 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) # convert to NoneDict, which returns None for missing keys opt = option.dict_to_nonedict(opt) # choose small opt for SFTMD test, fill path of pre-trained model_F #### set random seed seed = opt["train"]["manual_seed"] if seed is None: seed = random.randint(1, 10000) util.set_random_seed(seed) # load PCA matrix of enough kernel print("load PCA matrix") pca_matrix = torch.load(opt["pca_matrix_path"], map_location=lambda storage, loc: storage) print("PCA matrix shape: {}".format(pca_matrix.shape)) #### distributed training settings if args.launcher == "none": # disabled distributed training opt["dist"] = False opt["dist"] = False rank = -1 print("Disabled distributed training.") else: opt["dist"] = True opt["dist"] = True init_dist() world_size = ( torch.distributed.get_world_size() ) # Returns the number of processes in the current process group rank = torch.distributed.get_rank( ) # Returns the rank of current process group torch.backends.cudnn.benchmark = True # torch.backends.cudnn.deterministic = True ###### Predictor&Corrector train ###### #### 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-7) if resume_state is None: # Predictor path 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)) os.system("rm ./log") os.symlink(os.path.join(opt["path"]["experiments_root"], ".."), "./log") # 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( "val", opt["path"]["log"], "val_" + 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="log/tb_logger/" + opt["name"]) else: util.setup_logger("base", opt["path"]["log"], "train", level=logging.INFO, screen=True) logger = logging.getLogger("base") 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 assert val_loader is not None #### create model model = create_model(opt) # load pretrained model of SFTMD #### 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 prepro = util.SRMDPreprocessing( opt["scale"], pca_matrix, random=True, para_input=opt["code_length"], kernel=opt["kernel_size"], noise=False, cuda=True, sig=None, sig_min=opt["sig_min"], sig_max=opt["sig_max"], rate_iso=1.0, scaling=3, rate_cln=0.2, noise_high=0.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 #### preprocessing for LR_img and kernel map LR_img, ker_map = prepro(train_data["GT"]) LR_img = (LR_img * 255).round() / 255 #### training Predictor model.feed_data(LR_img, train_data["GT"], ker_map) model.optimize_parameters(current_step) model.update_learning_rate(current_step, warmup_iter=opt["train"]["warmup_iter"]) visuals = model.get_current_visuals() #### log of model_P if current_step % opt["logger"]["print_freq"] == 0: logs = model.get_current_log() message = "Predictor <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) # validation, to produce ker_map_list(fake) if current_step % opt["train"]["val_freq"] == 0 and rank <= 0: avg_psnr = 0.0 idx = 0 for _, val_data in enumerate(val_loader): # LR_img, ker_map = prepro(val_data['GT']) LR_img = val_data["LQ"] lr_img = util.tensor2img( LR_img) # save LR image for reference # valid Predictor model.feed_data(LR_img, val_data["GT"]) model.test() visuals = model.get_current_visuals() # Save images for reference 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) # img_dir = os.path.join(opt['path']['val_images'], str(current_step), '_', str(step)) util.mkdir(img_dir) save_lr_path = os.path.join(img_dir, "{:s}_LR.png".format(img_name)) util.save_img(lr_img, save_lr_path) sr_img = util.tensor2img(visuals["SR"]) # uint8 gt_img = util.tensor2img(visuals["GT"]) # uint8 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 crop_size = opt["scale"] gt_img = gt_img / 255.0 sr_img = sr_img / 255.0 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) idx += 1 avg_psnr = avg_psnr / idx # log logger.info("# Validation # PSNR: {:.6f}".format(avg_psnr)) logger_val = logging.getLogger("val") # validation logger logger_val.info( "<epoch:{:3d}, iter:{:8,d}, psnr: {:.6f}".format( epoch, current_step, avg_psnr)) # tensorboard logger if opt["use_tb_logger"] and "debug" not in opt["name"]: tb_logger.add_scalar("psnr", avg_psnr, 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 Predictor and Corrector training.") tb_logger.close()
"-opt_P", type=str, default="options/test/Predictor/test_Predictor_x4.yml", help="Path to options YMAL file.", ) parser.add_argument( "-opt_C", type=str, default="options/test/Corrector/test_Corrector_x4.yml", help="Path to options YMAL file.", ) opt_F = option.parse(parser.parse_args().opt_F, is_train=False) opt_P = option.parse(parser.parse_args().opt_P, is_train=False) opt_C = option.parse(parser.parse_args().opt_C, is_train=False) opt_F = option.dict_to_nonedict(opt_F) opt_P = option.dict_to_nonedict(opt_P) opt_C = option.dict_to_nonedict(opt_C) #### mkdir and logger util.mkdirs((path for key, path in opt_P["path"].items() if not key == "experiments_root" and "pretrain_model" not in key and "resume" not in key)) util.mkdirs((path for key, path in opt_C["path"].items() if not key == "experiments_root" and "pretrain_model" not in key and "resume" not in key)) os.system("rm ./result") os.symlink(os.path.join(opt_P["path"]["results_root"], ".."), "./result") util.setup_logger(
def main(): #### options parser = argparse.ArgumentParser() parser.add_argument('-opt', type=str, 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) #### 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: print(opt['path']) 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 and path is not None)) # 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('val', opt['path']['log'], 'val_' + 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 trial = 0 while os.path.isdir('../Loggers/' + opt['name'] + '/' + str(trial)): trial += 1 tb_logger = SummaryWriter(log_dir='../Loggers/' + opt['name'] + '/' + str(trial)) 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) # -------------------------------------------- ADDED -------------------------------------------- l1_loss = torch.nn.L1Loss() mse_loss = torch.nn.MSELoss() calc_lpips = PerceptualLossLPIPS() if torch.cuda.is_available(): l1_loss = l1_loss.cuda() mse_loss = mse_loss.cuda() # ----------------------------------------------------------------------------------------------- #### 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.benckmark = 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 = 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) train_bar = tqdm(train_loader, desc='[%d/%d]' % (epoch, total_epochs)) for bus, train_data in enumerate(train_bar): # validation if epoch % opt['train']['val_freq'] == 0 and bus == 0 and rank <= 0: avg_ssim = avg_psnr = avg_lpips = val_pix_err_f = val_pix_err_nf = val_mean_color_err = 0.0 print("into validation!") idx = 0 val_bar = tqdm(val_loader, desc='[%d/%d]' % (epoch, total_epochs)) for val_data in val_bar: 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 lq_img = util.tensor2img(visuals['LQ']) # uint8 #nr_img = util.tensor2img(visuals['NR']) # uint8 #nf_img = util.tensor2img(visuals['NF']) # uint8 #nh_img = util.tensor2img(visuals['NH']) # uint8 #print("Great! images got into here.") # Save SR images for reference save_sr_img_path = os.path.join(img_dir, '{:s}_{:d}_sr.png'.format(img_name, current_step)) save_nr_img_path = os.path.join(img_dir, '{:s}_{:d}_lq.png'.format(img_name, current_step)) #save_nf_img_path = os.path.join(img_dir, # 'bs_{:s}_{:d}_nr.png'.format(img_name, current_step)) #save_nh_img_path = os.path.join(img_dir, # 'bs_{:s}_{:d}_nh.png'.format(img_name, current_step)) util.save_img(sr_img, save_sr_img_path) util.save_img(lq_img, save_nr_img_path) #util.save_img(nf_img, save_nf_img_path) #util.save_img(nh_img, save_nh_img_path) #print("Saved") # calculate PSNR gt_img = gt_img / 255. sr_img = sr_img / 255. #nf_img = nf_img / 255. lq_img = lq_img / 255. #cropped_lq_img = lq_img[crop_size:-crop_size, crop_size:-crop_size, :] #cropped_nr_img = nr_img[crop_size:-crop_size, crop_size:-crop_size, :] avg_psnr += util.calculate_psnr(sr_img * 255, gt_img * 255) avg_ssim += util.calculate_ssim(sr_img * 255, gt_img * 255) avg_lpips += calc_lpips(visuals['SR'], visuals['GT']) #avg_psnr_n += util.calculate_psnr(cropped_lq_img * 255, cropped_nr_img * 255) # ----------------------------------------- ADDED ----------------------------------------- val_pix_err_nf += l1_loss(visuals['SR'], visuals['GT']) val_mean_color_err += mse_loss(visuals['SR'].mean(2).mean(1), visuals['GT'].mean(2).mean(1)) # ----------------------------------------------------------------------------------------- avg_psnr = avg_psnr / idx avg_ssim = avg_ssim / idx avg_lpips = avg_lpips / idx val_pix_err_f /= idx val_pix_err_nf /= idx val_mean_color_err /= idx # log logger.info('# Validation # PSNR: {:.4e},'.format(avg_psnr)) logger.info('# Validation # SSIM: {:.4e},'.format(avg_ssim)) logger.info('# Validation # LPIPS: {:.4e},'.format(avg_lpips)) logger_val = logging.getLogger('val') # validation logger logger_val.info('<epoch:{:3d}, iter:{:8,d}> psnr: {:.4e} ssim: {:.4e} lpips: {:.4e}'.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('val_psnr', avg_psnr, current_step) tb_logger.add_scalar('val_ssim', avg_ssim, current_step) tb_logger.add_scalar('val_lpips', avg_lpips, 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) 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) model.clear_data() #### tb_logger if current_step % opt['logger']['tb_freq'] == 0: logs = model.get_current_log() if opt['use_tb_logger'] and 'debug' not in opt['name']: for k, v in logs.items(): if rank <= 0: tb_logger.add_scalar(k, v, current_step) #### logger if epoch % opt['logger']['print_freq'] == 0 and epoch != 0 and bus == 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 rank <= 0: logger.info(message) #### save models and training states if epoch % opt['logger']['save_checkpoint_freq'] == 0 and epoch != 0 and bus == 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.')