def comput_PSNR_SSIM(self, pred, gt, shave_border=0):
if isinstance(pred, torch.Tensor):
pred = tensor2np(pred, self.opt.rgb_range)
pred = pred.astype(np.float32)
if isinstance(gt, torch.Tensor):
gt = tensor2np(gt, self.opt.rgb_range)
gt = gt.astype(np.float32)
height, width = pred.shape[:2]
pred = pred[shave_border:height - shave_border,
shave_border:width - shave_border]
gt = gt[shave_border:height - shave_border,
shave_border:width - shave_border]
if pred.shape[2] == 3 and gt.shape[2] == 3:
pred_y = rgb2ycbcr(pred)[:, :, 0]
gt_y = rgb2ycbcr(gt)[:, :, 0]
elif pred.shape[2] == 1 and gt.shape[2] == 1:
pred_y = pred[:, :, 0]
gt_y = gt[:, :, 0]
else:
raise ValueError('Input or output channel is not 1 or 3!')
psnr_ = calc_PSNR(pred_y, gt_y)
ssim_ = calc_ssim(pred_y, gt_y)
return psnr_, ssim_
def test(self):
epoch = self.scheduler.last_epoch + 1
self.ckp.write_log('\nEvaluation:')
self.ckp.add_log(torch.zeros(1, len(self.scale)), False)
self.model.eval()
# We can use custom forward function
def _test_forward(x, scale):
if self.args.self_ensemble:
return utils.x8_forward(x, self.model, self.args.precision)
elif self.args.chop_forward:
return utils.chop_forward(x, self.model, scale)
else:
return self.model(x)
timer_test = utils.timer()
set_name = type(self.loader_test.dataset).__name__
for idx_scale, scale in enumerate(self.scale):
eval_acc = 0
self._scale_change(idx_scale, self.loader_test)
for idx_img, (input, target, _) in enumerate(self.loader_test):
input, target = self._prepare(input, target, volatile=True)
output = _test_forward(input, scale)
eval_acc += utils.calc_PSNR(
output, target, set_name, self.args.rgb_range, scale)
self.ckp.save_results(idx_img, input, output, target, scale)
self.ckp.log_test[-1, idx_scale] = eval_acc / len(self.loader_test)
best = self.ckp.log_test.max(0)
performance = 'PSNR: {:.3f}'.format(
self.ckp.log_test[-1, idx_scale])
self.ckp.write_log(
'[{} x{}]\t{} (Best: {:.3f} from epoch {})'.format(
set_name,
scale,
performance,
best[0][idx_scale],
best[1][idx_scale] + 1))
if best[1][0] + 1 == epoch:
is_best = True
else:
is_best = False
self.ckp.write_log(
'Time: {:.2f}s\n'.format(timer_test.toc()), refresh=True)
self.ckp.save(self, epoch, is_best=is_best)
def test_generator(self, test_num_patch = 200, test_num_image = 5, load = False):
self.generator.eval()
# test for patches
start = time.time()
test_list_phone = sorted(glob(self.config.test_path_phone_patch))
PSNR_phone_enhanced_list = np.zeros([test_num_patch])
indexes = []
for i in range(test_num_patch):
index = np.random.randint(len(test_list_phone))
indexes.append(index)
test_img = scipy.misc.imread(test_list_phone[index], mode = "RGB").astype("float32")
test_patch_phone = get_patch(test_img, self.config.patch_size)
test_patch_phone = preprocess(test_patch_phone)
with torch.no_grad():
test_patch_phone = torch.from_numpy(np.transpose(test_patch_phone, (2,1,0))).float().unsqueeze(0)
if torch.cuda.is_available():
test_patch_phone = test_patch_phone.cuda()
test_patch_enhanced = self.generator(test_patch_phone)
test_patch_enhanced = test_patch_enhanced.cpu().data.numpy()
test_patch_enhanced = np.transpose(test_patch_enhanced.cpu().data.numpy(), (0,2,3,1))
test_patch_phone = np.transpose(test_patch_phone.cpu().data.numpy(), (0,2,3,1))
if i % 50 == 0:
imageio.imwrite(("%s/phone_%d.png" %(self.result_img_dir, i)), postprocess(test_patch_phone[0]))
imageio.imwrite(("%s/enhanced_%d.png" %(self.result_img_dir,i)), postprocess(test_patch_enhanced[0]))
PSNR = calc_PSNR(postprocess(test_patch_enhanced[0]), postprocess(test_patch_phone))
PSNR_phone_enhanced_list[i] = PSNR
print("(runtime: %.3f s) Average test PSNR for %d random test image patches: phone-enhanced %.3f" %(time.time()-start, test_num_patch, np.mean(PSNR_phone_enhanced_list)))
# test for images
start = time.time()
test_list_phone = sorted(glob(self.config.test_path_phone_image))
PSNR_phone_enhanced_list = np.zeros([test_num_image])
indexes = []
for i in range(test_num_image):
index = i
indexes.append(index)
test_image_phone = preprocess(scipy.misc.imread(test_list_phone[index], mode = "RGB").astype("float32"))
with torch.no_grad():
test_image_phone = torch.from_numpy(np.transpose(test_image_phone, (2,1,0))).float().unsqueeze(0)
if torch.cuda.is_available():
test_image_phone = test_image_phone.cuda()
test_image_enhanced = self.generator(test_image_phone)
test_image_enhanced = test_image_enhanced.cpu().data.numpy()
test_image_enhanced = np.transpose(test_image_enhanced.cpu().data.numpy(), (0,2,3,1))
test_image_phone = np.transpose(test_image_phone.cpu().data.numpy(), (0,2,3,1))
imageio.imwrite(("%s/phone_%d.png" %(self.sample_dir, i)), postprocess(test_image_phone[0]))
imageio.imwrite(("%s/enhanced_%d.png" %(self.sample_dir, i)), postprocess(test_image_enhanced[0]))
PSNR = calc_PSNR(postprocess(test_image_enhanced[0]), postprocess(test_image_phone[0]))
PSNR_phone_enhanced_list[i] = PSNR
if test_num_image > 0:
print("(runtime: %.3f s) Average test PSNR for %d random full test images: original-enhanced %.3f" %(time.time()-start, test_num_image, np.mean(PSNR_phone_enhanced_list)))
def test(self):
epoch = self.scheduler.last_epoch + 1
self.ckp.write_log('\nEvaluation:')
self.ckp.add_log(torch.zeros(1, len(self.scale)), False)
self.model.eval()
# We can use custom forward function
def _test_forward(x, scale):
if self.args.self_ensemble:
return utils.x8_forward(x, self.model, self.args.precision)
elif self.args.chop_forward:
return utils.chop_forward(x, self.model, scale)
else:
return self.model(x)
timer_test = utils.timer()
set_name = type(self.loader_test.dataset).__name__
## creation of log file
old_stdout = sys.stdout
TooLargeImageslog_file = open("TooLargeImages.log", "w")
for idx_scale, scale in enumerate(self.scale):
eval_acc = 0
validation_loss = 0
self._scale_change(idx_scale, self.loader_test)
for idx_img, (input, target, _) in enumerate(self.loader_test):
input, target = self._prepare(input, target, volatile=True)
try:
output = _test_forward(input, scale)
eval_acc += utils.calc_PSNR(output, target, set_name,
self.args.rgb_range, scale)
self.ckp.save_results(idx_img, input, output, target,
scale)
validation_loss += self._calc_loss_for_validation(
output, target)
except:
sys.stdout = TooLargeImageslog_file
print(
"image: ", idx_img + 1 + self.args.n_train
) # note: idx_img + 1 + self.args.n_train (used to be just idx_img)
sys.stdout = old_stdout
self.ckp.log_test_loss[-1, idx_scale] = validation_loss / len(
self.loader_test)
self.ckp.log_test[-1, idx_scale] = eval_acc / len(self.loader_test)
best = self.ckp.log_test.max(0)
performance = 'PSNR: {:.3f}'.format(self.ckp.log_test[-1,
idx_scale])
self.ckp.write_log(
'[{} x{}]\t{} (Best: {:.3f} from epoch {})'.format(
set_name, scale, performance, best[0][idx_scale],
best[1][idx_scale] + 1))
## CLOSING LOG FILE ##
TooLargeImageslog_file.close()
######################
if best[1][0] + 1 == epoch:
is_best = True
else:
is_best = False
self.ckp.write_log('Time: {:.2f}s\n'.format(timer_test.toc()),
refresh=True)
self.ckp.save(
self, epoch, is_best=is_best
) #this is where we save (in utils) , in save function we also plot.
def validation(valid_path, result_path, model, scale):
model.eval()
count = 0
PSNR = 0
SSIM = 0
# RGB mean for ImageNet
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
unnormalize = transforms.Normalize(mean=[-2.118, -2.036, -1.804],
std=[4.367, 4.464, 4.444])
filepath = valid_path.split('_LR')[0]
file = os.listdir(filepath)
lr_file = os.listdir(valid_path + '/X' + str(scale))
if lr_file[0].find('jpeg') != -1:
lr_type = '.jpeg'
elif lr_file[0].find('jpg') != -1:
lr_type = '.jpg'
elif lr_file[0].find('bmp') != -1:
lr_type = '.bmp'
else:
lr_type = '.png'
if file[0].find('jpeg') != -1:
hr_type = '.jpeg'
elif file[0].find('jpg') != -1:
hr_type = '.jpg'
elif file[0].find('bmp') != -1:
hr_type = '.bmp'
else:
hr_type = '.png'
file.sort()
length = file.__len__()
with torch.no_grad():
with tqdm(total=length) as pbar:
for idx_img in range(length):
time.sleep(0.01)
pbar.update(1)
img_name = file[idx_img].split(hr_type)[0]
img_hr_rgb = imageio.imread(filepath + '/' + img_name +
hr_type)
img_lr_rgb = imageio.imread(valid_path + '/X' + str(scale) +
'/' + img_name + 'x' + str(scale) +
lr_type)
img_lr_rgb, img_hr_rgb = common.set_channel(
img_lr_rgb, img_hr_rgb, 3)
img_lr_rgb, img_hr_rgb = common.np2Tensor(
img_lr_rgb, img_hr_rgb, 255)
img_lr_rgb = normalize(img_lr_rgb) # Normalize
# img_hr_rgb = normalize(img_hr_rgb)
img_lr_rgb = Variable(img_lr_rgb).view(1, img_lr_rgb.shape[0],
img_lr_rgb.shape[1],
img_lr_rgb.shape[2])
img_hr_rgb = Variable(img_hr_rgb).view(1, img_hr_rgb.shape[0],
img_hr_rgb.shape[1],
img_hr_rgb.shape[2])
img_lr_rgb = img_lr_rgb.cuda()
# SR = F.interpolate(img_lr_rgb, scale_factor=scale)
SR = model(img_lr_rgb)
# SR = model(img_lr_rgb, scale)
SR = unnormalize(SR.data[0].cpu())
# plt.figure()
# plt.subplot(1,3, 1)
# plt.imshow(img_lr_rgb.data[0].cpu().numpy().transpose(1,2,0))
# plt.subplot(1,3, 2)
# plt.imshow(SR.numpy().transpose(1,2,0))
# plt.subplot(1,3, 3)
# plt.imshow(img_hr_rgb.data[0].cpu().numpy().transpose(1,2,0))
# plt.show()
PSNR += utils.calc_PSNR(SR,
img_hr_rgb.data[0],
rgb_range=255,
shave=scale)
SSIM += utils.calc_SSIM(SR,
img_hr_rgb.data[0],
rgb_range=255,
shave=scale)
count = count + 1
result = SR.mul(255).clamp(0, 255).round()
result = result.numpy().astype(np.uint8)
result = result.transpose((1, 2, 0))
result = Image.fromarray(result)
result.save(result_path + '/' + img_name + '_DeFiAN_x' +
str(scale) + '.png')
Avg_PSNR = PSNR / count
Avg_SSIM = SSIM / count
return Avg_PSNR, Avg_SSIM
