-
Notifications
You must be signed in to change notification settings - Fork 6
/
fid.py
298 lines (259 loc) · 11.5 KB
/
fid.py
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
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
import torch
from torch import nn
from torchvision.models import inception_v3
import cv2
import multiprocessing
import numpy as np
import glob
import os
from scipy import linalg
from torch.utils.data import DataLoader
import torchvision.datasets as datasets
import torchvision.transforms as transforms
from tqdm import tqdm
def to_cuda(elements):
"""
Transfers elements to cuda if GPU is available
Args:
elements: torch.tensor or torch.nn.module
--
Returns:
elements: same as input on GPU memory, if available
"""
if torch.cuda.is_available():
return elements.cuda()
return elements
class PartialInceptionNetwork(nn.Module):
def __init__(self, transform_input=True):
super().__init__()
self.inception_network = inception_v3(pretrained=True)
self.inception_network.Mixed_7c.register_forward_hook(self.output_hook)
self.transform_input = transform_input
def output_hook(self, module, input, output):
# N x 2048 x 8 x 8
self.mixed_7c_output = output
def forward(self, x):
"""
Args:
x: shape (N, 3, 299, 299) dtype: torch.float32 in range 0-1
Returns:
inception activations: torch.tensor, shape: (N, 2048), dtype: torch.float32
"""
assert x.shape[1:] == (3, 299, 299), "Expected input shape to be: (N,3,299,299)" +\
", but got {}".format(x.shape)
x = x * 2 -1 # Normalize to [-1, 1]
# Trigger output hook
self.inception_network(x)
# Output: N x 2048 x 1 x 1
activations = self.mixed_7c_output
activations = torch.nn.functional.adaptive_avg_pool2d(activations, (1,1))
activations = activations.view(x.shape[0], 2048)
return activations
def get_activations(dataloader, num_images):
"""
Calculates activations for last pool layer for all iamges
--
Images: torch.array shape: (N, 3, 299, 299), dtype: torch.float32
batch size: batch size used for inception network
--
Returns: np array shape: (N, 2048), dtype: np.float32
"""
inception_network = PartialInceptionNetwork()
inception_network = inception_network.cuda()
inception_network.eval()
inception_activations = np.zeros((num_images, 2048), dtype=np.float32)
for count, data in enumerate(tqdm(dataloader)):
images, _, _=data
images=images.cuda()
activations=inception_network(images)
activations=activations.detach().cpu().numpy()
batch_size=activations.shape[0]
start_idx = batch_size * count
end_idx = batch_size * (count + 1)
inception_activations[start_idx:end_idx, :]=activations
return inception_activations
def calculate_activation_statistics(dataloader, num_img):
"""Calculates the statistics used by FID
Args:
images: torch.tensor, shape: (N, 3, H, W), dtype: torch.float32 in range 0 - 1
batch_size: batch size to use to calculate inception scores
Returns:
mu: mean over all activations from the last pool layer of the inception model
sigma: covariance matrix over all activations from the last pool layer
of the inception model.
"""
act = get_activations(dataloader, num_img)
mu = np.mean(act, axis=0)
sigma = np.cov(act, rowvar=False)
return mu, sigma
# Modified from: https://github.com/bioinf-jku/TTUR/blob/master/fid.py
def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
"""Numpy implementation of the Frechet Distance.
The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)
and X_2 ~ N(mu_2, C_2) is
d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).
Stable version by Dougal J. Sutherland.
Params:
-- mu1 : Numpy array containing the activations of the pool_3 layer of the
inception net ( like returned by the function 'get_predictions')
for generated samples.
-- mu2 : The sample mean over activations of the pool_3 layer, precalcualted
on an representive data set.
-- sigma1: The covariance matrix over activations of the pool_3 layer for
generated samples.
-- sigma2: The covariance matrix over activations of the pool_3 layer,
precalcualted on an representive data set.
Returns:
-- : The Frechet Distance.
"""
mu1 = np.atleast_1d(mu1)
mu2 = np.atleast_1d(mu2)
sigma1 = np.atleast_2d(sigma1)
sigma2 = np.atleast_2d(sigma2)
assert mu1.shape == mu2.shape, "Training and test mean vectors have different lengths"
assert sigma1.shape == sigma2.shape, "Training and test covariances have different dimensions"
diff = mu1 - mu2
# product might be almost singular
covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
if not np.isfinite(covmean).all():
msg = "fid calculation produces singular product; adding %s to diagonal of cov estimates" % eps
warnings.warn(msg)
offset = np.eye(sigma1.shape[0]) * eps
covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
# numerical error might give slight imaginary component
if np.iscomplexobj(covmean):
if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
m = np.max(np.abs(covmean.imag))
raise ValueError("Imaginary component {}".format(m))
covmean = covmean.real
tr_covmean = np.trace(covmean)
return diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean
def preprocess_image(im):
"""Resizes and shifts the dynamic range of image to 0-1
Args:
im: np.array, shape: (H, W, 3), dtype: float32 between 0-1 or np.uint8
Return:
im: torch.tensor, shape: (3, 299, 299), dtype: torch.float32 between 0-1
"""
assert im.shape[2] == 3
assert len(im.shape) == 3
if im.dtype == np.uint8:
im = im.astype(np.float32) / 255
im = cv2.resize(im, (299, 299))
im = np.rollaxis(im, axis=2)
im = torch.from_numpy(im)
assert im.max() <= 1.0
assert im.min() >= 0.0
assert im.dtype == torch.float32
assert im.shape == (3, 299, 299)
return im
def preprocess_images(images, use_multiprocessing):
"""Resizes and shifts the dynamic range of image to 0-1
Args:
images: np.array, shape: (N, H, W, 3), dtype: float32 between 0-1 or np.uint8
use_multiprocessing: If multiprocessing should be used to pre-process the images
Return:
final_images: torch.tensor, shape: (N, 3, 299, 299), dtype: torch.float32 between 0-1
"""
if use_multiprocessing:
with multiprocessing.Pool(multiprocessing.cpu_count()) as pool:
jobs = []
for im in images:
job = pool.apply_async(preprocess_image, (im,))
jobs.append(job)
final_images = torch.zeros(images.shape[0], 3, 299, 299)
for idx, job in enumerate(jobs):
im = job.get()
final_images[idx] = im#job.get()
else:
final_images = torch.stack([preprocess_image(im) for im in images], dim=0)
assert final_images.shape == (images.shape[0], 3, 299, 299)
assert final_images.max() <= 1.0
assert final_images.min() >= 0.0
assert final_images.dtype == torch.float32
return final_images
def calculate_fid(dataloader1, dataloader2, use_multiprocessing, num_img1, num_img2):
""" Calculate FID between images1 and images2
Args:
images1: np.array, shape: (N, H, W, 3), dtype: np.float32 between 0-1 or np.uint8
images2: np.array, shape: (N, H, W, 3), dtype: np.float32 between 0-1 or np.uint8
use_multiprocessing: If multiprocessing should be used to pre-process the images
batch size: batch size used for inception network
Returns:
FID (scalar)
"""
#images1 = preprocess_images(images1, use_multiprocessing)
#images2 = preprocess_images(images2, use_multiprocessing)
print("calculating FID features of original dataset")
mu1, sigma1 = calculate_activation_statistics(dataloader1, num_img1)
print("calculating FID features of super-resolved dataset")
mu2, sigma2 = calculate_activation_statistics(dataloader2, num_img2)
fid = calculate_frechet_distance(mu1, sigma1, mu2, sigma2)
return fid
def load_images(path):
""" Loads all .png or .jpg images from a given path
Warnings: Expects all images to be of same dtype and shape.
Args:
path: relative path to directory
Returns:
final_images: np.array of image dtype and shape.
"""
image_paths = []
image_extensions = ["png", "jpg", "jpeg"]
for ext in image_extensions:
print("Looking for images in", os.path.join(path, "*.{}".format(ext)))
for impath in glob.glob(os.path.join(path, "*.{}".format(ext))):
image_paths.append(impath)
print(image_paths[0])
first_image = cv2.imread(image_paths[0])
W, H = first_image.shape[:2]
image_paths.sort()
image_paths = image_paths
final_images = np.zeros((len(image_paths), H, W, 3), dtype=first_image.dtype)
for idx, impath in enumerate(image_paths):
im = cv2.imread(impath)
im = im[:, :, ::-1] # Convert from BGR to RGB
assert im.dtype == final_images.dtype
final_images[idx] = im
return final_images
celeba="C:/Users/Asus/PycharmProjects/AutoEncoder/LOGS/Original_Images_Test"
lfw="C:/Users/Asus/PycharmProjects/datasets/aligned_lfw_all_in_one"
helen="C:/Users/Asus/PycharmProjects/datasets/Hellen_clean_aligned"
logs_path="C:/Users/Asus/PycharmProjects/AutoEncoder/LOGS/"
scale=16
algorithm="bicubic"
test_dataset_name="helen"
dest_path = logs_path+algorithm.upper()+str(scale)+"/sr_images_"+ test_dataset_name
if test_dataset_name=="celeba":
src_path=celeba
elif test_dataset_name=="lfw":
src_path=lfw
elif test_dataset_name=="helen":
src_path=helen
if __name__ == "__main__":
from optparse import OptionParser
parser = OptionParser()
parser.add_option("--p1", "--path1", dest="path1", default=src_path,
help="Path to directory containing the real images")
parser.add_option("--p2", "--path2", dest="path2", default=dest_path,
help="Path to directory containing the generated images")
parser.add_option("--multiprocessing", dest="use_multiprocessing",
help="Toggle use of multiprocessing for image pre-processing. Defaults to use all cores",
default=False,
action="store_true")
parser.add_option("-b", "--batch-size", dest="batch_size", default=16,
help="Set batch size to use for InceptionV3 network",
type=int)
options, _ = parser.parse_args()
assert options.path1 is not None, "--path1 is an required option"
assert options.path2 is not None, "--path2 is an required option"
assert options.batch_size is not None, "--batch_size is an required option"
transform=transforms.Compose([transforms.Resize(299), transforms.ToTensor()])
dataset1=datasets.ImageFolder(options.path1, transform=transform)
dataloader1=DataLoader(dataset1, batch_size=options.batch_size, shuffle=False)
num_image1=len(dataset1)
dataset2 = datasets.ImageFolder(options.path2, transform=transform)
dataloader2 = DataLoader(dataset2, batch_size=options.batch_size, shuffle=False)
num_image2 = len(dataset2)
fid_value = calculate_fid(dataloader1, dataloader2, options.use_multiprocessing, num_image1, num_image2)
print(fid_value)