test.py

import warnings

warnings.filterwarnings("ignore")

import os
import numpy as np
from tqdm import tqdm
from imageio import imsave
from PIL import Image
import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
from torch.utils import data
from torchvision.utils import save_image
from skimage.metrics import mean_squared_error, peak_signal_noise_ratio
from skimage.metrics import structural_similarity as ssim
from models.generator.generator import Generator
from datasets.dataset import create_image_dataset
from options.test_options import TestOptions
from utils.misc import sample_data, postprocess
import cv2
from criteria.lpips1 import util_of_lpips


def mae(image1, image2):
    mae = np.mean(abs(image1 - image2))
    return mae


is_cuda = torch.cuda.is_available()
if is_cuda:
    print('Cuda is available')
    cudnn.enable = True
    cudnn.benchmark = True

opts = TestOptions().parse

os.makedirs('{:s}'.format(opts.result_root), exist_ok=True)

# model & load model
generator = Generator(image_in_channels=3, edge_in_channels=2, out_channels=3)
if opts.pre_trained != '':
    generator.load_state_dict(torch.load(opts.pre_trained)['generator'])
else:
    print('Please provide pre-trained model!')

if is_cuda:
    generator = generator.cuda()

# dataset
image_dataset = create_image_dataset(opts)
image_data_loader = data.DataLoader(
    image_dataset,
    batch_size=opts.batch_size,
    shuffle=True,
    num_workers=opts.num_workers,
    drop_last=False
)
image_data_loader = sample_data(image_data_loader)

print('start test...')
with torch.no_grad():
    generator.eval()
    A = 0
    B = 0
    C = 0
    D = 0
    for _ in tqdm(range(opts.number_eval)):

        ground_truth, mask, edge, gray_image = next(image_data_loader)
        if is_cuda:
            ground_truth, mask, edge, gray_image = ground_truth.cuda(), mask.cuda(), edge.cuda(), gray_image.cuda()

        input_image, input_edge, input_gray_image = ground_truth * mask, edge * mask, gray_image * mask

        output, __, __ = generator(input_image, torch.cat((input_edge, input_gray_image), dim=1), mask)
        output_comp = ground_truth * mask + output * (1 - mask)

        # output_comp = postprocess(output_comp)
        output_comp1 = torch.cat([ground_truth, input_image, output_comp], dim=3)
        img1 = output_comp1.add(1).div(2).mul(255).clamp(0, 255).byte()
        img1 = img1[0].permute(1, 2, 0).data.cpu().numpy()
        output_comp1 = Image.fromarray(img1)

        ground_truth = ground_truth.add(1).div(2).mul(255).clamp(0, 255).byte()
        ground_truth = ground_truth[0].permute(1, 2, 0).data.cpu().numpy()
        ground_truth = Image.fromarray(ground_truth)
        output_comp = output_comp.add(1).div(2).mul(255).clamp(0, 255).byte()
        output_comp = output_comp[0].permute(1, 2, 0).data.cpu().numpy()
        output_comp = Image.fromarray(output_comp)
        output_comp.save(opts.result_root + '/{:05d}.png'.format(_))

        #psnr指标
        ground_truth_split, output_comp_split = ground_truth, output_comp
        MSE2=0
        for i in range(3):
            MSE1 = 0
            j = np.asarray(ground_truth_split.split()[i])
            k = np.asarray(output_comp_split.split()[i])
            MSE1 = mean_squared_error(j, k)
            MSE2 += MSE1
        MSE1 = MSE2 / 3

        psnr2  = 10 * np.log10((255 ** 2) / MSE1)

        # lpips指标
        LPIPS = util_of_lpips(net='alex').calc_lpips(np.asarray(ground_truth), np.asarray(output_comp)).squeeze().detach().numpy()

        output_comp= cv2.cvtColor(np.asarray(output_comp), cv2.COLOR_RGB2BGR)
        ground_truth = cv2.cvtColor(np.asarray(ground_truth), cv2.COLOR_RGB2BGR)
        SSIM = ssim(ground_truth , output_comp, win_size=3, multichannel=True)

        # MAE = mae(ground_truth, output_comp)
        A += psnr2
        B += SSIM
        # C += MAE

        D+=LPIPS
        fileName = 'note_KCA.txt'
        with open(fileName, 'a', encoding='utf-8') as file:
            file.write(f'{_:05d}.png,PSNR:{psnr2},SSIM:{SSIM},LPIPS:{LPIPS}\n')
            if _ == opts.number_eval - 1:
                file.write(
                    f'PSNR_mean:{A / opts.number_eval},SSIM_mean:{B / opts.number_eval},LPIPS_mean:{D / opts.number_eval}\n')
                file.write(opts.pre_trained+'\n')
        # print(PSNR, ' ', SSIM,'',MAE)

        output_comp1.save(opts.result_root + '/{:05d}.png'.format(_))
    file.close()