Spaces:
Running
Running
File size: 5,157 Bytes
8a6df40 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 |
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()
|