def valid(dataset, loader, model, criterion, device):
losses = AverageMeter()
psnrs = AverageMeter()
model.eval()
with tqdm(total=len(dataset)) as t:
t.set_description("valid")
for data in loader:
lr, hr, sigma = data
lr = lr.to(device)
hr = hr.to(device)
with torch.no_grad():
sr = model(lr)
loss = criterion(sr, hr)
losses.update(loss.item(), lr.shape[0])
sr = quantize(sr, [0, 255])
psnr = calc_psnr(sr, hr)
psnrs.update(psnr.item(), lr.shape[0])
t.set_postfix(loss='{:.4f}'.format(losses.avg))
t.update(lr.shape[0])
return psnrs.avg
def main():
global args, model
args = parser.parse_args()
print(args)
if args.gpu and not torch.cuda.is_available():
raise Exception("No GPU found!")
if not os.path.exists(args.test_result):
os.makedirs(args.test_result)
if not is_ready(args):
prepare_data(args)
cudnn.benchmark = True
device = torch.device(('cuda:' + args.gpu_id) if args.gpu else 'cpu')
model = Grad_none.GRAD(feats=args.feats,
basic_conv=args.basic_conv,
tail_conv=args.tail_conv)
checkpoint_file = torch.load(args.test_checkpoint)
model.load_state_dict(checkpoint_file['model'])
model.eval()
model = model.to(device)
psnrs = AverageMeter()
with tqdm(total=100) as t:
t.set_description("test")
for idx in range(0, 100):
with h5py.File(
"{}/DIV2K_np_test_{}.h5".format(args.h5file_dir,
args.test_sigma),
'r') as h5:
l_image, h_image = h5['l'][str(idx)][()], h5['h'][str(idx)][()]
l_image = np2tensor(l_image)
h_image = np2tensor(h_image)
l_image = l_image.unsqueeze(0)
h_image = h_image.unsqueeze(0)
l_image = l_image.to(device)
h_image = h_image.to(device)
with torch.no_grad():
output = model(l_image)
output = quantize(output, [0, 255])
psnr = calc_psnr(output, h_image)
psnrs.update(psnr.item(), 1)
if args.test_save:
save_image_path = "{}/{:04d}.png".format(
args.test_result, idx)
output = output.squeeze(0)
output = output.data.permute(1, 2, 0)
save_image = pil_image.fromarray(
output.byte().cpu().numpy())
save_image.save(save_image_path)
t.update(1)
print("PSNR: {:.4f}".format(psnrs.avg))
