chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,333 @@
|
||||
# Adapted from https://github.com/mseitzer/pytorch-fid/blob/0a754fb8e66021700478fd365b79c2eaa316e31b/src/pytorch_fid/fid_score.py
|
||||
"""Calculates the Frechet Inception Distance (FID) to evalulate GANs
|
||||
|
||||
The FID metric calculates the distance between two distributions of images.
|
||||
Typically, we have summary statistics (mean & covariance matrix) of one
|
||||
of these distributions, while the 2nd distribution is given by a GAN.
|
||||
|
||||
When run as a stand-alone program, it compares the distribution of
|
||||
images that are stored as PNG/JPEG at a specified location with a
|
||||
distribution given by summary statistics (in pickle format).
|
||||
|
||||
The FID is calculated by assuming that X_1 and X_2 are the activations of
|
||||
the pool_3 layer of the inception net for generated samples and real world
|
||||
samples respectively.
|
||||
|
||||
See --help to see further details.
|
||||
|
||||
Code apapted from https://github.com/bioinf-jku/TTUR to use PyTorch instead
|
||||
of Tensorflow
|
||||
|
||||
Copyright 2018 Institute of Bioinformatics, JKU Linz
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
"""
|
||||
import os
|
||||
import pathlib
|
||||
from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torchvision.transforms as TF
|
||||
from PIL import Image
|
||||
from scipy import linalg
|
||||
from torch.nn.functional import adaptive_avg_pool2d
|
||||
|
||||
try:
|
||||
from tqdm import tqdm
|
||||
except ImportError:
|
||||
# If tqdm is not available, provide a mock version of it
|
||||
def tqdm(x):
|
||||
return x
|
||||
|
||||
from pytorch_fid.inception import InceptionV3
|
||||
|
||||
parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
|
||||
parser.add_argument('--batch-size', type=int, default=50,
|
||||
help='Batch size to use')
|
||||
parser.add_argument('--num-workers', type=int,
|
||||
help=('Number of processes to use for data loading. '
|
||||
'Defaults to `min(8, num_cpus)`'))
|
||||
parser.add_argument('--device', type=str, default=None,
|
||||
help='Device to use. Like cuda, cuda:0 or cpu')
|
||||
parser.add_argument('--dims', type=int, default=2048,
|
||||
choices=list(InceptionV3.BLOCK_INDEX_BY_DIM),
|
||||
help=('Dimensionality of Inception features to use. '
|
||||
'By default, uses pool3 features'))
|
||||
parser.add_argument('--save-stats', action='store_true',
|
||||
help=('Generate an npz archive from a directory of samples. '
|
||||
'The first path is used as input and the second as output.'))
|
||||
parser.add_argument('path', type=str, nargs=2,
|
||||
help=('Paths to the generated images or '
|
||||
'to .npz statistic files'))
|
||||
|
||||
IMAGE_EXTENSIONS = {'bmp', 'jpg', 'jpeg', 'pgm', 'png', 'ppm',
|
||||
'tif', 'tiff', 'webp'}
|
||||
|
||||
|
||||
class ImagePathDataset(torch.utils.data.Dataset):
|
||||
def __init__(self, files, transforms=None):
|
||||
self.files = files
|
||||
self.transforms = transforms
|
||||
|
||||
def __len__(self):
|
||||
return len(self.files)
|
||||
|
||||
def __getitem__(self, i):
|
||||
path = self.files[i]
|
||||
img = Image.open(path).convert('RGB')
|
||||
if self.transforms is not None:
|
||||
img = self.transforms(img)
|
||||
return img
|
||||
|
||||
|
||||
def get_activations(files, model, batch_size=50, dims=2048, device='cpu',
|
||||
num_workers=1):
|
||||
"""Calculates the activations of the pool_3 layer for all images.
|
||||
|
||||
Params:
|
||||
-- files : List of image files paths
|
||||
-- model : Instance of inception model
|
||||
-- batch_size : Batch size of images for the model to process at once.
|
||||
Make sure that the number of samples is a multiple of
|
||||
the batch size, otherwise some samples are ignored. This
|
||||
behavior is retained to match the original FID score
|
||||
implementation.
|
||||
-- dims : Dimensionality of features returned by Inception
|
||||
-- device : Device to run calculations
|
||||
-- num_workers : Number of parallel dataloader workers
|
||||
|
||||
Returns:
|
||||
-- A numpy array of dimension (num images, dims) that contains the
|
||||
activations of the given tensor when feeding inception with the
|
||||
query tensor.
|
||||
"""
|
||||
model.eval()
|
||||
|
||||
if batch_size > len(files):
|
||||
print(('Warning: batch size is bigger than the data size. '
|
||||
'Setting batch size to data size'))
|
||||
batch_size = len(files)
|
||||
|
||||
dataset = ImagePathDataset(files, transforms=TF.ToTensor())
|
||||
dataloader = torch.utils.data.DataLoader(dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=False,
|
||||
drop_last=False,
|
||||
num_workers=num_workers)
|
||||
|
||||
pred_arr = np.empty((len(files), dims))
|
||||
|
||||
start_idx = 0
|
||||
|
||||
for batch in tqdm(dataloader):
|
||||
batch = batch.to(device)
|
||||
|
||||
with torch.no_grad():
|
||||
pred = model(batch)[0]
|
||||
|
||||
# If model output is not scalar, apply global spatial average pooling.
|
||||
# This happens if you choose a dimensionality not equal 2048.
|
||||
if pred.size(2) != 1 or pred.size(3) != 1:
|
||||
pred = adaptive_avg_pool2d(pred, output_size=(1, 1))
|
||||
|
||||
pred = pred.squeeze(3).squeeze(2).cpu().numpy()
|
||||
|
||||
pred_arr[start_idx:start_idx + pred.shape[0]] = pred
|
||||
|
||||
start_idx = start_idx + pred.shape[0]
|
||||
|
||||
return pred_arr
|
||||
|
||||
|
||||
def calculate_frechet_distance(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.
|
||||
|
||||
Params:
|
||||
-- mu1 : Numpy array containing the activations of a layer of the
|
||||
inception net (like returned by the function 'get_predictions')
|
||||
for generated samples.
|
||||
-- mu2 : The sample mean over activations, precalculated on an
|
||||
representative data set.
|
||||
-- sigma1: The covariance matrix over activations for generated samples.
|
||||
-- sigma2: The covariance matrix over activations, precalculated on an
|
||||
representative data set.
|
||||
|
||||
Returns:
|
||||
-- : The Frechet Distance.
|
||||
"""
|
||||
|
||||
mu1 = np.atleast_1d(mu1)
|
||||
mu2 = np.atleast_1d(mu2)
|
||||
|
||||
sigma1 = np.atleast_2d(sigma1)
|
||||
sigma2 = np.atleast_2d(sigma2)
|
||||
|
||||
assert mu1.shape == mu2.shape, \
|
||||
'Training and test mean vectors have different lengths'
|
||||
assert sigma1.shape == sigma2.shape, \
|
||||
'Training and test covariances have different dimensions'
|
||||
|
||||
diff = mu1 - mu2
|
||||
|
||||
# Product might be almost singular
|
||||
covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
|
||||
if not np.isfinite(covmean).all():
|
||||
msg = ('fid calculation produces singular product; '
|
||||
'adding %s to diagonal of cov estimates') % eps
|
||||
print(msg)
|
||||
offset = np.eye(sigma1.shape[0]) * eps
|
||||
covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
|
||||
|
||||
# Numerical error might give slight imaginary component
|
||||
if np.iscomplexobj(covmean):
|
||||
if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
|
||||
m = np.max(np.abs(covmean.imag))
|
||||
raise ValueError('Imaginary component {}'.format(m))
|
||||
covmean = covmean.real
|
||||
|
||||
tr_covmean = np.trace(covmean)
|
||||
|
||||
return (diff.dot(diff) + np.trace(sigma1)
|
||||
+ np.trace(sigma2) - 2 * tr_covmean)
|
||||
|
||||
|
||||
def calculate_activation_statistics(files, model, batch_size=50, dims=2048,
|
||||
device='cpu', num_workers=1):
|
||||
"""Calculation of the statistics used by the FID.
|
||||
Params:
|
||||
-- files : List of image files paths
|
||||
-- model : Instance of inception model
|
||||
-- batch_size : The images numpy array is split into batches with
|
||||
batch size batch_size. A reasonable batch size
|
||||
depends on the hardware.
|
||||
-- dims : Dimensionality of features returned by Inception
|
||||
-- device : Device to run calculations
|
||||
-- num_workers : Number of parallel dataloader workers
|
||||
|
||||
Returns:
|
||||
-- mu : The mean over samples of the activations of the pool_3 layer of
|
||||
the inception model.
|
||||
-- sigma : The covariance matrix of the activations of the pool_3 layer of
|
||||
the inception model.
|
||||
"""
|
||||
act = get_activations(files, model, batch_size, dims, device, num_workers)
|
||||
mu = np.mean(act, axis=0)
|
||||
sigma = np.cov(act, rowvar=False)
|
||||
return mu, sigma
|
||||
|
||||
|
||||
def compute_statistics_of_path(path, model, batch_size, dims, device,
|
||||
num_workers=1):
|
||||
if type(path) is not list and path.endswith('.npz'):
|
||||
with np.load(path) as f:
|
||||
m, s = f['mu'][:], f['sigma'][:]
|
||||
else:
|
||||
if type(path) is list:
|
||||
files = []
|
||||
for p in path:
|
||||
p = pathlib.Path(p)
|
||||
files += sorted([file for ext in IMAGE_EXTENSIONS for file in p.glob('*.{}'.format(ext))])
|
||||
files = sorted(files)
|
||||
else:
|
||||
path = pathlib.Path(path)
|
||||
files = sorted([file for ext in IMAGE_EXTENSIONS for file in path.glob('*.{}'.format(ext))])
|
||||
m, s = calculate_activation_statistics(files, model, batch_size, dims, device, num_workers)
|
||||
|
||||
return m, s
|
||||
|
||||
|
||||
def calculate_fid_given_paths(paths, batch_size=50, device="cuda:0", dims=2048, num_workers=1):
|
||||
"""Calculates the FID of two paths"""
|
||||
for p in paths:
|
||||
if type(p) is list:
|
||||
for subp in p:
|
||||
if not os.path.exists(subp):
|
||||
raise RuntimeError('Invalid path: %s' % subp)
|
||||
else:
|
||||
if not os.path.exists(p):
|
||||
raise RuntimeError('Invalid path: %s' % p)
|
||||
|
||||
block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[dims]
|
||||
|
||||
model = InceptionV3([block_idx]).to(device)
|
||||
|
||||
m1, s1 = compute_statistics_of_path(paths[0], model, batch_size,
|
||||
dims, device, num_workers)
|
||||
m2, s2 = compute_statistics_of_path(paths[1], model, batch_size,
|
||||
dims, device, num_workers)
|
||||
fid_value = calculate_frechet_distance(m1, s1, m2, s2)
|
||||
|
||||
return fid_value
|
||||
|
||||
|
||||
def save_fid_stats(paths, batch_size, device, dims, num_workers=1):
|
||||
"""Calculates the FID of two paths"""
|
||||
if not os.path.exists(paths[0]):
|
||||
raise RuntimeError('Invalid path: %s' % paths[0])
|
||||
|
||||
if os.path.exists(paths[1]):
|
||||
raise RuntimeError('Existing output file: %s' % paths[1])
|
||||
|
||||
block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[dims]
|
||||
|
||||
model = InceptionV3([block_idx]).to(device)
|
||||
|
||||
print(f"Saving statistics for {paths[0]}")
|
||||
|
||||
m1, s1 = compute_statistics_of_path(paths[0], model, batch_size,
|
||||
dims, device, num_workers)
|
||||
|
||||
np.savez_compressed(paths[1], mu=m1, sigma=s1)
|
||||
|
||||
|
||||
def main():
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.device is None:
|
||||
device = torch.device('cuda' if (torch.cuda.is_available()) else 'cpu')
|
||||
else:
|
||||
device = torch.device(args.device)
|
||||
|
||||
if args.num_workers is None:
|
||||
try:
|
||||
num_cpus = len(os.sched_getaffinity(0))
|
||||
except AttributeError:
|
||||
# os.sched_getaffinity is not available under Windows, use
|
||||
# os.cpu_count instead (which may not return the *available* number
|
||||
# of CPUs).
|
||||
num_cpus = os.cpu_count()
|
||||
|
||||
num_workers = min(num_cpus, 8) if num_cpus is not None else 0
|
||||
else:
|
||||
num_workers = args.num_workers
|
||||
|
||||
if args.save_stats:
|
||||
save_fid_stats(args.path, args.batch_size, device, args.dims, num_workers)
|
||||
return
|
||||
|
||||
fid_value = calculate_fid_given_paths(args.path,
|
||||
args.batch_size,
|
||||
device,
|
||||
args.dims,
|
||||
num_workers)
|
||||
print('FID: ', fid_value)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in New Issue
Block a user