fid_score.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pathlib\n",
    "\n",
    "import numpy as np\n",
    "import cv2\n",
    "import torch\n",
    "from scipy import linalg\n",
    "from torch.nn.functional import adaptive_avg_pool2d\n",
    "\n",
    "try:\n",
    "    from tqdm import tqdm\n",
    "except ImportError:\n",
    "    # If not tqdm is not available, provide a mock version of it\n",
    "    def tqdm(x): return x\n",
    "\n",
    "def get_activations(files, model, batch_size=50, dims=2048,\n",
    "                    cuda=False, verbose=False):\n",
    "    \"\"\"Calculates the activations of the pool_3 layer for all images.\n",
    "    Params:\n",
    "    -- files       : List of image files paths\n",
    "    -- model       : Instance of inception model\n",
    "    -- batch_size  : Batch size of images for the model to process at once.\n",
    "                     Make sure that the number of samples is a multiple of\n",
    "                     the batch size, otherwise some samples are ignored. This\n",
    "                     behavior is retained to match the original FID score\n",
    "                     implementation.\n",
    "    -- dims        : Dimensionality of features returned by Inception\n",
    "    -- cuda        : If set to True, use GPU\n",
    "    -- verbose     : If set to True and parameter out_step is given, the number\n",
    "                     of calculated batches is reported.\n",
    "    Returns:\n",
    "    -- A numpy array of dimension (num images, dims) that contains the\n",
    "       activations of the given tensor when feeding inception with the\n",
    "       query tensor.\n",
    "    \"\"\"\n",
    "    model.eval()\n",
    "\n",
    "    if len(files) % batch_size != 0:\n",
    "        print(('Warning: number of images is not a multiple of the '\n",
    "               'batch size. Some samples are going to be ignored.'))\n",
    "    if batch_size > len(files):\n",
    "        print(('Warning: batch size is bigger than the data size. '\n",
    "               'Setting batch size to data size'))\n",
    "        batch_size = len(files)\n",
    "\n",
    "    n_batches = len(files) // batch_size\n",
    "    n_used_imgs = n_batches * batch_size\n",
    "\n",
    "    pred_arr = np.empty((n_used_imgs, dims))\n",
    "\n",
    "    for i in tqdm(range(n_batches)):\n",
    "        if verbose:\n",
    "            print('\\rPropagating batch %d/%d' % (i + 1, n_batches),\n",
    "                  end='', flush=True)\n",
    "        start = i * batch_size\n",
    "        end = start + batch_size\n",
    "\n",
    "        images = np.array([cv2.imread(str(f)).astype(np.float32)\n",
    "                           for f in files[start:end]])\n",
    "\n",
    "        # Reshape to (n_images, 3, height, width)\n",
    "        images = images.transpose((0, 3, 1, 2))\n",
    "        images /= 255\n",
    "\n",
    "        batch = torch.from_numpy(images).type(torch.FloatTensor)\n",
    "        if cuda:\n",
    "            batch = batch.cuda()\n",
    "\n",
    "        pred = model(batch)[0]\n",
    "\n",
    "        # If model output is not scalar, apply global spatial average pooling.\n",
    "        # This happens if you choose a dimensionality not equal 2048.\n",
    "        if pred.shape[2] != 1 or pred.shape[3] != 1:\n",
    "            pred = adaptive_avg_pool2d(pred, output_size=(1, 1))\n",
    "\n",
    "        pred_arr[start:end] = pred.cpu().data.numpy().reshape(batch_size, -1)\n",
    "\n",
    "    if verbose:\n",
    "        print(' done')\n",
    "\n",
    "    return pred_arr\n",
    "\n",
    "\n",
    "def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):\n",
    "    \"\"\"Numpy implementation of the Frechet Distance.\n",
    "    The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)\n",
    "    and X_2 ~ N(mu_2, C_2) is\n",
    "            d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).\n",
    "    Stable version by Dougal J. Sutherland.\n",
    "    Params:\n",
    "    -- mu1   : Numpy array containing the activations of a layer of the\n",
    "               inception net (like returned by the function 'get_predictions')\n",
    "               for generated samples.\n",
    "    -- mu2   : The sample mean over activations, precalculated on an\n",
    "               representative data set.\n",
    "    -- sigma1: The covariance matrix over activations for generated samples.\n",
    "    -- sigma2: The covariance matrix over activations, precalculated on an\n",
    "               representative data set.\n",
    "    Returns:\n",
    "    --   : The Frechet Distance.\n",
    "    \"\"\"\n",
    "\n",
    "    mu1 = np.atleast_1d(mu1)\n",
    "    mu2 = np.atleast_1d(mu2)\n",
    "\n",
    "    sigma1 = np.atleast_2d(sigma1)\n",
    "    sigma2 = np.atleast_2d(sigma2)\n",
    "\n",
    "    assert mu1.shape == mu2.shape, \\\n",
    "        'Training and test mean vectors have different lengths'\n",
    "    assert sigma1.shape == sigma2.shape, \\\n",
    "        'Training and test covariances have different dimensions'\n",
    "\n",
    "    diff = mu1 - mu2\n",
    "\n",
    "    # Product might be almost singular\n",
    "    covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)\n",
    "    if not np.isfinite(covmean).all():\n",
    "        msg = ('fid calculation produces singular product; '\n",
    "               'adding %s to diagonal of cov estimates') % eps\n",
    "        print(msg)\n",
    "        offset = np.eye(sigma1.shape[0]) * eps\n",
    "        covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))\n",
    "\n",
    "    # Numerical error might give slight imaginary component\n",
    "    if np.iscomplexobj(covmean):\n",
    "        if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):\n",
    "            m = np.max(np.abs(covmean.imag))\n",
    "            raise ValueError('Imaginary component {}'.format(m))\n",
    "        covmean = covmean.real\n",
    "\n",
    "    tr_covmean = np.trace(covmean)\n",
    "\n",
    "    return (diff.dot(diff) + np.trace(sigma1) +\n",
    "            np.trace(sigma2) - 2 * tr_covmean)\n",
    "\n",
    "\n",
    "def calculate_activation_statistics(files, model, batch_size=50,\n",
    "                                    dims=2048, cuda=False, verbose=False):\n",
    "    \"\"\"Calculation of the statistics used by the FID.\n",
    "    Params:\n",
    "    -- files       : List of image files paths\n",
    "    -- model       : Instance of inception model\n",
    "    -- batch_size  : The images numpy array is split into batches with\n",
    "                     batch size batch_size. A reasonable batch size\n",
    "                     depends on the hardware.\n",
    "    -- dims        : Dimensionality of features returned by Inception\n",
    "    -- cuda        : If set to True, use GPU\n",
    "    -- verbose     : If set to True and parameter out_step is given, the\n",
    "                     number of calculated batches is reported.\n",
    "    Returns:\n",
    "    -- mu    : The mean over samples of the activations of the pool_3 layer of\n",
    "               the inception model.\n",
    "    -- sigma : The covariance matrix of the activations of the pool_3 layer of\n",
    "               the inception model.\n",
    "    \"\"\"\n",
    "    act = get_activations(files, model, batch_size, dims, cuda, verbose)\n",
    "    mu = np.mean(act, axis=0)\n",
    "    sigma = np.cov(act, rowvar=False)\n",
    "    return mu, sigma\n",
    "\n",
    "\n",
    "def _compute_statistics_of_path(path, model, batch_size, dims, cuda):\n",
    "    if path.endswith('.npz'):\n",
    "        f = np.load(path)\n",
    "        m, s = f['mu'][:], f['sigma'][:]\n",
    "        f.close()\n",
    "    else:\n",
    "        path = pathlib.Path(path)\n",
    "        files = list(path.glob('*.jpg')) + list(path.glob('*.png'))\n",
    "        m, s = calculate_activation_statistics(files, model, batch_size,\n",
    "                                               dims, cuda)\n",
    "\n",
    "    return m, s\n",
    "\n",
    "\n",
    "def calculate_fid_given_paths(paths, batch_size, cuda, dims):\n",
    "    \"\"\"Calculates the FID of two paths\"\"\"\n",
    "    for p in paths:\n",
    "        if not os.path.exists(p):\n",
    "            raise RuntimeError('Invalid path: %s' % p)\n",
    "\n",
    "    block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[dims]\n",
    "\n",
    "    model = InceptionV3([block_idx])\n",
    "    if cuda:\n",
    "        model.cuda()\n",
    "\n",
    "    m1, s1 = _compute_statistics_of_path(paths[0], model, batch_size,\n",
    "                                         dims, cuda)\n",
    "    m2, s2 = _compute_statistics_of_path(paths[1], model, batch_size,\n",
    "                                         dims, cuda)\n",
    "    fid_value = calculate_frechet_distance(m1, s1, m2, s2)\n",
    "\n",
    "    return fid_value"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}