facebook/StyleNeRF

1# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.2#3# NVIDIA CORPORATION and its licensors retain all intellectual property4# and proprietary rights in and to this software, related documentation5# and any modifications thereto.  Any use, reproduction, disclosure or6# distribution of this software and related documentation without an express7# license agreement from NVIDIA CORPORATION is strictly prohibited.8 9"""Kernel Inception Distance (KID) from the paper "Demystifying MMD10GANs". Matches the original implementation by Binkowski et al. at11https://github.com/mbinkowski/MMD-GAN/blob/master/gan/compute_scores.py"""12 13import numpy as np14from . import metric_utils15 16#----------------------------------------------------------------------------17 18def compute_kid(opts, max_real, num_gen, num_subsets, max_subset_size):19    # Direct TorchScript translation of http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz20    detector_url = 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metrics/inception-2015-12-05.pt'21    detector_kwargs = dict(return_features=True) # Return raw features before the softmax layer.22 23    real_features = metric_utils.compute_feature_stats_for_dataset(24        opts=opts, detector_url=detector_url, detector_kwargs=detector_kwargs,25        rel_lo=0, rel_hi=0, capture_all=True, max_items=max_real).get_all()26 27    gen_features = metric_utils.compute_feature_stats_for_generator(28        opts=opts, detector_url=detector_url, detector_kwargs=detector_kwargs,29        rel_lo=0, rel_hi=1, capture_all=True, max_items=num_gen).get_all()30 31    if opts.rank != 0:32        return float('nan')33 34    n = real_features.shape[1]35    m = min(min(real_features.shape[0], gen_features.shape[0]), max_subset_size)36    t = 037    for _subset_idx in range(num_subsets):38        x = gen_features[np.random.choice(gen_features.shape[0], m, replace=False)]39        y = real_features[np.random.choice(real_features.shape[0], m, replace=False)]40        a = (x @ x.T / n + 1) ** 3 + (y @ y.T / n + 1) ** 341        b = (x @ y.T / n + 1) ** 342        t += (a.sum() - np.diag(a).sum()) / (m - 1) - b.sum() * 2 / m43    kid = t / num_subsets / m44    return float(kid)45 46#----------------------------------------------------------------------------47