631 lines
20 KiB
Python
631 lines
20 KiB
Python
|
import argparse
|
||
|
|
import glob
|
||
|
|
import os
|
||
|
|
from tqdm import tqdm
|
||
|
|
from collections import namedtuple
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import torch
|
||
|
|
import torchvision.transforms as transforms
|
||
|
|
from torchvision import models
|
||
|
|
from PIL import Image
|
||
|
|
|
||
|
|
from ldm.modules.evaluate.ssim import ssim
|
||
|
|
|
||
|
|
|
||
|
|
transform = transforms.Compose([transforms.ToTensor()])
|
||
|
|
|
||
|
|
def normalize_tensor(in_feat, eps=1e-10):
|
||
|
|
norm_factor = torch.sqrt(torch.sum(in_feat ** 2, dim=1)).view(
|
||
|
|
in_feat.size()[0], 1, in_feat.size()[2], in_feat.size()[3]
|
||
|
|
)
|
||
|
|
return in_feat / (norm_factor.expand_as(in_feat) + eps)
|
||
|
|
|
||
|
|
|
||
|
|
def cos_sim(in0, in1):
|
||
|
|
in0_norm = normalize_tensor(in0)
|
||
|
|
in1_norm = normalize_tensor(in1)
|
||
|
|
N = in0.size()[0]
|
||
|
|
X = in0.size()[2]
|
||
|
|
Y = in0.size()[3]
|
||
|
|
|
||
|
|
return torch.mean(
|
||
|
|
torch.mean(
|
||
|
|
torch.sum(in0_norm * in1_norm, dim=1).view(N, 1, X, Y), dim=2
|
||
|
|
).view(N, 1, 1, Y),
|
||
|
|
dim=3,
|
||
|
|
).view(N)
|
||
|
|
|
||
|
|
|
||
|
|
class squeezenet(torch.nn.Module):
|
||
|
|
def __init__(self, requires_grad=False, pretrained=True):
|
||
|
|
super(squeezenet, self).__init__()
|
||
|
|
pretrained_features = models.squeezenet1_1(
|
||
|
|
pretrained=pretrained
|
||
|
|
).features
|
||
|
|
self.slice1 = torch.nn.Sequential()
|
||
|
|
self.slice2 = torch.nn.Sequential()
|
||
|
|
self.slice3 = torch.nn.Sequential()
|
||
|
|
self.slice4 = torch.nn.Sequential()
|
||
|
|
self.slice5 = torch.nn.Sequential()
|
||
|
|
self.slice6 = torch.nn.Sequential()
|
||
|
|
self.slice7 = torch.nn.Sequential()
|
||
|
|
self.N_slices = 7
|
||
|
|
for x in range(2):
|
||
|
|
self.slice1.add_module(str(x), pretrained_features[x])
|
||
|
|
for x in range(2, 5):
|
||
|
|
self.slice2.add_module(str(x), pretrained_features[x])
|
||
|
|
for x in range(5, 8):
|
||
|
|
self.slice3.add_module(str(x), pretrained_features[x])
|
||
|
|
for x in range(8, 10):
|
||
|
|
self.slice4.add_module(str(x), pretrained_features[x])
|
||
|
|
for x in range(10, 11):
|
||
|
|
self.slice5.add_module(str(x), pretrained_features[x])
|
||
|
|
for x in range(11, 12):
|
||
|
|
self.slice6.add_module(str(x), pretrained_features[x])
|
||
|
|
for x in range(12, 13):
|
||
|
|
self.slice7.add_module(str(x), pretrained_features[x])
|
||
|
|
if not requires_grad:
|
||
|
|
for param in self.parameters():
|
||
|
|
param.requires_grad = False
|
||
|
|
|
||
|
|
def forward(self, X):
|
||
|
|
h = self.slice1(X)
|
||
|
|
h_relu1 = h
|
||
|
|
h = self.slice2(h)
|
||
|
|
h_relu2 = h
|
||
|
|
h = self.slice3(h)
|
||
|
|
h_relu3 = h
|
||
|
|
h = self.slice4(h)
|
||
|
|
h_relu4 = h
|
||
|
|
h = self.slice5(h)
|
||
|
|
h_relu5 = h
|
||
|
|
h = self.slice6(h)
|
||
|
|
h_relu6 = h
|
||
|
|
h = self.slice7(h)
|
||
|
|
h_relu7 = h
|
||
|
|
vgg_outputs = namedtuple(
|
||
|
|
"SqueezeOutputs",
|
||
|
|
["relu1", "relu2", "relu3", "relu4", "relu5", "relu6", "relu7"],
|
||
|
|
)
|
||
|
|
out = vgg_outputs(
|
||
|
|
h_relu1, h_relu2, h_relu3, h_relu4, h_relu5, h_relu6, h_relu7
|
||
|
|
)
|
||
|
|
|
||
|
|
return out
|
||
|
|
|
||
|
|
|
||
|
|
class alexnet(torch.nn.Module):
|
||
|
|
def __init__(self, requires_grad=False, pretrained=True):
|
||
|
|
super(alexnet, self).__init__()
|
||
|
|
alexnet_pretrained_features = models.alexnet(
|
||
|
|
pretrained=pretrained
|
||
|
|
).features
|
||
|
|
self.slice1 = torch.nn.Sequential()
|
||
|
|
self.slice2 = torch.nn.Sequential()
|
||
|
|
self.slice3 = torch.nn.Sequential()
|
||
|
|
self.slice4 = torch.nn.Sequential()
|
||
|
|
self.slice5 = torch.nn.Sequential()
|
||
|
|
self.N_slices = 5
|
||
|
|
for x in range(2):
|
||
|
|
self.slice1.add_module(str(x), alexnet_pretrained_features[x])
|
||
|
|
for x in range(2, 5):
|
||
|
|
self.slice2.add_module(str(x), alexnet_pretrained_features[x])
|
||
|
|
for x in range(5, 8):
|
||
|
|
self.slice3.add_module(str(x), alexnet_pretrained_features[x])
|
||
|
|
for x in range(8, 10):
|
||
|
|
self.slice4.add_module(str(x), alexnet_pretrained_features[x])
|
||
|
|
for x in range(10, 12):
|
||
|
|
self.slice5.add_module(str(x), alexnet_pretrained_features[x])
|
||
|
|
if not requires_grad:
|
||
|
|
for param in self.parameters():
|
||
|
|
param.requires_grad = False
|
||
|
|
|
||
|
|
def forward(self, X):
|
||
|
|
h = self.slice1(X)
|
||
|
|
h_relu1 = h
|
||
|
|
h = self.slice2(h)
|
||
|
|
h_relu2 = h
|
||
|
|
h = self.slice3(h)
|
||
|
|
h_relu3 = h
|
||
|
|
h = self.slice4(h)
|
||
|
|
h_relu4 = h
|
||
|
|
h = self.slice5(h)
|
||
|
|
h_relu5 = h
|
||
|
|
alexnet_outputs = namedtuple(
|
||
|
|
"AlexnetOutputs", ["relu1", "relu2", "relu3", "relu4", "relu5"]
|
||
|
|
)
|
||
|
|
out = alexnet_outputs(h_relu1, h_relu2, h_relu3, h_relu4, h_relu5)
|
||
|
|
|
||
|
|
return out
|
||
|
|
|
||
|
|
|
||
|
|
class vgg16(torch.nn.Module):
|
||
|
|
def __init__(self, requires_grad=False, pretrained=True):
|
||
|
|
super(vgg16, self).__init__()
|
||
|
|
vgg_pretrained_features = models.vgg16(pretrained=pretrained).features
|
||
|
|
self.slice1 = torch.nn.Sequential()
|
||
|
|
self.slice2 = torch.nn.Sequential()
|
||
|
|
self.slice3 = torch.nn.Sequential()
|
||
|
|
self.slice4 = torch.nn.Sequential()
|
||
|
|
self.slice5 = torch.nn.Sequential()
|
||
|
|
self.N_slices = 5
|
||
|
|
for x in range(4):
|
||
|
|
self.slice1.add_module(str(x), vgg_pretrained_features[x])
|
||
|
|
for x in range(4, 9):
|
||
|
|
self.slice2.add_module(str(x), vgg_pretrained_features[x])
|
||
|
|
for x in range(9, 16):
|
||
|
|
self.slice3.add_module(str(x), vgg_pretrained_features[x])
|
||
|
|
for x in range(16, 23):
|
||
|
|
self.slice4.add_module(str(x), vgg_pretrained_features[x])
|
||
|
|
for x in range(23, 30):
|
||
|
|
self.slice5.add_module(str(x), vgg_pretrained_features[x])
|
||
|
|
if not requires_grad:
|
||
|
|
for param in self.parameters():
|
||
|
|
param.requires_grad = False
|
||
|
|
|
||
|
|
def forward(self, X):
|
||
|
|
h = self.slice1(X)
|
||
|
|
h_relu1_2 = h
|
||
|
|
h = self.slice2(h)
|
||
|
|
h_relu2_2 = h
|
||
|
|
h = self.slice3(h)
|
||
|
|
h_relu3_3 = h
|
||
|
|
h = self.slice4(h)
|
||
|
|
h_relu4_3 = h
|
||
|
|
h = self.slice5(h)
|
||
|
|
h_relu5_3 = h
|
||
|
|
vgg_outputs = namedtuple(
|
||
|
|
"VggOutputs",
|
||
|
|
["relu1_2", "relu2_2", "relu3_3", "relu4_3", "relu5_3"],
|
||
|
|
)
|
||
|
|
out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3, h_relu5_3)
|
||
|
|
|
||
|
|
return out
|
||
|
|
|
||
|
|
|
||
|
|
class resnet(torch.nn.Module):
|
||
|
|
def __init__(self, requires_grad=False, pretrained=True, num=18):
|
||
|
|
super(resnet, self).__init__()
|
||
|
|
if num == 18:
|
||
|
|
self.net = models.resnet18(pretrained=pretrained)
|
||
|
|
elif num == 34:
|
||
|
|
self.net = models.resnet34(pretrained=pretrained)
|
||
|
|
elif num == 50:
|
||
|
|
self.net = models.resnet50(pretrained=pretrained)
|
||
|
|
elif num == 101:
|
||
|
|
self.net = models.resnet101(pretrained=pretrained)
|
||
|
|
elif num == 152:
|
||
|
|
self.net = models.resnet152(pretrained=pretrained)
|
||
|
|
self.N_slices = 5
|
||
|
|
|
||
|
|
self.conv1 = self.net.conv1
|
||
|
|
self.bn1 = self.net.bn1
|
||
|
|
self.relu = self.net.relu
|
||
|
|
self.maxpool = self.net.maxpool
|
||
|
|
self.layer1 = self.net.layer1
|
||
|
|
self.layer2 = self.net.layer2
|
||
|
|
self.layer3 = self.net.layer3
|
||
|
|
self.layer4 = self.net.layer4
|
||
|
|
|
||
|
|
def forward(self, X):
|
||
|
|
h = self.conv1(X)
|
||
|
|
h = self.bn1(h)
|
||
|
|
h = self.relu(h)
|
||
|
|
h_relu1 = h
|
||
|
|
h = self.maxpool(h)
|
||
|
|
h = self.layer1(h)
|
||
|
|
h_conv2 = h
|
||
|
|
h = self.layer2(h)
|
||
|
|
h_conv3 = h
|
||
|
|
h = self.layer3(h)
|
||
|
|
h_conv4 = h
|
||
|
|
h = self.layer4(h)
|
||
|
|
h_conv5 = h
|
||
|
|
|
||
|
|
outputs = namedtuple(
|
||
|
|
"Outputs", ["relu1", "conv2", "conv3", "conv4", "conv5"]
|
||
|
|
)
|
||
|
|
out = outputs(h_relu1, h_conv2, h_conv3, h_conv4, h_conv5)
|
||
|
|
|
||
|
|
return out
|
||
|
|
|
||
|
|
# Off-the-shelf deep network
|
||
|
|
class PNet(torch.nn.Module):
|
||
|
|
"""Pre-trained network with all channels equally weighted by default"""
|
||
|
|
|
||
|
|
def __init__(self, pnet_type="vgg", pnet_rand=False, use_gpu=True):
|
||
|
|
super(PNet, self).__init__()
|
||
|
|
|
||
|
|
self.use_gpu = use_gpu
|
||
|
|
|
||
|
|
self.pnet_type = pnet_type
|
||
|
|
self.pnet_rand = pnet_rand
|
||
|
|
|
||
|
|
self.shift = torch.Tensor([-0.030, -0.088, -0.188]).view(1, 3, 1, 1)
|
||
|
|
self.scale = torch.Tensor([0.458, 0.448, 0.450]).view(1, 3, 1, 1)
|
||
|
|
|
||
|
|
if self.pnet_type in ["vgg", "vgg16"]:
|
||
|
|
self.net = vgg16(pretrained=not self.pnet_rand, requires_grad=False)
|
||
|
|
elif self.pnet_type == "alex":
|
||
|
|
self.net = alexnet(
|
||
|
|
pretrained=not self.pnet_rand, requires_grad=False
|
||
|
|
)
|
||
|
|
elif self.pnet_type[:-2] == "resnet":
|
||
|
|
self.net = resnet(
|
||
|
|
pretrained=not self.pnet_rand,
|
||
|
|
requires_grad=False,
|
||
|
|
num=int(self.pnet_type[-2:]),
|
||
|
|
)
|
||
|
|
elif self.pnet_type == "squeeze":
|
||
|
|
self.net = squeezenet(
|
||
|
|
pretrained=not self.pnet_rand, requires_grad=False
|
||
|
|
)
|
||
|
|
|
||
|
|
self.L = self.net.N_slices
|
||
|
|
|
||
|
|
if use_gpu:
|
||
|
|
self.net.cuda()
|
||
|
|
self.shift = self.shift.cuda()
|
||
|
|
self.scale = self.scale.cuda()
|
||
|
|
|
||
|
|
def forward(self, in0, in1, retPerLayer=False):
|
||
|
|
in0_sc = (in0 - self.shift.expand_as(in0)) / self.scale.expand_as(in0)
|
||
|
|
in1_sc = (in1 - self.shift.expand_as(in0)) / self.scale.expand_as(in0)
|
||
|
|
|
||
|
|
outs0 = self.net.forward(in0_sc)
|
||
|
|
outs1 = self.net.forward(in1_sc)
|
||
|
|
|
||
|
|
if retPerLayer:
|
||
|
|
all_scores = []
|
||
|
|
for (kk, out0) in enumerate(outs0):
|
||
|
|
cur_score = 1.0 - cos_sim(outs0[kk], outs1[kk])
|
||
|
|
if kk == 0:
|
||
|
|
val = 1.0 * cur_score
|
||
|
|
else:
|
||
|
|
val = val + cur_score
|
||
|
|
if retPerLayer:
|
||
|
|
all_scores += [cur_score]
|
||
|
|
|
||
|
|
if retPerLayer:
|
||
|
|
return (val, all_scores)
|
||
|
|
else:
|
||
|
|
return val
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
# The SSIM metric
|
||
|
|
def ssim_metric(img1, img2, mask=None):
|
||
|
|
return ssim(img1, img2, mask=mask, size_average=False)
|
||
|
|
|
||
|
|
|
||
|
|
# The PSNR metric
|
||
|
|
def psnr(img1, img2, mask=None,reshape=False):
|
||
|
|
b = img1.size(0)
|
||
|
|
if not (mask is None):
|
||
|
|
b = img1.size(0)
|
||
|
|
mse_err = (img1 - img2).pow(2) * mask
|
||
|
|
if reshape:
|
||
|
|
mse_err = mse_err.reshape(b, -1).sum(dim=1) / (
|
||
|
|
3 * mask.reshape(b, -1).sum(dim=1).clamp(min=1)
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
mse_err = mse_err.view(b, -1).sum(dim=1) / (
|
||
|
|
3 * mask.view(b, -1).sum(dim=1).clamp(min=1)
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
if reshape:
|
||
|
|
mse_err = (img1 - img2).pow(2).reshape(b, -1).mean(dim=1)
|
||
|
|
else:
|
||
|
|
mse_err = (img1 - img2).pow(2).view(b, -1).mean(dim=1)
|
||
|
|
|
||
|
|
psnr = 10 * (1 / mse_err).log10()
|
||
|
|
return psnr
|
||
|
|
|
||
|
|
|
||
|
|
# The perceptual similarity metric
|
||
|
|
def perceptual_sim(img1, img2, vgg16):
|
||
|
|
# First extract features
|
||
|
|
dist = vgg16(img1 * 2 - 1, img2 * 2 - 1)
|
||
|
|
|
||
|
|
return dist
|
||
|
|
|
||
|
|
def load_img(img_name, size=None):
|
||
|
|
try:
|
||
|
|
img = Image.open(img_name)
|
||
|
|
|
||
|
|
if type(size) == int:
|
||
|
|
img = img.resize((size, size))
|
||
|
|
elif size is not None:
|
||
|
|
img = img.resize((size[1], size[0]))
|
||
|
|
|
||
|
|
img = transform(img).cuda()
|
||
|
|
img = img.unsqueeze(0)
|
||
|
|
except Exception as e:
|
||
|
|
print("Failed at loading %s " % img_name)
|
||
|
|
print(e)
|
||
|
|
img = torch.zeros(1, 3, 256, 256).cuda()
|
||
|
|
raise
|
||
|
|
return img
|
||
|
|
|
||
|
|
|
||
|
|
def compute_perceptual_similarity(folder, pred_img, tgt_img, take_every_other):
|
||
|
|
|
||
|
|
# Load VGG16 for feature similarity
|
||
|
|
vgg16 = PNet().to("cuda")
|
||
|
|
vgg16.eval()
|
||
|
|
vgg16.cuda()
|
||
|
|
|
||
|
|
values_percsim = []
|
||
|
|
values_ssim = []
|
||
|
|
values_psnr = []
|
||
|
|
folders = os.listdir(folder)
|
||
|
|
for i, f in tqdm(enumerate(sorted(folders))):
|
||
|
|
pred_imgs = glob.glob(folder + f + "/" + pred_img)
|
||
|
|
tgt_imgs = glob.glob(folder + f + "/" + tgt_img)
|
||
|
|
assert len(tgt_imgs) == 1
|
||
|
|
|
||
|
|
perc_sim = 10000
|
||
|
|
ssim_sim = -10
|
||
|
|
psnr_sim = -10
|
||
|
|
for p_img in pred_imgs:
|
||
|
|
t_img = load_img(tgt_imgs[0])
|
||
|
|
p_img = load_img(p_img, size=t_img.shape[2:])
|
||
|
|
t_perc_sim = perceptual_sim(p_img, t_img, vgg16).item()
|
||
|
|
perc_sim = min(perc_sim, t_perc_sim)
|
||
|
|
|
||
|
|
ssim_sim = max(ssim_sim, ssim_metric(p_img, t_img).item())
|
||
|
|
psnr_sim = max(psnr_sim, psnr(p_img, t_img).item())
|
||
|
|
|
||
|
|
values_percsim += [perc_sim]
|
||
|
|
values_ssim += [ssim_sim]
|
||
|
|
values_psnr += [psnr_sim]
|
||
|
|
|
||
|
|
if take_every_other:
|
||
|
|
n_valuespercsim = []
|
||
|
|
n_valuesssim = []
|
||
|
|
n_valuespsnr = []
|
||
|
|
for i in range(0, len(values_percsim) // 2):
|
||
|
|
n_valuespercsim += [
|
||
|
|
min(values_percsim[2 * i], values_percsim[2 * i + 1])
|
||
|
|
]
|
||
|
|
n_valuespsnr += [max(values_psnr[2 * i], values_psnr[2 * i + 1])]
|
||
|
|
n_valuesssim += [max(values_ssim[2 * i], values_ssim[2 * i + 1])]
|
||
|
|
|
||
|
|
values_percsim = n_valuespercsim
|
||
|
|
values_ssim = n_valuesssim
|
||
|
|
values_psnr = n_valuespsnr
|
||
|
|
|
||
|
|
avg_percsim = np.mean(np.array(values_percsim))
|
||
|
|
std_percsim = np.std(np.array(values_percsim))
|
||
|
|
|
||
|
|
avg_psnr = np.mean(np.array(values_psnr))
|
||
|
|
std_psnr = np.std(np.array(values_psnr))
|
||
|
|
|
||
|
|
avg_ssim = np.mean(np.array(values_ssim))
|
||
|
|
std_ssim = np.std(np.array(values_ssim))
|
||
|
|
|
||
|
|
return {
|
||
|
|
"Perceptual similarity": (avg_percsim, std_percsim),
|
||
|
|
"PSNR": (avg_psnr, std_psnr),
|
||
|
|
"SSIM": (avg_ssim, std_ssim),
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def compute_perceptual_similarity_from_list(pred_imgs_list, tgt_imgs_list,
|
||
|
|
take_every_other,
|
||
|
|
simple_format=True):
|
||
|
|
|
||
|
|
# Load VGG16 for feature similarity
|
||
|
|
vgg16 = PNet().to("cuda")
|
||
|
|
vgg16.eval()
|
||
|
|
vgg16.cuda()
|
||
|
|
|
||
|
|
values_percsim = []
|
||
|
|
values_ssim = []
|
||
|
|
values_psnr = []
|
||
|
|
equal_count = 0
|
||
|
|
ambig_count = 0
|
||
|
|
for i, tgt_img in enumerate(tqdm(tgt_imgs_list)):
|
||
|
|
pred_imgs = pred_imgs_list[i]
|
||
|
|
tgt_imgs = [tgt_img]
|
||
|
|
assert len(tgt_imgs) == 1
|
||
|
|
|
||
|
|
if type(pred_imgs) != list:
|
||
|
|
pred_imgs = [pred_imgs]
|
||
|
|
|
||
|
|
perc_sim = 10000
|
||
|
|
ssim_sim = -10
|
||
|
|
psnr_sim = -10
|
||
|
|
assert len(pred_imgs)>0
|
||
|
|
for p_img in pred_imgs:
|
||
|
|
t_img = load_img(tgt_imgs[0])
|
||
|
|
p_img = load_img(p_img, size=t_img.shape[2:])
|
||
|
|
t_perc_sim = perceptual_sim(p_img, t_img, vgg16).item()
|
||
|
|
perc_sim = min(perc_sim, t_perc_sim)
|
||
|
|
|
||
|
|
ssim_sim = max(ssim_sim, ssim_metric(p_img, t_img).item())
|
||
|
|
psnr_sim = max(psnr_sim, psnr(p_img, t_img).item())
|
||
|
|
|
||
|
|
values_percsim += [perc_sim]
|
||
|
|
values_ssim += [ssim_sim]
|
||
|
|
if psnr_sim != np.float("inf"):
|
||
|
|
values_psnr += [psnr_sim]
|
||
|
|
else:
|
||
|
|
if torch.allclose(p_img, t_img):
|
||
|
|
equal_count += 1
|
||
|
|
print("{} equal src and wrp images.".format(equal_count))
|
||
|
|
else:
|
||
|
|
ambig_count += 1
|
||
|
|
print("{} ambiguous src and wrp images.".format(ambig_count))
|
||
|
|
|
||
|
|
if take_every_other:
|
||
|
|
n_valuespercsim = []
|
||
|
|
n_valuesssim = []
|
||
|
|
n_valuespsnr = []
|
||
|
|
for i in range(0, len(values_percsim) // 2):
|
||
|
|
n_valuespercsim += [
|
||
|
|
min(values_percsim[2 * i], values_percsim[2 * i + 1])
|
||
|
|
]
|
||
|
|
n_valuespsnr += [max(values_psnr[2 * i], values_psnr[2 * i + 1])]
|
||
|
|
n_valuesssim += [max(values_ssim[2 * i], values_ssim[2 * i + 1])]
|
||
|
|
|
||
|
|
values_percsim = n_valuespercsim
|
||
|
|
values_ssim = n_valuesssim
|
||
|
|
values_psnr = n_valuespsnr
|
||
|
|
|
||
|
|
avg_percsim = np.mean(np.array(values_percsim))
|
||
|
|
std_percsim = np.std(np.array(values_percsim))
|
||
|
|
|
||
|
|
avg_psnr = np.mean(np.array(values_psnr))
|
||
|
|
std_psnr = np.std(np.array(values_psnr))
|
||
|
|
|
||
|
|
avg_ssim = np.mean(np.array(values_ssim))
|
||
|
|
std_ssim = np.std(np.array(values_ssim))
|
||
|
|
|
||
|
|
if simple_format:
|
||
|
|
# just to make yaml formatting readable
|
||
|
|
return {
|
||
|
|
"Perceptual similarity": [float(avg_percsim), float(std_percsim)],
|
||
|
|
"PSNR": [float(avg_psnr), float(std_psnr)],
|
||
|
|
"SSIM": [float(avg_ssim), float(std_ssim)],
|
||
|
|
}
|
||
|
|
else:
|
||
|
|
return {
|
||
|
|
"Perceptual similarity": (avg_percsim, std_percsim),
|
||
|
|
"PSNR": (avg_psnr, std_psnr),
|
||
|
|
"SSIM": (avg_ssim, std_ssim),
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def compute_perceptual_similarity_from_list_topk(pred_imgs_list, tgt_imgs_list,
|
||
|
|
take_every_other, resize=False):
|
||
|
|
|
||
|
|
# Load VGG16 for feature similarity
|
||
|
|
vgg16 = PNet().to("cuda")
|
||
|
|
vgg16.eval()
|
||
|
|
vgg16.cuda()
|
||
|
|
|
||
|
|
values_percsim = []
|
||
|
|
values_ssim = []
|
||
|
|
values_psnr = []
|
||
|
|
individual_percsim = []
|
||
|
|
individual_ssim = []
|
||
|
|
individual_psnr = []
|
||
|
|
for i, tgt_img in enumerate(tqdm(tgt_imgs_list)):
|
||
|
|
pred_imgs = pred_imgs_list[i]
|
||
|
|
tgt_imgs = [tgt_img]
|
||
|
|
assert len(tgt_imgs) == 1
|
||
|
|
|
||
|
|
if type(pred_imgs) != list:
|
||
|
|
assert False
|
||
|
|
pred_imgs = [pred_imgs]
|
||
|
|
|
||
|
|
perc_sim = 10000
|
||
|
|
ssim_sim = -10
|
||
|
|
psnr_sim = -10
|
||
|
|
sample_percsim = list()
|
||
|
|
sample_ssim = list()
|
||
|
|
sample_psnr = list()
|
||
|
|
for p_img in pred_imgs:
|
||
|
|
if resize:
|
||
|
|
t_img = load_img(tgt_imgs[0], size=(256,256))
|
||
|
|
else:
|
||
|
|
t_img = load_img(tgt_imgs[0])
|
||
|
|
p_img = load_img(p_img, size=t_img.shape[2:])
|
||
|
|
|
||
|
|
t_perc_sim = perceptual_sim(p_img, t_img, vgg16).item()
|
||
|
|
sample_percsim.append(t_perc_sim)
|
||
|
|
perc_sim = min(perc_sim, t_perc_sim)
|
||
|
|
|
||
|
|
t_ssim = ssim_metric(p_img, t_img).item()
|
||
|
|
sample_ssim.append(t_ssim)
|
||
|
|
ssim_sim = max(ssim_sim, t_ssim)
|
||
|
|
|
||
|
|
t_psnr = psnr(p_img, t_img).item()
|
||
|
|
sample_psnr.append(t_psnr)
|
||
|
|
psnr_sim = max(psnr_sim, t_psnr)
|
||
|
|
|
||
|
|
values_percsim += [perc_sim]
|
||
|
|
values_ssim += [ssim_sim]
|
||
|
|
values_psnr += [psnr_sim]
|
||
|
|
individual_percsim.append(sample_percsim)
|
||
|
|
individual_ssim.append(sample_ssim)
|
||
|
|
individual_psnr.append(sample_psnr)
|
||
|
|
|
||
|
|
if take_every_other:
|
||
|
|
assert False, "Do this later, after specifying topk to get proper results"
|
||
|
|
n_valuespercsim = []
|
||
|
|
n_valuesssim = []
|
||
|
|
n_valuespsnr = []
|
||
|
|
for i in range(0, len(values_percsim) // 2):
|
||
|
|
n_valuespercsim += [
|
||
|
|
min(values_percsim[2 * i], values_percsim[2 * i + 1])
|
||
|
|
]
|
||
|
|
n_valuespsnr += [max(values_psnr[2 * i], values_psnr[2 * i + 1])]
|
||
|
|
n_valuesssim += [max(values_ssim[2 * i], values_ssim[2 * i + 1])]
|
||
|
|
|
||
|
|
values_percsim = n_valuespercsim
|
||
|
|
values_ssim = n_valuesssim
|
||
|
|
values_psnr = n_valuespsnr
|
||
|
|
|
||
|
|
avg_percsim = np.mean(np.array(values_percsim))
|
||
|
|
std_percsim = np.std(np.array(values_percsim))
|
||
|
|
|
||
|
|
avg_psnr = np.mean(np.array(values_psnr))
|
||
|
|
std_psnr = np.std(np.array(values_psnr))
|
||
|
|
|
||
|
|
avg_ssim = np.mean(np.array(values_ssim))
|
||
|
|
std_ssim = np.std(np.array(values_ssim))
|
||
|
|
|
||
|
|
individual_percsim = np.array(individual_percsim)
|
||
|
|
individual_psnr = np.array(individual_psnr)
|
||
|
|
individual_ssim = np.array(individual_ssim)
|
||
|
|
|
||
|
|
return {
|
||
|
|
"avg_of_best": {
|
||
|
|
"Perceptual similarity": [float(avg_percsim), float(std_percsim)],
|
||
|
|
"PSNR": [float(avg_psnr), float(std_psnr)],
|
||
|
|
"SSIM": [float(avg_ssim), float(std_ssim)],
|
||
|
|
},
|
||
|
|
"individual": {
|
||
|
|
"PSIM": individual_percsim,
|
||
|
|
"PSNR": individual_psnr,
|
||
|
|
"SSIM": individual_ssim,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
args = argparse.ArgumentParser()
|
||
|
|
args.add_argument("--folder", type=str, default="")
|
||
|
|
args.add_argument("--pred_image", type=str, default="")
|
||
|
|
args.add_argument("--target_image", type=str, default="")
|
||
|
|
args.add_argument("--take_every_other", action="store_true", default=False)
|
||
|
|
args.add_argument("--output_file", type=str, default="")
|
||
|
|
|
||
|
|
opts = args.parse_args()
|
||
|
|
|
||
|
|
folder = opts.folder
|
||
|
|
pred_img = opts.pred_image
|
||
|
|
tgt_img = opts.target_image
|
||
|
|
|
||
|
|
results = compute_perceptual_similarity(
|
||
|
|
folder, pred_img, tgt_img, opts.take_every_other
|
||
|
|
)
|
||
|
|
|
||
|
|
f = open(opts.output_file, 'w')
|
||
|
|
for key in results:
|
||
|
|
print("%s for %s: \n" % (key, opts.folder))
|
||
|
|
print(
|
||
|
|
"\t {:0.4f} | {:0.4f} \n".format(results[key][0], results[key][1])
|
||
|
|
)
|
||
|
|
|
||
|
|
f.write("%s for %s: \n" % (key, opts.folder))
|
||
|
|
f.write(
|
||
|
|
"\t {:0.4f} | {:0.4f} \n".format(results[key][0], results[key][1])
|
||
|
|
)
|
||
|
|
|
||
|
|
f.close()
|