File size: 5,785 Bytes
27ca8b3 |
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 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 |
"""
Adopted from https://github.com/cvpr2022-stylegan-v/stylegan-v
Verified to be the same as tf version by https://github.com/universome/fvd-comparison
"""
import io
import re
import requests
import html
import hashlib
import urllib
import urllib.request
from typing import Any, List, Tuple, Union, Dict
import scipy
import torch
import torch.nn as nn
import numpy as np
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 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
class FrechetVideoDistance(nn.Module):
def __init__(self):
super().__init__()
detector_url = (
"https://www.dropbox.com/s/ge9e5ujwgetktms/i3d_torchscript.pt?dl=1"
)
# Return raw features before the softmax layer.
self.detector_kwargs = dict(rescale=False, resize=True, return_features=True)
with open_url(detector_url, verbose=False) as f:
self.detector = torch.jit.load(f).eval()
@torch.no_grad()
def compute(self, videos_fake: torch.Tensor, videos_real: torch.Tensor):
"""
:param videos_fake: predicted video tensor of shape (frame, batch, channel, height, width)
:param videos_real: ground-truth observation tensor of shape (frame, batch, channel, height, width)
:return:
"""
n_frames, batch_size, c, h, w = videos_fake.shape
if n_frames < 2:
raise ValueError("Video must have more than 1 frame for FVD")
videos_fake = videos_fake.permute(1, 2, 0, 3, 4).contiguous()
videos_real = videos_real.permute(1, 2, 0, 3, 4).contiguous()
# detector takes in tensors of shape [batch_size, c, video_len, h, w] with range -1 to 1
feats_fake = self.detector(videos_fake, **self.detector_kwargs).cpu().numpy()
feats_real = self.detector(videos_real, **self.detector_kwargs).cpu().numpy()
return compute_fvd(feats_fake, feats_real)
|