295 lines
10 KiB
Python
295 lines
10 KiB
Python
|
# based on https://github.com/universome/fvd-comparison/blob/master/compare_models.py; huge thanks!
|
||
|
|
import os
|
||
|
|
import numpy as np
|
||
|
|
import io
|
||
|
|
import re
|
||
|
|
import requests
|
||
|
|
import html
|
||
|
|
import hashlib
|
||
|
|
import urllib
|
||
|
|
import urllib.request
|
||
|
|
import scipy.linalg
|
||
|
|
import multiprocessing as mp
|
||
|
|
import glob
|
||
|
|
|
||
|
|
|
||
|
|
from tqdm import tqdm
|
||
|
|
from typing import Any, List, Tuple, Union, Dict, Callable
|
||
|
|
|
||
|
|
from torchvision.io import read_video
|
||
|
|
import torch; torch.set_grad_enabled(False)
|
||
|
|
from einops import rearrange
|
||
|
|
|
||
|
|
from nitro.util import isvideo
|
||
|
|
|
||
|
|
def compute_frechet_distance(mu_sample,sigma_sample,mu_ref,sigma_ref) -> float:
|
||
|
|
print('Calculate frechet distance...')
|
||
|
|
m = np.square(mu_sample - mu_ref).sum()
|
||
|
|
s, _ = scipy.linalg.sqrtm(np.dot(sigma_sample, sigma_ref), disp=False) # pylint: disable=no-member
|
||
|
|
fid = np.real(m + np.trace(sigma_sample + sigma_ref - 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
|
||
|
|
|
||
|
|
|
||
|
|
def open_url(url: str, num_attempts: int = 10, verbose: bool = True, return_filename: bool = False) -> Any:
|
||
|
|
"""Download the given URL and return a binary-mode file object to access the data."""
|
||
|
|
assert num_attempts >= 1
|
||
|
|
|
||
|
|
# Doesn't look like an URL scheme so interpret it as a local filename.
|
||
|
|
if not re.match('^[a-z]+://', url):
|
||
|
|
return url if return_filename else open(url, "rb")
|
||
|
|
|
||
|
|
# Handle file URLs. This code handles unusual file:// patterns that
|
||
|
|
# arise on Windows:
|
||
|
|
#
|
||
|
|
# file:///c:/foo.txt
|
||
|
|
#
|
||
|
|
# which would translate to a local '/c:/foo.txt' filename that's
|
||
|
|
# invalid. Drop the forward slash for such pathnames.
|
||
|
|
#
|
||
|
|
# If you touch this code path, you should test it on both Linux and
|
||
|
|
# Windows.
|
||
|
|
#
|
||
|
|
# Some internet resources suggest using urllib.request.url2pathname() but
|
||
|
|
# but that converts forward slashes to backslashes and this causes
|
||
|
|
# its own set of problems.
|
||
|
|
if url.startswith('file://'):
|
||
|
|
filename = urllib.parse.urlparse(url).path
|
||
|
|
if re.match(r'^/[a-zA-Z]:', filename):
|
||
|
|
filename = filename[1:]
|
||
|
|
return filename if return_filename else open(filename, "rb")
|
||
|
|
|
||
|
|
url_md5 = hashlib.md5(url.encode("utf-8")).hexdigest()
|
||
|
|
|
||
|
|
# Download.
|
||
|
|
url_name = None
|
||
|
|
url_data = None
|
||
|
|
with requests.Session() as session:
|
||
|
|
if verbose:
|
||
|
|
print("Downloading %s ..." % url, end="", flush=True)
|
||
|
|
for attempts_left in reversed(range(num_attempts)):
|
||
|
|
try:
|
||
|
|
with session.get(url) as res:
|
||
|
|
res.raise_for_status()
|
||
|
|
if len(res.content) == 0:
|
||
|
|
raise IOError("No data received")
|
||
|
|
|
||
|
|
if len(res.content) < 8192:
|
||
|
|
content_str = res.content.decode("utf-8")
|
||
|
|
if "download_warning" in res.headers.get("Set-Cookie", ""):
|
||
|
|
links = [html.unescape(link) for link in content_str.split('"') if "export=download" in link]
|
||
|
|
if len(links) == 1:
|
||
|
|
url = requests.compat.urljoin(url, links[0])
|
||
|
|
raise IOError("Google Drive virus checker nag")
|
||
|
|
if "Google Drive - Quota exceeded" in content_str:
|
||
|
|
raise IOError("Google Drive download quota exceeded -- please try again later")
|
||
|
|
|
||
|
|
match = re.search(r'filename="([^"]*)"', res.headers.get("Content-Disposition", ""))
|
||
|
|
url_name = match[1] if match else url
|
||
|
|
url_data = res.content
|
||
|
|
if verbose:
|
||
|
|
print(" done")
|
||
|
|
break
|
||
|
|
except KeyboardInterrupt:
|
||
|
|
raise
|
||
|
|
except:
|
||
|
|
if not attempts_left:
|
||
|
|
if verbose:
|
||
|
|
print(" failed")
|
||
|
|
raise
|
||
|
|
if verbose:
|
||
|
|
print(".", end="", flush=True)
|
||
|
|
|
||
|
|
# Return data as file object.
|
||
|
|
assert not return_filename
|
||
|
|
return io.BytesIO(url_data)
|
||
|
|
|
||
|
|
def load_video(ip):
|
||
|
|
vid, *_ = read_video(ip)
|
||
|
|
vid = rearrange(vid, 't h w c -> t c h w').to(torch.uint8)
|
||
|
|
return vid
|
||
|
|
|
||
|
|
def get_data_from_str(input_str,nprc = None):
|
||
|
|
assert os.path.isdir(input_str), f'Specified input folder "{input_str}" is not a directory'
|
||
|
|
vid_filelist = glob.glob(os.path.join(input_str,'*.mp4'))
|
||
|
|
print(f'Found {len(vid_filelist)} videos in dir {input_str}')
|
||
|
|
|
||
|
|
if nprc is None:
|
||
|
|
try:
|
||
|
|
nprc = mp.cpu_count()
|
||
|
|
except NotImplementedError:
|
||
|
|
print('WARNING: cpu_count() not avlailable, using only 1 cpu for video loading')
|
||
|
|
nprc = 1
|
||
|
|
|
||
|
|
pool = mp.Pool(processes=nprc)
|
||
|
|
|
||
|
|
vids = []
|
||
|
|
for v in tqdm(pool.imap_unordered(load_video,vid_filelist),total=len(vid_filelist),desc='Loading videos...'):
|
||
|
|
vids.append(v)
|
||
|
|
|
||
|
|
|
||
|
|
vids = torch.stack(vids,dim=0).float()
|
||
|
|
|
||
|
|
return vids
|
||
|
|
|
||
|
|
def get_stats(stats):
|
||
|
|
assert os.path.isfile(stats) and stats.endswith('.npz'), f'no stats found under {stats}'
|
||
|
|
|
||
|
|
print(f'Using precomputed statistics under {stats}')
|
||
|
|
stats = np.load(stats)
|
||
|
|
stats = {key: stats[key] for key in stats.files}
|
||
|
|
|
||
|
|
return stats
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
@torch.no_grad()
|
||
|
|
def compute_fvd(ref_input, sample_input, bs=32,
|
||
|
|
ref_stats=None,
|
||
|
|
sample_stats=None,
|
||
|
|
nprc_load=None):
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
calc_stats = ref_stats is None or sample_stats is None
|
||
|
|
|
||
|
|
if calc_stats:
|
||
|
|
|
||
|
|
only_ref = sample_stats is not None
|
||
|
|
only_sample = ref_stats is not None
|
||
|
|
|
||
|
|
|
||
|
|
if isinstance(ref_input,str) and not only_sample:
|
||
|
|
ref_input = get_data_from_str(ref_input,nprc_load)
|
||
|
|
|
||
|
|
if isinstance(sample_input, str) and not only_ref:
|
||
|
|
sample_input = get_data_from_str(sample_input, nprc_load)
|
||
|
|
|
||
|
|
stats = compute_statistics(sample_input,ref_input,
|
||
|
|
device='cuda' if torch.cuda.is_available() else 'cpu',
|
||
|
|
bs=bs,
|
||
|
|
only_ref=only_ref,
|
||
|
|
only_sample=only_sample)
|
||
|
|
|
||
|
|
if only_ref:
|
||
|
|
stats.update(get_stats(sample_stats))
|
||
|
|
elif only_sample:
|
||
|
|
stats.update(get_stats(ref_stats))
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
else:
|
||
|
|
stats = get_stats(sample_stats)
|
||
|
|
stats.update(get_stats(ref_stats))
|
||
|
|
|
||
|
|
fvd = compute_frechet_distance(**stats)
|
||
|
|
|
||
|
|
return {'FVD' : fvd,}
|
||
|
|
|
||
|
|
|
||
|
|
@torch.no_grad()
|
||
|
|
def compute_statistics(videos_fake, videos_real, device: str='cuda', bs=32, only_ref=False,only_sample=False) -> Dict:
|
||
|
|
detector_url = 'https://www.dropbox.com/s/ge9e5ujwgetktms/i3d_torchscript.pt?dl=1'
|
||
|
|
detector_kwargs = dict(rescale=True, resize=True, return_features=True) # Return raw features before the softmax layer.
|
||
|
|
|
||
|
|
with open_url(detector_url, verbose=False) as f:
|
||
|
|
detector = torch.jit.load(f).eval().to(device)
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
assert not (only_sample and only_ref), 'only_ref and only_sample arguments are mutually exclusive'
|
||
|
|
|
||
|
|
ref_embed, sample_embed = [], []
|
||
|
|
|
||
|
|
info = f'Computing I3D activations for FVD score with batch size {bs}'
|
||
|
|
|
||
|
|
if only_ref:
|
||
|
|
|
||
|
|
if not isvideo(videos_real):
|
||
|
|
# if not is video we assume to have numpy arrays pf shape (n_vids, t, h, w, c) in range [0,255]
|
||
|
|
videos_real = torch.from_numpy(videos_real).permute(0, 4, 1, 2, 3).float()
|
||
|
|
print(videos_real.shape)
|
||
|
|
|
||
|
|
if videos_real.shape[0] % bs == 0:
|
||
|
|
n_secs = videos_real.shape[0] // bs
|
||
|
|
else:
|
||
|
|
n_secs = videos_real.shape[0] // bs + 1
|
||
|
|
|
||
|
|
videos_real = torch.tensor_split(videos_real, n_secs, dim=0)
|
||
|
|
|
||
|
|
for ref_v in tqdm(videos_real, total=len(videos_real),desc=info):
|
||
|
|
|
||
|
|
feats_ref = detector(ref_v.to(device).contiguous(), **detector_kwargs).cpu().numpy()
|
||
|
|
ref_embed.append(feats_ref)
|
||
|
|
|
||
|
|
elif only_sample:
|
||
|
|
|
||
|
|
if not isvideo(videos_fake):
|
||
|
|
# if not is video we assume to have numpy arrays pf shape (n_vids, t, h, w, c) in range [0,255]
|
||
|
|
videos_fake = torch.from_numpy(videos_fake).permute(0, 4, 1, 2, 3).float()
|
||
|
|
print(videos_fake.shape)
|
||
|
|
|
||
|
|
if videos_fake.shape[0] % bs == 0:
|
||
|
|
n_secs = videos_fake.shape[0] // bs
|
||
|
|
else:
|
||
|
|
n_secs = videos_fake.shape[0] // bs + 1
|
||
|
|
|
||
|
|
videos_real = torch.tensor_split(videos_real, n_secs, dim=0)
|
||
|
|
|
||
|
|
for sample_v in tqdm(videos_fake, total=len(videos_real),desc=info):
|
||
|
|
feats_sample = detector(sample_v.to(device).contiguous(), **detector_kwargs).cpu().numpy()
|
||
|
|
sample_embed.append(feats_sample)
|
||
|
|
|
||
|
|
|
||
|
|
else:
|
||
|
|
|
||
|
|
if not isvideo(videos_real):
|
||
|
|
# if not is video we assume to have numpy arrays pf shape (n_vids, t, h, w, c) in range [0,255]
|
||
|
|
videos_real = torch.from_numpy(videos_real).permute(0, 4, 1, 2, 3).float()
|
||
|
|
|
||
|
|
if not isvideo(videos_fake):
|
||
|
|
videos_fake = torch.from_numpy(videos_fake).permute(0, 4, 1, 2, 3).float()
|
||
|
|
|
||
|
|
if videos_fake.shape[0] % bs == 0:
|
||
|
|
n_secs = videos_fake.shape[0] // bs
|
||
|
|
else:
|
||
|
|
n_secs = videos_fake.shape[0] // bs + 1
|
||
|
|
|
||
|
|
videos_real = torch.tensor_split(videos_real, n_secs, dim=0)
|
||
|
|
videos_fake = torch.tensor_split(videos_fake, n_secs, dim=0)
|
||
|
|
|
||
|
|
for ref_v, sample_v in tqdm(zip(videos_real,videos_fake),total=len(videos_fake),desc=info):
|
||
|
|
# print(ref_v.shape)
|
||
|
|
# ref_v = torch.nn.functional.interpolate(ref_v, size=(sample_v.shape[2], 256, 256), mode='trilinear', align_corners=False)
|
||
|
|
# sample_v = torch.nn.functional.interpolate(sample_v, size=(sample_v.shape[2], 256, 256), mode='trilinear', align_corners=False)
|
||
|
|
|
||
|
|
|
||
|
|
feats_sample = detector(sample_v.to(device).contiguous(), **detector_kwargs).cpu().numpy()
|
||
|
|
feats_ref = detector(ref_v.to(device).contiguous(), **detector_kwargs).cpu().numpy()
|
||
|
|
sample_embed.append(feats_sample)
|
||
|
|
ref_embed.append(feats_ref)
|
||
|
|
|
||
|
|
out = dict()
|
||
|
|
if len(sample_embed) > 0:
|
||
|
|
sample_embed = np.concatenate(sample_embed,axis=0)
|
||
|
|
mu_sample, sigma_sample = compute_stats(sample_embed)
|
||
|
|
out.update({'mu_sample': mu_sample,
|
||
|
|
'sigma_sample': sigma_sample})
|
||
|
|
|
||
|
|
if len(ref_embed) > 0:
|
||
|
|
ref_embed = np.concatenate(ref_embed,axis=0)
|
||
|
|
mu_ref, sigma_ref = compute_stats(ref_embed)
|
||
|
|
out.update({'mu_ref': mu_ref,
|
||
|
|
'sigma_ref': sigma_ref})
|
||
|
|
|
||
|
|
|
||
|
|
return out
|