File size: 2,431 Bytes
3be620b |
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 |
import numpy as np
import tensorflow as tf
from scipy import linalg
from tensorflow.keras.applications.inception_v3 import InceptionV3, preprocess_input
from tqdm.auto import tqdm
inceptionv3 = InceptionV3(include_top=False, weights="imagenet", pooling="avg")
def resize_images(images, new_shape):
images = tf.image.resize(images, new_shape)
return images
def calculate_fid(real_embeddings, generated_embeddings):
# calculate mean and covariance statistics
mu1, sigma1 = real_embeddings.mean(axis=0), np.cov(real_embeddings, rowvar=False)
mu2, sigma2 = generated_embeddings.mean(axis=0), np.cov(
generated_embeddings, rowvar=False
)
# calculate sum squared difference between means
ssdiff = np.sum((mu1 - mu2) ** 2.0)
# calculate sqrt of product between cov
covmean = linalg.sqrtm(sigma1.dot(sigma2))
# check and correct imaginary numbers from sqrt
if np.iscomplexobj(covmean):
covmean = covmean.real
# calculate score
fid = ssdiff + np.trace(sigma1 + sigma2 - 2.0 * covmean)
return fid
def calculate_images_metrics(dataset, model, total_length):
fake_embeddings = []
real_embeddings = []
psnrs = []
ssims = []
for sample in tqdm(dataset, total=total_length):
generated = model(sample[0], training=False)[0]
generated, real = generated, sample[0]
real_resized = resize_images(real, (299, 299))
generated_resized = resize_images(generated, (299, 299))
real_activations = inceptionv3(real_resized, training=False)
generated_activations = inceptionv3(generated_resized, training=False)
fake_embeddings.append(generated_activations)
real_embeddings.append(real_activations)
fake_scaled = tf.cast(((generated * 0.5) + 1) * 255, tf.uint8)
real_scaled = tf.cast(((real * 0.5) + 1) * 255, tf.uint8)
psnrs.append(tf.image.psnr(fake_scaled, real_scaled, 255).numpy())
ssims.append(tf.image.ssim(fake_scaled, real_scaled, 255).numpy())
fid = calculate_fid(
tf.concat(fake_embeddings, axis=0).numpy(),
tf.concat(real_embeddings, axis=0).numpy(),
)
# kid = calculate_kid(
# tf.concat(fake_embeddings, axis=0).numpy(),
# tf.concat(real_embeddings, axis=0).numpy(),
# )
psnr = np.array(psnrs).mean()
ssim = np.array(ssims).mean()
return {"fid": fid, "ssim": ssim, "psnr": psnr}
|