Spaces:
Running
Running
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| from scipy import linalg | |
| from tqdm import tqdm | |
| from basicsr.archs.inception import InceptionV3 | |
| def load_patched_inception_v3(device='cuda', resize_input=True, normalize_input=False): | |
| # we may not resize the input, but in [rosinality/stylegan2-pytorch] it | |
| # does resize the input. | |
| inception = InceptionV3([3], resize_input=resize_input, normalize_input=normalize_input) | |
| inception = nn.DataParallel(inception).eval().to(device) | |
| return inception | |
| def extract_inception_features(data_generator, inception, len_generator=None, device='cuda'): | |
| """Extract inception features. | |
| Args: | |
| data_generator (generator): A data generator. | |
| inception (nn.Module): Inception model. | |
| len_generator (int): Length of the data_generator to show the | |
| progressbar. Default: None. | |
| device (str): Device. Default: cuda. | |
| Returns: | |
| Tensor: Extracted features. | |
| """ | |
| if len_generator is not None: | |
| pbar = tqdm(total=len_generator, unit='batch', desc='Extract') | |
| else: | |
| pbar = None | |
| features = [] | |
| for data in data_generator: | |
| if pbar: | |
| pbar.update(1) | |
| data = data.to(device) | |
| feature = inception(data)[0].view(data.shape[0], -1) | |
| features.append(feature.to('cpu')) | |
| if pbar: | |
| pbar.close() | |
| features = torch.cat(features, 0) | |
| return features | |
| def calculate_fid(mu1, sigma1, mu2, sigma2, eps=1e-6): | |
| """Numpy implementation of the Frechet Distance. | |
| The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1) and X_2 ~ N(mu_2, C_2) is: | |
| d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)). | |
| Stable version by Dougal J. Sutherland. | |
| Args: | |
| mu1 (np.array): The sample mean over activations. | |
| sigma1 (np.array): The covariance matrix over activations for generated samples. | |
| mu2 (np.array): The sample mean over activations, precalculated on an representative data set. | |
| sigma2 (np.array): The covariance matrix over activations, precalculated on an representative data set. | |
| Returns: | |
| float: The Frechet Distance. | |
| """ | |
| assert mu1.shape == mu2.shape, 'Two mean vectors have different lengths' | |
| assert sigma1.shape == sigma2.shape, ('Two covariances have different dimensions') | |
| cov_sqrt, _ = linalg.sqrtm(sigma1 @ sigma2, disp=False) | |
| # Product might be almost singular | |
| if not np.isfinite(cov_sqrt).all(): | |
| print('Product of cov matrices is singular. Adding {eps} to diagonal of cov estimates') | |
| offset = np.eye(sigma1.shape[0]) * eps | |
| cov_sqrt = linalg.sqrtm((sigma1 + offset) @ (sigma2 + offset)) | |
| # Numerical error might give slight imaginary component | |
| if np.iscomplexobj(cov_sqrt): | |
| if not np.allclose(np.diagonal(cov_sqrt).imag, 0, atol=1e-3): | |
| m = np.max(np.abs(cov_sqrt.imag)) | |
| raise ValueError(f'Imaginary component {m}') | |
| cov_sqrt = cov_sqrt.real | |
| mean_diff = mu1 - mu2 | |
| mean_norm = mean_diff @ mean_diff | |
| trace = np.trace(sigma1) + np.trace(sigma2) - 2 * np.trace(cov_sqrt) | |
| fid = mean_norm + trace | |
| return fid | |