Spaces:
Running
on
L4
Running
on
L4
File size: 1,769 Bytes
3650c12 |
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 |
# Adapted from https://github.com/universome/fvd-comparison/blob/master/our_fvd.py
from typing import Tuple
import scipy
import numpy as np
import torch
def compute_fvd(feats_fake: np.ndarray, feats_real: np.ndarray) -> float:
mu_gen, sigma_gen = compute_stats(feats_fake)
mu_real, sigma_real = compute_stats(feats_real)
m = np.square(mu_gen - mu_real).sum()
s, _ = scipy.linalg.sqrtm(np.dot(sigma_gen, sigma_real), disp=False) # pylint: disable=no-member
fid = np.real(m + np.trace(sigma_gen + sigma_real - s * 2))
return float(fid)
def compute_stats(feats: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
mu = feats.mean(axis=0) # [d]
sigma = np.cov(feats, rowvar=False) # [d, d]
return mu, sigma
@torch.no_grad()
def compute_our_fvd(videos_fake: np.ndarray, videos_real: np.ndarray, device: str = "cuda") -> float:
i3d_path = "checkpoints/auxiliary/i3d_torchscript.pt"
i3d_kwargs = dict(
rescale=False, resize=False, return_features=True
) # Return raw features before the softmax layer.
with open(i3d_path, "rb") as f:
i3d_model = torch.jit.load(f).eval().to(device)
videos_fake = videos_fake.permute(0, 4, 1, 2, 3).to(device)
videos_real = videos_real.permute(0, 4, 1, 2, 3).to(device)
feats_fake = i3d_model(videos_fake, **i3d_kwargs).cpu().numpy()
feats_real = i3d_model(videos_real, **i3d_kwargs).cpu().numpy()
return compute_fvd(feats_fake, feats_real)
def main():
# input shape: (b, f, h, w, c)
videos_fake = torch.rand(10, 16, 224, 224, 3)
videos_real = torch.rand(10, 16, 224, 224, 3)
our_fvd_result = compute_our_fvd(videos_fake, videos_real)
print(f"[FVD scores] Ours: {our_fvd_result}")
if __name__ == "__main__":
main()
|