fea94d9805d3be49f5ae55d0ab3f35f07a5c6d12,test.py,,,#,15

Before Change

        if torch.cuda.is_available():
            image = image.cuda()

        out = model(image.unsqueeze(0)).cpu().data[0]
        mse = ((target - out) ** 2).mean()
        psnr = 10 * log10(1 / mse)
        out_img = ToPILImage()(out)
        out_img.save(out_path + "psnr_%.4f_" % psnr + image_name)

After Change

        image = Image.open(data_path + image_name)
        image = Variable(ToTensor()(image))
        target = Image.open(target_path + image_name)
        target = Variable(ToTensor()(target))
        if torch.cuda.is_available():
            image = image.cuda()
            target = target.cuda()

        out = model(image.unsqueeze(0))[0]
        mse = ((target - out) ** 2).mean()
        psnr = 10 * log10(1 / mse)
        ssim = pytorch_ssim.ssim(out, target)
        out_img = ToPILImage()(out.data)
        out_img.save(out_path + "psnr_%.4f_ssim_%.4f_" % (psnr, ssim) + image_name)

Instances