def eval(self):
with tqdm(total=len(self.val_loader),
miniters=1,
desc='Val Epoch: [{}/{}]'.format(self.epoch,
self.nEpochs)) as t1:
psnr_list, ssim_list = [], []
for iteration, batch in enumerate(self.val_loader, 1):
lr, hr = Variable(batch[0]), Variable(batch[1])
#for lr, hr in self.val_loader:
if self.cuda:
lr, hr = lr.cuda(), hr.cuda()
self.model.eval()
with torch.no_grad():
if self.cfg['algorithm'] == 'SMSR':
sr, sparsity = self.model(lr)
else:
sr = self.model(lr)
loss = self.loss(sr, hr)
batch_psnr, batch_ssim = [], []
for c in range(sr.shape[0]):
if not self.cfg['data']['normalize']:
predict_sr = (sr[c, ...].cpu().numpy().transpose(
(1, 2, 0))) * 255
ground_truth = (hr[c, ...].cpu().numpy().transpose(
(1, 2, 0))) * 255
else:
predict_sr = (sr[c, ...].cpu().numpy().transpose(
(1, 2, 0)) + 1) * 127.5
ground_truth = (hr[c, ...].cpu().numpy().transpose(
(1, 2, 0)) + 1) * 127.5
psnr = calculate_psnr(predict_sr, ground_truth, 255)
ssim = calculate_ssim(predict_sr, ground_truth, 255)
batch_psnr.append(psnr)
batch_ssim.append(ssim)
avg_psnr = np.array(batch_psnr).mean()
avg_ssim = np.array(batch_ssim).mean()
psnr_list.extend(batch_psnr)
ssim_list.extend(batch_ssim)
t1.set_postfix_str(
'Batch loss: {:.4f}, PSNR: {:.4f}, SSIM: {:.4f}'.format(
loss.item(), avg_psnr, avg_ssim))
t1.update()
self.records['Epoch'].append(self.epoch)
self.records['PSNR'].append(np.array(psnr_list).mean())
self.records['SSIM'].append(np.array(ssim_list).mean())
save_config(
self.log_name, 'Val Epoch {}: PSNR={:.4f}, SSIM={:.4f}'.format(
self.epoch, self.records['PSNR'][-1],
self.records['SSIM'][-1]))
self.writer.add_scalar('PSNR_epoch', self.records['PSNR'][-1],
self.epoch)
self.writer.add_scalar('SSIM_epoch', self.records['SSIM'][-1],
self.epoch)
self.writer.add_image('image_SR', sr[0], self.epoch)
self.writer.add_image('image_LR', lr[0], self.epoch)
self.writer.add_image('image_HR', hr[0], self.epoch)
def eval(self):
with tqdm(total=len(self.val_loader),
miniters=1,
desc='Val Epoch: [{}/{}]'.format(self.epoch,
self.max_epochs)) as t:
psnr_list, ssim_list, loss_list = [], [], []
for lr, hr in self.val_loader:
lr, hr = lr.cuda(), hr.cuda()
self.srntt.eval()
with torch.no_grad():
sr, _ = self.srntt(lr, None, None)
loss = self.reconst_loss(sr, hr)
batch_psnr, batch_ssim = [], []
for c in range(sr.shape[0]):
predict_sr = (sr[c, ...].cpu().numpy().transpose(
(1, 2, 0)) + 1) * 127.5
ground_truth = (hr[c, ...].cpu().numpy().transpose(
(1, 2, 0)) + 1) * 127.5
psnr = utils.calculate_psnr(predict_sr, ground_truth, 255)
ssim = utils.calculate_ssim(predict_sr, ground_truth, 255)
batch_psnr.append(psnr)
batch_ssim.append(ssim)
avg_psnr = np.array(batch_psnr).mean()
avg_ssim = np.array(batch_ssim).mean()
psnr_list.extend(batch_psnr)
ssim_list.extend(batch_ssim)
t.set_postfix_str(
'Batch loss: {:.4f}, PSNR: {:.4f}, SSIM: {:.4f}'.format(
loss.item(), avg_psnr, avg_ssim))
t.update()
self.records['Epoch'].append(self.epoch)
self.records['PSNR'].append(np.array(psnr_list).mean())
self.records['SSIM'].append(np.array(ssim_list).mean())
self.logger.log('Val Epoch {}: PSNR={}, SSIM={}'.format(
self.epoch, self.records['PSNR'][-1],
self.records['SSIM'][-1]))
def test():
"""
Test(Zooming SloMo) - inference on set of input data or Vid4 data
"""
# set context and load the model
ctx = get_extension_context(args.context)
nn.set_default_context(ctx)
nn.load_parameters(args.model)
input_dir = args.input_dir
n_ot = 7
# list all input sequence folders containing input frames
inp_dir_list = sorted(glob.glob(input_dir + '/*'))
inp_dir_name_list = []
avg_psnr_l = []
avg_psnr_y_l = []
avg_ssim_y_l = []
sub_folder_name_l = []
save_folder = 'results'
# for each sub-folder
for inp_dir in inp_dir_list:
gt_tested_list = []
inp_dir_name = inp_dir.split('/')[-1]
sub_folder_name_l.append(inp_dir_name)
inp_dir_name_list.append(inp_dir_name)
save_inp_folder = osp.join(save_folder, inp_dir_name)
img_low_res_list = sorted(glob.glob(inp_dir + '/*'))
util.mkdirs(save_inp_folder)
imgs = util.read_seq_imgs_(inp_dir)
img_gt_l = []
if args.metrics:
replace_str = 'LR'
for img_gt_path in sorted(glob.glob(osp.join(inp_dir.replace(replace_str, 'HR'), '*'))):
img_gt_l.append(util.read_image(img_gt_path))
avg_psnr, avg_psnr_sum, cal_n = 0, 0, 0
avg_psnr_y, avg_psnr_sum_y = 0, 0
avg_ssim_y, avg_ssim_sum_y = 0, 0
skip = args.metrics
select_idx_list = util.test_index_generation(
skip, n_ot, len(img_low_res_list))
# process each image
for select_idxs in select_idx_list:
# get input images
select_idx = [select_idxs[0]]
gt_idx = select_idxs[1]
imgs_in = F.gather_nd(
imgs, indices=nn.Variable.from_numpy_array(select_idx))
imgs_in = F.reshape(x=imgs_in, shape=(1,) + imgs_in.shape)
output = zooming_slo_mo_network(imgs_in, args.only_slomo)
outputs = output[0]
outputs.forward(clear_buffer=True)
for idx, name_idx in enumerate(gt_idx):
if name_idx in gt_tested_list:
continue
gt_tested_list.append(name_idx)
output_f = outputs.d[idx, :, :, :]
output = util.tensor2img(output_f)
cv2.imwrite(osp.join(save_inp_folder,
'{:08d}.png'.format(name_idx + 1)), output)
print("Saving :", osp.join(save_inp_folder,
'{:08d}.png'.format(name_idx + 1)))
if args.metrics:
# calculate PSNR
output = output / 255.
ground_truth = np.copy(img_gt_l[name_idx])
cropped_output = output
cropped_gt = ground_truth
crt_psnr = util.calculate_psnr(
cropped_output * 255, cropped_gt * 255)
cropped_gt_y = util.bgr2ycbcr(cropped_gt, only_y=True)
cropped_output_y = util.bgr2ycbcr(
cropped_output, only_y=True)
crt_psnr_y = util.calculate_psnr(
cropped_output_y * 255, cropped_gt_y * 255)
crt_ssim_y = util.calculate_ssim(
cropped_output_y * 255, cropped_gt_y * 255)
avg_psnr_sum += crt_psnr
avg_psnr_sum_y += crt_psnr_y
avg_ssim_sum_y += crt_ssim_y
cal_n += 1
if args.metrics:
avg_psnr = avg_psnr_sum / cal_n
avg_psnr_y = avg_psnr_sum_y / cal_n
avg_ssim_y = avg_ssim_sum_y / cal_n
avg_psnr_l.append(avg_psnr)
avg_psnr_y_l.append(avg_psnr_y)
avg_ssim_y_l.append(avg_ssim_y)
if args.metrics:
print('################ Tidy Outputs ################')
for name, ssim, psnr_y in zip(sub_folder_name_l, avg_ssim_y_l, avg_psnr_y_l):
print(
'Folder {} - Average SSIM: {:.6f} PSNR-Y: {:.6f} dB. '.format(name, ssim, psnr_y))
print('################ Final Results ################')
print('Total Average SSIM: {:.6f} PSNR-Y: {:.6f} dB for {} clips. '.format(
sum(avg_ssim_y_l) / len(avg_ssim_y_l), sum(avg_psnr_y_l) /
len(avg_psnr_y_l),
len(inp_dir_list)))
def eval(self):
with tqdm(total=len(self.val_loader),
miniters=1,
desc='Val Epoch: [{}/{}]'.format(self.epoch,
self.nEpochs)) as t1:
psnr_list, ssim_list = [], []
for iteration, batch in enumerate(self.val_loader, 1):
ms_image, lms_image, pan_image, bms_image, file = Variable(
batch[0]), Variable(batch[1]), Variable(
batch[2]), Variable(batch[3]), (batch[4])
if self.cuda:
ms_image, lms_image, pan_image, bms_image = ms_image.cuda(
self.gpu_ids[0]), lms_image.cuda(
self.gpu_ids[0]), pan_image.cuda(
self.gpu_ids[0]), bms_image.cuda(
self.gpu_ids[0])
self.model.eval()
with torch.no_grad():
y = self.model(lms_image, bms_image, pan_image)
loss = self.loss(y, ms_image)
batch_psnr, batch_ssim = [], []
y = y[:, 0:3, :, :]
ms_image = ms_image[:, 0:3, :, :]
for c in range(y.shape[0]):
if not self.cfg['data']['normalize']:
predict_y = (y[c, ...].cpu().numpy().transpose(
(1, 2, 0))) * 255
ground_truth = (ms_image[c,
...].cpu().numpy().transpose(
(1, 2, 0))) * 255
else:
predict_y = (y[c, ...].cpu().numpy().transpose(
(1, 2, 0)) + 1) * 127.5
ground_truth = (
ms_image[c, ...].cpu().numpy().transpose(
(1, 2, 0)) + 1) * 127.5
psnr = calculate_psnr(predict_y, ground_truth, 255)
ssim = calculate_ssim(predict_y, ground_truth, 255)
batch_psnr.append(psnr)
batch_ssim.append(ssim)
avg_psnr = np.array(batch_psnr).mean()
avg_ssim = np.array(batch_ssim).mean()
psnr_list.extend(batch_psnr)
ssim_list.extend(batch_ssim)
t1.set_postfix_str(
'Batch loss: {:.4f}, PSNR: {:.4f}, SSIM: {:.4f}'.format(
loss.item(), avg_psnr, avg_ssim))
t1.update()
self.records['Epoch'].append(self.epoch)
self.records['PSNR'].append(np.array(psnr_list).mean())
self.records['SSIM'].append(np.array(ssim_list).mean())
save_config(
self.log_name, 'Val Epoch {}: PSNR={:.4f}, SSIM={:.4f}'.format(
self.epoch, self.records['PSNR'][-1],
self.records['SSIM'][-1]))
self.writer.add_scalar('PSNR_epoch', self.records['PSNR'][-1],
self.epoch)
self.writer.add_scalar('SSIM_epoch', self.records['SSIM'][-1],
self.epoch)