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.')
