def tmp_vis(img_t, to_np=True, rgb2bgr=True, remove_batch=False, save_dir=''):
'''
Visualization function that can be inserted at any point
in the code, works with tensor or np images
img_t: image
save_dir: path to save image
'''
import cv2
from dataops.common import tensor2np
if to_np:
img = tensor2np(img_t.detach(),
rgb2bgr=rgb2bgr,
remove_batch=remove_batch)
else:
img = img_t
print("out: ", img.shape)
cv2.imshow('image', img)
cv2.waitKey(0)
if save_dir != '':
cv2.imwrite(save_dir, img)
cv2.destroyAllWindows()
return None
def save_image(image=None, num_rep=0, sufix=None, random=False):
'''
Output images to a directory instead of visualizing them,
may be easier to compare multiple batches of images
'''
import uuid, cv2
from dataops.common import tensor2np
img = tensor2np(image, remove_batch=False) # uint8
if random:
#random name to save + had to multiply by 255, else getting all black image
hex = uuid.uuid4().hex
cv2.imwrite(
"D:/tmp_test/fake_" + sufix + "_" + str(num_rep) + hex + ".png",
img)
else:
cv2.imwrite("D:/tmp_test/fake_" + sufix + "_" + str(num_rep) + ".png",
img)
return None
def get_sp_transform(train_opt: dict, znorm: bool = True):
n_segments = train_opt.get('sp_n_segments', 200) # 500
max_size = train_opt.get('sp_max_size', None) # crop_size
# 'selective' 'cluster' 'rag' None
reduction = train_opt.get('sp_reduction', 'selective')
# 'seeds', 'slic', 'slico', 'mslic', 'sk_slic', 'sk_felzenszwalb'
algo = train_opt.get('sp_algo', 'sk_felzenszwalb')
gamma_range = train_opt.get('sp_gamma_range', (100, 120))
superpixel_fn = transforms.Compose([
transforms.Lambda(lambda img: tensor2np(
img, rgb2bgr=True, denormalize=znorm, remove_batch=False)),
transforms.Superpixels(p_replace=1,
n_segments=n_segments,
algo=algo,
reduction=reduction,
max_size=max_size,
p=1),
transforms.RandomGamma(gamma_range=gamma_range, gain=1, p=1),
transforms.Lambda(lambda img: np2tensor(
img, bgr2rgb=True, normalize=znorm, add_batch=False))
])
return superpixel_fn
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 test_loop(model, opt, dataloaders, data_params):
logger = util.get_root_logger()
# read data_params
znorms = data_params['znorm']
# prepare the metric calculation classes for RGB and Y_only images
calc_metrics = opt.get('metrics', None)
if calc_metrics:
test_metrics = metrics.MetricsDict(metrics = calc_metrics)
test_metrics_y = metrics.MetricsDict(metrics = calc_metrics)
for phase, dataloader in dataloaders.items():
name = dataloader.dataset.opt['name']
logger.info('\nTesting [{:s}]...'.format(name))
dataset_dir = os.path.join(opt['path']['results_root'], name)
util.mkdir(dataset_dir)
for data in dataloader:
znorm = znorms[name]
need_HR = False if dataloader.dataset.opt['dataroot_HR'] is None else True
# set up per image CEM wrapper if configured
CEM_net = get_CEM(opt, data)
model.feed_data(data, need_HR=need_HR) # unpack data from data loader
# img_path = data['LR_path'][0]
img_path = get_img_path(data)
img_name = os.path.splitext(os.path.basename(img_path))[0]
# test with eval mode. This only affects layers like batchnorm and dropout.
test_mode = opt.get('test_mode', None)
if test_mode == 'x8':
# geometric self-ensemble
model.test_x8(CEM_net=CEM_net)
elif test_mode == 'chop':
# chop images in patches/crops, to reduce VRAM usage
model.test_chop(patch_size=opt.get('chop_patch_size', 100),
step=opt.get('chop_step', 0.9),
CEM_net=CEM_net)
else:
# normal inference
model.test(CEM_net=CEM_net) # run inference
# get image results
visuals = model.get_current_visuals(need_HR=need_HR)
# post-process options if using CEM
if opt.get('use_cem', None) and opt['cem_config'].get('out_orig', False):
# run regular inference
if test_mode == 'x8':
model.test_x8()
elif test_mode == 'chop':
model.test_chop(patch_size=opt.get('chop_patch_size', 100),
step=opt.get('chop_step', 0.9))
else:
model.test()
orig_visuals = model.get_current_visuals(need_HR=need_HR)
if opt['cem_config'].get('out_filter', False):
GF = GuidedFilter(ks=opt['cem_config'].get('out_filter_ks', 7))
filt = GF(visuals['SR'].unsqueeze(0), (visuals['SR']-orig_visuals['SR']).unsqueeze(0)).squeeze(0)
visuals['SR'] = orig_visuals['SR']+filt
if opt['cem_config'].get('out_keepY', False):
out_regY = rgb_to_ycbcr(orig_visuals['SR']).unsqueeze(0)
out_cemY = rgb_to_ycbcr(visuals['SR']).unsqueeze(0)
visuals['SR'] = ycbcr_to_rgb(torch.cat([out_regY[:, 0:1, :, :], out_cemY[:, 1:2, :, :], out_cemY[:, 2:3, :, :]], 1)).squeeze(0)
res_options = visuals_check(visuals.keys(), opt.get('val_comparison', None))
# save images
save_img_path = os.path.join(dataset_dir, img_name + opt.get('suffix', ''))
# save single images or lr / sr comparison
if opt['val_comparison'] and len(res_options['save_imgs']) > 1:
comp_images = [tensor2np(visuals[save_img_name], denormalize=znorm) for save_img_name in res_options['save_imgs']]
util.save_img_comp(comp_images, save_img_path + '.png')
else:
for save_img_name in res_options['save_imgs']:
imn = '_' + save_img_name if len(res_options['save_imgs']) > 1 else ''
util.save_img(tensor2np(visuals[save_img_name], denormalize=znorm), save_img_path + imn + '.png')
# calculate metrics if HR dataset is provided and metrics are configured in options
if need_HR and calc_metrics and res_options['aligned_metrics']:
metric_imgs = [tensor2np(visuals[x], denormalize=znorm) for x in res_options['compare_imgs']]
test_results = test_metrics.calculate_metrics(metric_imgs[0], metric_imgs[1],
crop_size=opt['scale'])
# prepare single image metrics log message
logger_m = '{:20s} -'.format(img_name)
for k, v in test_results:
formatted_res = k.upper() + ': {:.6f}, '.format(v)
logger_m += formatted_res
if gt_img.shape[2] == 3: # RGB image, calculate y_only metrics
test_results_y = test_metrics_y.calculate_metrics(metric_imgs[0], metric_imgs[1],
crop_size=opt['scale'], only_y=True)
# add the y only results to the single image log message
for k, v in test_results_y:
formatted_res = k.upper() + ': {:.6f}, '.format(v)
logger_m += formatted_res
logger.info(logger_m)
else:
logger.info(img_name)
# average metrics results for the dataset
if need_HR and calc_metrics:
# aggregate the metrics results (automatically resets the metric classes)
avg_metrics = test_metrics.get_averages()
avg_metrics_y = test_metrics_y.get_averages()
# prepare log average metrics message
agg_logger_m = ''
for r in avg_metrics:
formatted_res = r['name'].upper() + ': {:.6f}, '.format(r['average'])
agg_logger_m += formatted_res
logger.info('----Average metrics results for {}----\n\t'.format(name) + agg_logger_m[:-2])
if len(avg_metrics_y > 0):
# prepare log average Y channel metrics message
agg_logger_m = ''
for r in avg_metrics_y:
formatted_res = r['name'].upper() + '_Y' + ': {:.6f}, '.format(r['average'])
agg_logger_m += formatted_res
logger.info('----Y channel, average metrics ----\n\t' + agg_logger_m[:-2])
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 fit(model, opt, dataloaders, steps_states, data_params, loggers):
# read data_params
batch_size = data_params['batch_size']
virtual_batch_size = data_params['virtual_batch_size']
total_iters = data_params['total_iters']
total_epochs = data_params['total_epochs']
# read steps_states
start_epoch = steps_states["start_epoch"]
current_step = steps_states["current_step"]
virtual_step = steps_states["virtual_step"]
# read loggers
logger = util.get_root_logger()
tb_logger = loggers["tb_logger"]
# training
logger.info('Start training from epoch: {:d}, iter: {:d}'.format(start_epoch, current_step))
try:
timer = metrics.Timer() # iteration timer
timerData = metrics.TickTock() # data timer
timerEpoch = metrics.TickTock() # epoch timer
# outer loop for different epochs
for epoch in range(start_epoch, (total_epochs * (virtual_batch_size // batch_size))+1):
timerData.tick()
timerEpoch.tick()
# inner iteration loop within one epoch
for n, train_data in enumerate(dataloaders['train'], start=1):
timerData.tock()
virtual_step += 1
take_step = False
if virtual_step > 0 and virtual_step * batch_size % virtual_batch_size == 0:
current_step += 1
take_step = True
if current_step > total_iters:
break
# training
model.feed_data(train_data) # unpack data from dataset and apply preprocessing
model.optimize_parameters(virtual_step) # calculate loss functions, get gradients, update network weights
# log
def eta(t_iter):
# calculate training ETA in hours
return (t_iter * (opt['train']['niter'] - current_step)) / 3600 if t_iter > 0 else 0
if current_step % opt['logger']['print_freq'] == 0 and take_step:
# iteration end time
avg_time = timer.get_average_and_reset()
avg_data_time = timerData.get_average_and_reset()
# print training losses and save logging information to disk
message = '<epoch:{:3d}, iter:{:8,d}, lr:{:.3e}, t:{:.4f}s, td:{:.4f}s, eta:{:.4f}h> '.format(
epoch, current_step, model.get_current_learning_rate(current_step), avg_time,
avg_data_time, eta(avg_time))
# tensorboard training logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
if current_step % opt['logger'].get('tb_sample_rate', 1) == 0: # Reduce rate of tb logs
# tb_logger.add_scalar('loss/nll', nll, current_step)
tb_logger.add_scalar('lr/base', model.get_current_learning_rate(), current_step)
tb_logger.add_scalar('time/iteration', timer.get_last_iteration(), current_step)
tb_logger.add_scalar('time/data', timerData.get_last_iteration(), current_step)
logs = model.get_current_log()
for k, v in logs.items():
message += '{:s}: {:.4e} '.format(k, v)
# tensorboard loss logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
if current_step % opt['logger'].get('tb_sample_rate', 1) == 0: # Reduce rate of tb logs
tb_logger.add_scalar(k, v, current_step)
# tb_logger.flush()
logger.info(message)
# start time for next iteration #TODO:skip the validation time from calculation
timer.tick()
# update learning rate
if model.optGstep and model.optDstep and take_step:
model.update_learning_rate(current_step, warmup_iter=opt['train'].get('warmup_iter', -1))
# save latest models and training states every <save_checkpoint_freq> iterations
if current_step % opt['logger']['save_checkpoint_freq'] == 0 and take_step:
if model.swa:
model.save(current_step, opt['logger']['overwrite_chkp'], loader=dataloaders['train'])
else:
model.save(current_step, opt['logger']['overwrite_chkp'])
model.save_training_state(
epoch=epoch + (n >= len(dataloaders['train'])),
iter_step=current_step,
latest=opt['logger']['overwrite_chkp']
)
logger.info('Models and training states saved.')
# validation
if dataloaders.get('val', None) and current_step % opt['train']['val_freq'] == 0 and take_step:
val_metrics = metrics.MetricsDict(metrics=opt['train'].get('metrics', None))
nlls = []
for val_data in dataloaders['val']:
model.feed_data(val_data) # unpack data from data loader
model.test() # run inference
if hasattr(model, 'nll'):
nll = model.nll if model.nll else 0
nlls.append(nll)
"""
Get Visuals
"""
visuals = model.get_current_visuals() # get image results
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)
# Save SR images for reference
sr_img = None
if hasattr(model, 'heats'): # SRFlow
opt['train']['val_comparison'] = False
for heat in model.heats:
for i in range(model.n_sample):
sr_img = tensor2np(visuals['SR', heat, i], denormalize=opt['datasets']['train']['znorm'])
if opt['train']['overwrite_val_imgs']:
save_img_path = os.path.join(img_dir,
'{:s}_h{:03d}_s{:d}.png'.format(img_name, int(heat * 100), i))
else:
save_img_path = os.path.join(img_dir,
'{:s}_{:09d}_h{:03d}_s{:d}.png'.format(img_name,
current_step,
int(heat * 100), i))
util.save_img(sr_img, save_img_path)
else: # regular SR
sr_img = tensor2np(visuals['SR'], denormalize=opt['datasets']['train']['znorm'])
if opt['train']['overwrite_val_imgs']:
save_img_path = os.path.join(img_dir, '{:s}.png'.format(img_name))
else:
save_img_path = os.path.join(img_dir, '{:s}_{:d}.png'.format(img_name, current_step))
if not opt['train']['val_comparison']:
util.save_img(sr_img, save_img_path)
assert sr_img is not None
# Save GT images for reference
gt_img = tensor2np(visuals['HR'], denormalize=opt['datasets']['train']['znorm'])
if opt['train']['save_gt']:
save_img_path_gt = os.path.join(img_dir,
'{:s}_GT.png'.format(img_name))
if not os.path.isfile(save_img_path_gt):
util.save_img(gt_img, save_img_path_gt)
# Save LQ images for reference
if opt['train']['save_lr']:
save_img_path_lq = os.path.join(img_dir,
'{:s}_LQ.png'.format(img_name))
if not os.path.isfile(save_img_path_lq):
lq_img = tensor2np(visuals['LR'], denormalize=opt['datasets']['train']['znorm'])
util.save_img(lq_img, save_img_path_lq, scale=opt['scale'])
# save single images or LQ / SR comparison
if opt['train']['val_comparison']:
lr_img = tensor2np(visuals['LR'], denormalize=opt['datasets']['train']['znorm'])
util.save_img_comp([lr_img, sr_img], save_img_path)
# else:
# util.save_img(sr_img, save_img_path)
"""
Get Metrics
# TODO: test using tensor based metrics (batch) instead of numpy.
"""
val_metrics.calculate_metrics(sr_img, gt_img, crop_size=opt['scale']) # , only_y=True)
avg_metrics = val_metrics.get_averages()
if nlls:
avg_nll = sum(nlls) / len(nlls)
del val_metrics
# log
logger_m = ''
for r in avg_metrics:
formatted_res = r['name'].upper() + ': {:.5g}, '.format(r['average'])
logger_m += formatted_res
if nlls:
logger_m += 'avg_nll: {:.4e} '.format(avg_nll)
logger.info('# Validation # ' + logger_m[:-2])
logger_val = logging.getLogger('val') # validation logger
logger_val.info('<epoch:{:3d}, iter:{:8,d}> '.format(epoch, current_step) + logger_m[:-2])
# memory_usage = torch.cuda.memory_allocated()/(1024.0 ** 3) # in GB
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
for r in avg_metrics:
tb_logger.add_scalar(r['name'], r['average'], current_step)
if nlls:
tb_logger.add_scalar('average nll', avg_nll, current_step)
# tb_logger.flush()
# tb_logger_valid.add_scalar(r['name'], r['average'], current_step)
# tb_logger_valid.flush()
timerData.tick()
timerEpoch.tock()
logger.info('End of epoch {} / {} \t Time Taken: {:.4f} sec'.format(
epoch, total_epochs, timerEpoch.get_last_iteration()))
logger.info('Saving the final model.')
if model.swa:
model.save('latest', loader=dataloaders['train'])
else:
model.save('latest')
logger.info('End of training.')
except KeyboardInterrupt:
# catch a KeyboardInterrupt and save the model and state to resume later
if model.swa:
model.save(current_step, True, loader=dataloaders['train'])
else:
model.save(current_step, True)
model.save_training_state(epoch + (n >= len(dataloaders['train'])), current_step, True)
logger.info('Training interrupted. Latest models and training states saved.')
def tmp_vis(img_t,
to_np=True,
rgb2bgr=True,
remove_batch=False,
denormalize=False,
save_dir='',
tensor_shape='TCHW'):
'''
Visualization function that can be inserted at any point
in the code, works with tensor or np images
img_t: image (shape: [B, ..., W, H])
save_dir: path to save image
tensor_shape: in the case of n_dim == 5, needs to provide the order
of the dimensions
'''
import cv2
from dataops.common import tensor2np
if isinstance(img_t, torch.Tensor) and to_np:
n_dim = img_t.dim()
if n_dim == 5:
# "volumetric" tensor [B, _, _, H, W], where indexes [1] and [2]
# can be (for example) either channels or time. Reduce to 4D tensor
# for visualization
if tensor_shape == 'CTHW':
_, _, n_frames, _, _ = img_t.size()
frames = []
for frame in range(n_frames):
frames.append(img_t[:, :, frame:frame + 1, :, :])
img_t = torch.cat(frames, -1)
elif tensor_shape == 'TCHW':
_, n_frames, _, _, _ = img_t.size()
frames = []
for frame in range(n_frames):
frames.append(img_t[:, frame:frame + 1, :, :, :])
img_t = torch.cat(frames, -1)
elif tensor_shape == 'CTHW_m':
# select only the middle frame of CTHW tensor
_, _, n_frames, _, _ = img_t.size()
center = (n_frames - 1) // 2
img_t = img_t[:, :, center, :, :]
elif tensor_shape == 'TCHW_m':
# select only the middle frame of TCHW tensor
_, n_frames, _, _, _ = img_t.size()
center = (n_frames - 1) // 2
img_t = img_t[:, center, :, :, :]
else:
TypeError("Unrecognized tensor_shape: {}".format(tensor_shape))
img = tensor2np(img_t.detach(),
rgb2bgr=rgb2bgr,
remove_batch=remove_batch,
denormalize=denormalize)
elif isinstance(img_t, np.ndarray) and not to_np:
img = img_t
else:
raise TypeError("img_t type not supported, expected tensor or ndarray")
print("out: ", img.shape)
cv2.imshow('image', img)
cv2.waitKey(0)
if save_dir != '':
cv2.imwrite(save_dir, img)
cv2.destroyAllWindows()
return None
def test_loop(model, opt, dataloaders, data_params):
logger = util.get_root_logger()
# read data_params
znorms = data_params['znorm']
# prepare the metric calculation classes for RGB and Y_only images
calc_metrics = opt.get('metrics', None)
if calc_metrics:
test_metrics = metrics.MetricsDict(metrics=calc_metrics)
test_metrics_y = metrics.MetricsDict(metrics=calc_metrics)
for phase, dataloader in dataloaders.items():
name = dataloader.dataset.opt['name']
logger.info('\nTesting [{:s}]...'.format(name))
dataset_dir = os.path.join(opt['path']['results_root'], name)
util.mkdir(dataset_dir)
nlls = []
for data in dataloader:
znorm = znorms[name]
need_HR = False if dataloader.dataset.opt[
'dataroot_HR'] is None else True
# set up per image CEM wrapper if configured
CEM_net = get_CEM(opt, data)
model.feed_data(data,
need_HR=need_HR) # unpack data from data loader
# img_path = data['LR_path'][0]
img_path = get_img_path(data)
img_name = os.path.splitext(os.path.basename(img_path))[0]
# test with eval mode. This only affects layers like batchnorm and dropout.
test_mode = opt.get('test_mode', None)
if test_mode == 'x8':
# geometric self-ensemble
# model.test_x8(CEM_net=CEM_net)
break
elif test_mode == 'chop':
# chop images in patches/crops, to reduce VRAM usage
# model.test_chop(patch_size=opt.get('chop_patch_size', 100),
# step=opt.get('chop_step', 0.9),
# CEM_net=CEM_net)
break
else:
# normal inference
model.test(CEM_net=CEM_net) # run inference
if hasattr(model, 'nll'):
nll = model.nll if model.nll else 0
nlls.append(nll)
# get image results
visuals = model.get_current_visuals(need_HR=need_HR)
res_options = visuals_check(visuals.keys(),
opt.get('val_comparison', None))
# save images
save_img_path = os.path.join(dataset_dir,
img_name + opt.get('suffix', ''))
# Save SR images for reference
sr_img = None
if hasattr(model, 'heats'): # SRFlow
opt['val_comparison'] = False
for heat in model.heats:
for i in range(model.n_sample):
for save_img_name in res_options['save_imgs']:
imn = '_' + save_img_name if len(
res_options['save_imgs']) > 1 else ''
imn += '_h{:03d}_s{:d}'.format(int(heat * 100), i)
util.save_img(
tensor2np(visuals[save_img_name, heat, i],
denormalize=znorm),
save_img_path + imn + '.png')
else: # regular SR
if not opt['val_comparison']:
for save_img_name in res_options['save_imgs']:
imn = '_' + save_img_name if len(
res_options['save_imgs']) > 1 else ''
util.save_img(
tensor2np(visuals[save_img_name],
denormalize=znorm),
save_img_path + imn + '.png')
# save single images or lr / sr comparison
if opt['val_comparison'] and len(res_options['save_imgs']) > 1:
comp_images = [
tensor2np(visuals[save_img_name], denormalize=znorm)
for save_img_name in res_options['save_imgs']
]
util.save_img_comp(comp_images, save_img_path + '.png')
# else:
# util.save_img(sr_img, save_img_path)
# calculate metrics if HR dataset is provided and metrics are configured in options
if need_HR and calc_metrics and res_options['aligned_metrics']:
metric_imgs = [
tensor2np(visuals[x], denormalize=znorm)
for x in res_options['compare_imgs']
]
test_results = test_metrics.calculate_metrics(
metric_imgs[0], metric_imgs[1], crop_size=opt['scale'])
# prepare single image metrics log message
logger_m = '{:20s} -'.format(img_name)
for k, v in test_results:
formatted_res = k.upper() + ': {:.6f}, '.format(v)
logger_m += formatted_res
if gt_img.shape[2] == 3: # RGB image, calculate y_only metrics
test_results_y = test_metrics_y.calculate_metrics(
metric_imgs[0],
metric_imgs[1],
crop_size=opt['scale'],
only_y=True)
# add the y only results to the single image log message
for k, v in test_results_y:
formatted_res = k.upper() + ': {:.6f}, '.format(v)
logger_m += formatted_res
logger.info(logger_m)
else:
logger.info(img_name)
# average metrics results for the dataset
if need_HR and calc_metrics:
# aggregate the metrics results (automatically resets the metric classes)
avg_metrics = test_metrics.get_averages()
avg_metrics_y = test_metrics_y.get_averages()
# prepare log average metrics message
agg_logger_m = ''
for r in avg_metrics:
formatted_res = r['name'].upper() + ': {:.6f}, '.format(
r['average'])
agg_logger_m += formatted_res
logger.info(
'----Average metrics results for {}----\n\t'.format(name) +
agg_logger_m[:-2])
if len(avg_metrics_y > 0):
# prepare log average Y channel metrics message
agg_logger_m = ''
for r in avg_metrics_y:
formatted_res = r['name'].upper(
) + '_Y' + ': {:.6f}, '.format(r['average'])
agg_logger_m += formatted_res
logger.info('----Y channel, average metrics ----\n\t' +
agg_logger_m[:-2])
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']:
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
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('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:
tb_logger = SummaryWriter(logdir='../tb_logger/' + opt['name']) #for version tensorboardX >= 1.7
except:
tb_logger = SummaryWriter(log_dir='../tb_logger/' + opt['name']) #for version tensorboardX < 1.6
# 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)
# if the model does not change and input sizes remain the same during training then there may be benefit
# from setting torch.backends.cudnn.benchmark = True, otherwise it may stall training
torch.backends.cudnn.benchmark = True
# torch.backends.cudnn.deterministic = True
# create train and val dataloader
val_loader = False
for phase, dataset_opt in opt['datasets'].items():
if phase == 'train':
train_set = create_dataset(dataset_opt)
batch_size = dataset_opt.get('batch_size', 4)
virtual_batch_size = dataset_opt.get('virtual_batch_size', batch_size)
virtual_batch_size = virtual_batch_size if virtual_batch_size > batch_size else batch_size
train_size = int(math.ceil(len(train_set) / 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']
virtual_step = current_step * virtual_batch_size / batch_size \
if virtual_batch_size and virtual_batch_size > batch_size else current_step
model.resume_training(resume_state) # handle optimizers and schedulers
model.update_schedulers(opt['train']) # updated schedulers in case JSON configuration has changed
del resume_state
# start the iteration time when resuming
t0 = time.time()
else:
current_step = 0
virtual_step = 0
start_epoch = 0
# training
logger.info('Start training from epoch: {:d}, iter: {:d}'.format(start_epoch, current_step))
try:
for epoch in range(start_epoch, total_epochs*(virtual_batch_size//batch_size)):
for n, train_data in enumerate(train_loader,start=1):
if virtual_step == 0:
# first iteration start time
t0 = time.time()
virtual_step += 1
take_step = False
if virtual_step > 0 and virtual_step * batch_size % virtual_batch_size == 0:
current_step += 1
take_step = True
if current_step > total_iters:
break
# training
model.feed_data(train_data)
model.optimize_parameters(virtual_step)
# log
if current_step % opt['logger']['print_freq'] == 0 and take_step:
# iteration end time
t1 = time.time()
logs = model.get_current_log()
message = '<epoch:{:3d}, iter:{:8,d}, lr:{:.3e}, i_time: {:.4f} sec.> '.format(
epoch, current_step, model.get_current_learning_rate(current_step), (t1 - t0))
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)
# # start time for next iteration
# t0 = time.time()
# update learning rate
if model.optGstep and model.optDstep and take_step:
model.update_learning_rate(current_step, warmup_iter=opt['train'].get('warmup_iter', -1))
# save models and training states (changed to save models before validation)
if current_step % opt['logger']['save_checkpoint_freq'] == 0 and take_step:
if model.swa:
model.save(current_step, opt['logger']['overwrite_chkp'], loader=train_loader)
else:
model.save(current_step, opt['logger']['overwrite_chkp'])
model.save_training_state(epoch + (n >= len(train_loader)), current_step, opt['logger']['overwrite_chkp'])
logger.info('Models and training states saved.')
# validation
if val_loader and current_step % opt['train']['val_freq'] == 0 and take_step:
val_sr_imgs_list = []
val_gt_imgs_list = []
val_metrics = metrics.MetricsDict(metrics=opt['train'].get('metrics', None))
for val_data in val_loader:
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(val_data)
"""
Get Visuals
"""
visuals = model.get_current_visuals()
sr_img = tensor2np(visuals['SR'], denormalize=opt['datasets']['train']['znorm'])
gt_img = tensor2np(visuals['HR'], denormalize=opt['datasets']['train']['znorm'])
# Save SR images for reference
if opt['train']['overwrite_val_imgs']:
save_img_path = os.path.join(img_dir, '{:s}.png'.format(\
img_name))
else:
save_img_path = os.path.join(img_dir, '{:s}_{:d}.png'.format(\
img_name, current_step))
# save single images or lr / sr comparison
if opt['train']['val_comparison']:
lr_img = tensor2np(visuals['LR'], denormalize=opt['datasets']['train']['znorm'])
util.save_img_comp(lr_img, sr_img, save_img_path)
else:
util.save_img(sr_img, save_img_path)
"""
Get Metrics
# TODO: test using tensor based metrics (batch) instead of numpy.
"""
crop_size = opt['scale']
val_metrics.calculate_metrics(sr_img, gt_img, crop_size = crop_size) #, only_y=True)
avg_metrics = val_metrics.get_averages()
del val_metrics
# log
logger_m = ''
for r in avg_metrics:
#print(r)
formatted_res = r['name'].upper()+': {:.5g}, '.format(r['average'])
logger_m += formatted_res
logger.info('# Validation # '+logger_m[:-2])
logger_val = logging.getLogger('val') # validation logger
logger_val.info('<epoch:{:3d}, iter:{:8,d}> '.format(epoch, current_step)+logger_m[:-2])
# memory_usage = torch.cuda.memory_allocated()/(1024.0 ** 3) # in GB
# tensorboard logger
if opt['use_tb_logger'] and 'debug' not in opt['name']:
for r in avg_metrics:
tb_logger.add_scalar(r['name'], r['average'], current_step)
# # reset time for next iteration to skip the validation time from calculation
# t0 = time.time()
if current_step % opt['logger']['print_freq'] == 0 and take_step or \
(val_loader and current_step % opt['train']['val_freq'] == 0 and take_step):
# reset time for next iteration to skip the validation time from calculation
t0 = time.time()
logger.info('Saving the final model.')
if model.swa:
model.save('latest', loader=train_loader)
else:
model.save('latest')
logger.info('End of training.')
except KeyboardInterrupt:
# catch a KeyboardInterrupt and save the model and state to resume later
if model.swa:
model.save(current_step, True, loader=train_loader)
else:
model.save(current_step, True)
model.save_training_state(epoch + (n >= len(train_loader)), current_step, True)
logger.info('Training interrupted. Latest models and training states saved.')
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)
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
#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)
pred = tensor2np(pred,
denormalize=znorm).clip(0, 255) # uint8
pred_90 = tensor2np(pred_90,
denormalize=znorm).clip(0,
255) # uint8
pred_180 = tensor2np(pred_180,
denormalize=znorm).clip(0,
255) # uint8
pred_270 = tensor2np(pred_270,
denormalize=znorm).clip(0,
255) # uint8
pred_f = tensor2np(pred_f,
denormalize=znorm).clip(0, 255) # uint8
pred_90f = tensor2np(pred_90f,
denormalize=znorm).clip(0,
255) # uint8
pred_180f = tensor2np(pred_180f,
denormalize=znorm).clip(0,
255) # uint8
pred_270f = tensor2np(pred_270f,
denormalize=znorm).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
#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(highres_output,
denormalize=znorm) # 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 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(visuals['SR'], denormalize=znorm) # uint8
# 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))
