Spaces:
Running
Running
| import argparse | |
| import os | |
| import numpy as np | |
| from PIL import Image | |
| from scipy.stats import entropy | |
| from skimage.metrics import structural_similarity as ssim | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import torchvision.transforms as Transforms | |
| from torchvision.models.inception import inception_v3 | |
| import eval_models as models | |
| def get_opt(): | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument('--evaluation', default='LPIPS') | |
| parser.add_argument('--predict_dir', default='./result/bg_ver1/output/') | |
| parser.add_argument('--ground_truth_dir', default='./data/zalando-hd-resize/test/image') | |
| parser.add_argument('--resolution', type=int, default=1024) | |
| opt = parser.parse_args() | |
| return opt | |
| def Evaluation(opt, pred_list, gt_list): | |
| T1 = Transforms.ToTensor() | |
| T2 = Transforms.Compose([Transforms.Resize((128, 128)), | |
| Transforms.ToTensor(), | |
| Transforms.Normalize(mean=(0.5, 0.5, 0.5), | |
| std=(0.5, 0.5, 0.5))]) | |
| T3 = Transforms.Compose([Transforms.Resize((299, 299)), | |
| Transforms.ToTensor(), | |
| Transforms.Normalize(mean=(0.5, 0.5, 0.5), | |
| std=(0.5, 0.5, 0.5))]) | |
| splits = 1 # Hyper-parameter for IS score | |
| model = models.PerceptualLoss(model='net-lin',net='alex',use_gpu=True) | |
| model.eval() | |
| inception_model = inception_v3(pretrained=True, transform_input=False).type(torch.cuda.FloatTensor) | |
| inception_model.eval() | |
| avg_ssim, avg_mse, avg_distance = 0.0, 0.0, 0.0 | |
| preds = np.zeros((len(gt_list), 1000)) | |
| lpips_list = [] | |
| with torch.no_grad(): | |
| print("Calculate SSIM, MSE, LPIPS...") | |
| for i, img_pred in enumerate(pred_list): | |
| img = img_pred.split('_')[0] + '_00.jpg' | |
| # Calculate SSIM | |
| gt_img = Image.open(os.path.join(opt.ground_truth_dir, img)) | |
| if not opt.resolution == 1024: | |
| if opt.resolution == 512: | |
| gt_img = gt_img.resize((384,512), Image.BILINEAR) | |
| elif opt.resolution == 256: | |
| gt_img = gt_img.resize((192,256), Image.BILINEAR) | |
| else: | |
| raise NotImplementedError | |
| gt_np = np.asarray(gt_img.convert('L')) | |
| pred_img = Image.open(os.path.join(opt.predict_dir, img_pred)) | |
| assert gt_img.size == pred_img.size, f"{gt_img.size} vs {pred_img.size}" | |
| pred_np = np.asarray(pred_img.convert('L')) | |
| avg_ssim += ssim(gt_np, pred_np, data_range=255, gaussian_weights=True, use_sample_covariance=False) | |
| # Calculate LPIPS | |
| gt_img_LPIPS = T2(gt_img).unsqueeze(0).cuda() | |
| pred_img_LPIPS = T2(pred_img).unsqueeze(0).cuda() | |
| lpips_list.append((img_pred, model.forward(gt_img_LPIPS, pred_img_LPIPS).item())) | |
| avg_distance += lpips_list[-1][1] | |
| # Calculate Inception model prediction | |
| pred_img_IS = T3(pred_img).unsqueeze(0).cuda() | |
| preds[i] = F.softmax(inception_model(pred_img_IS)).data.cpu().numpy() | |
| gt_img_MSE = T1(gt_img).unsqueeze(0).cuda() | |
| pred_img_MSE = T1(pred_img).unsqueeze(0).cuda() | |
| avg_mse += F.mse_loss(gt_img_MSE, pred_img_MSE) | |
| print(f"step: {i+1} evaluation... lpips:{lpips_list[-1][1]}") | |
| avg_ssim /= len(gt_list) | |
| avg_mse = avg_mse / len(gt_list) | |
| avg_distance = avg_distance / len(gt_list) | |
| # Calculate Inception Score | |
| split_scores = [] # Now compute the mean kl-divergence | |
| lpips_list.sort(key=lambda x: x[1], reverse=True) | |
| for name, score in lpips_list: | |
| f = open(os.path.join(opt.predict_dir, 'lpips.txt'), 'a') | |
| f.write(f"{name} {score}\n") | |
| f.close() | |
| print("Calculate Inception Score...") | |
| for k in range(splits): | |
| part = preds[k * (len(gt_list) // splits): (k+1) * (len(gt_list) // splits), :] | |
| py = np.mean(part, axis=0) | |
| scores = [] | |
| for i in range(part.shape[0]): | |
| pyx = part[i, :] | |
| scores.append(entropy(pyx, py)) | |
| split_scores.append(np.exp(np.mean(scores))) | |
| IS_mean, IS_std = np.mean(split_scores), np.std(split_scores) | |
| f = open(os.path.join(opt.predict_dir, 'eval.txt'), 'a') | |
| f.write(f"SSIM : {avg_ssim} / MSE : {avg_mse} / LPIPS : {avg_distance}\n") | |
| f.write(f"IS_mean : {IS_mean} / IS_std : {IS_std}\n") | |
| f.close() | |
| return avg_ssim, avg_mse, avg_distance, IS_mean, IS_std | |
| def main(): | |
| opt = get_opt() | |
| # Output과 Ground Truth Data | |
| pred_list = os.listdir(opt.predict_dir) | |
| gt_list = os.listdir(opt.ground_truth_dir) | |
| pred_list.sort() | |
| gt_list.sort() | |
| avg_ssim, avg_mse, avg_distance, IS_mean, IS_std = Evaluation(opt, pred_list, gt_list) | |
| print("SSIM : %f / MSE : %f / LPIPS : %f" % (avg_ssim, avg_mse, avg_distance)) | |
| print("IS_mean : %f / IS_std : %f" % (IS_mean, IS_std)) | |
| if __name__ == '__main__': | |
| main() | |