acmyu/KeyframesAI
0
1import os2import pathlib3import torch4import numpy as np5import skimage6from imageio import imread7from scipy import linalg8from torch.nn.functional import adaptive_avg_pool2d9from skimage.metrics import structural_similarity as compare_ssim10from skimage.metrics import peak_signal_noise_ratio as compare_psnr11import glob12import argparse13import matplotlib.pyplot as plt14from inception import InceptionV315#from scripts.PerceptualSimilarity.models import dist_model as dm16import lpips17import pandas as pd18import json19import imageio20import cv221print(skimage.__version__)22 23class FID():24 """docstring for FID25 Calculates the Frechet Inception Distance (FID) to evalulate GANs26 The FID metric calculates the distance between two distributions of images.27 Typically, we have summary statistics (mean & covariance matrix) of one28 of these distributions, while the 2nd distribution is given by a GAN.29 When run as a stand-alone program, it compares the distribution of30 images that are stored as PNG/JPEG at a specified location with a31 distribution given by summary statistics (in pickle format).32 The FID is calculated by assuming that X_1 and X_2 are the activations of33 the pool_3 layer of the inception net for generated samples and real world34 samples respectivly.35 See --help to see further details.36 Code apapted from https://github.com/bioinf-jku/TTUR to use PyTorch instead37 of Tensorflow38 Copyright 2018 Institute of Bioinformatics, JKU Linz39 Licensed under the Apache License, Version 2.0 (the "License");40 you may not use this file except in compliance with the License.41 You may obtain a copy of the License at42 http://www.apache.org/licenses/LICENSE-2.043 Unless required by applicable law or agreed to in writing, software44 distributed under the License is distributed on an "AS IS" BASIS,45 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.46 See the License for the specific language governing permissions and47 limitations under the License.48 """49 def __init__(self):50 self.dims = 204851 self.batch_size = 12852 self.cuda = True53 self.verbose=False54 55 block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[self.dims]56 self.model = InceptionV3([block_idx])57 if self.cuda:58 # TODO: put model into specific GPU59 self.model.cuda()60 61 def __call__(self, images, gt_path):62 """ images: list of the generated image. The values must lie between 0 and 1.63 gt_path: the path of the ground truth images. The values must lie between 0 and 1.64 """65 if not os.path.exists(gt_path):66 raise RuntimeError('Invalid path: %s' % gt_path)67 68 69 print('calculate gt_path statistics...')70 m1, s1 = self.compute_statistics_of_path(gt_path, self.verbose)71 print('calculate generated_images statistics...')72 m2, s2 = self.calculate_activation_statistics(images, self.verbose)73 fid_value = self.calculate_frechet_distance(m1, s1, m2, s2)74 return fid_value75 76 77 def calculate_from_disk(self, generated_path, gt_path, img_size):78 """ 79 """80 if not os.path.exists(gt_path):81 raise RuntimeError('Invalid path: %s' % gt_path)82 if not os.path.exists(generated_path):83 raise RuntimeError('Invalid path: %s' % generated_path)84 85 print ('exp-path - '+generated_path)86 87 print('calculate gt_path statistics...')88 m1, s1 = self.compute_statistics_of_path(gt_path, self.verbose, img_size)89 print('calculate generated_path statistics...')90 m2, s2 = self.compute_statistics_of_path(generated_path, self.verbose, img_size)91 print('calculate frechet distance...')92 fid_value = self.calculate_frechet_distance(m1, s1, m2, s2)93 print('fid_distance %f' % (fid_value))94 return fid_value 95 96 97 def compute_statistics_of_path(self, path , verbose, img_size):98 99 size_flag = '{}_{}'.format(img_size[0], img_size[1])100 npz_file = os.path.join(path, size_flag + '_statistics.npz')101 if os.path.exists(npz_file):102 f = np.load(npz_file)103 m, s = f['mu'][:], f['sigma'][:]104 f.close()105 106 else:107 108 path = pathlib.Path(path)109 files = list(path.glob('*.jpg')) + list(path.glob('*.png'))110 111 imgs = (np.array([(cv2.resize(imread(str(fn)).astype(np.float32),img_size,interpolation=cv2.INTER_CUBIC)) for fn in files]))/255.0112 # Bring images to shape (B, 3, H, W)113 imgs = imgs.transpose((0, 3, 1, 2))114 115 # Rescale images to be between 0 and 1116 117 118 m, s = self.calculate_activation_statistics(imgs, verbose)119 np.savez(npz_file, mu=m, sigma=s)120 121 return m, s 122 123 def calculate_activation_statistics(self, images, verbose):124 """Calculation of the statistics used by the FID.125 Params:126 -- images : Numpy array of dimension (n_images, 3, hi, wi). The values127 must lie between 0 and 1.128 -- model : Instance of inception model129 -- batch_size : The images numpy array is split into batches with130 batch size batch_size. A reasonable batch size131 depends on the hardware.132 -- dims : Dimensionality of features returned by Inception133 -- cuda : If set to True, use GPU134 -- verbose : If set to True and parameter out_step is given, the135 number of calculated batches is reported.136 Returns:137 -- mu : The mean over samples of the activations of the pool_3 layer of138 the inception model.139 -- sigma : The covariance matrix of the activations of the pool_3 layer of140 the inception model.141 """142 act = self.get_activations(images, verbose)143 mu = np.mean(act, axis=0)144 sigma = np.cov(act, rowvar=False)145 return mu, sigma 146 147 148 149 def get_activations(self, images, verbose=False):150 """Calculates the activations of the pool_3 layer for all images.151 Params:152 -- images : Numpy array of dimension (n_images, 3, hi, wi). The values153 must lie between 0 and 1.154 -- model : Instance of inception model155 -- batch_size : the images numpy array is split into batches with156 batch size batch_size. A reasonable batch size depends157 on the hardware.158 -- dims : Dimensionality of features returned by Inception159 -- cuda : If set to True, use GPU160 -- verbose : If set to True and parameter out_step is given, the number161 of calculated batches is reported.162 Returns:163 -- A numpy array of dimension (num images, dims) that contains the164 activations of the given tensor when feeding inception with the165 query tensor.166 """167 self.model.eval()168 169 d0 = images.shape[0]170 if self.batch_size > d0:171 print(('Warning: batch size is bigger than the data size. '172 'Setting batch size to data size'))173 self.batch_size = d0174 175 n_batches = d0 // self.batch_size176 n_used_imgs = n_batches * self.batch_size177 178 pred_arr = np.empty((n_used_imgs, self.dims))179 for i in range(n_batches):180 if verbose:181 print('\rPropagating batch %d/%d' % (i + 1, n_batches))182 # end='', flush=True)183 start = i * self.batch_size184 end = start + self.batch_size185 186 batch = torch.from_numpy(images[start:end]).type(torch.FloatTensor)187 # batch = Variable(batch, volatile=True)188 if self.cuda:189 batch = batch.cuda()190 191 pred = self.model(batch)[0]192 193 # If model output is not scalar, apply global spatial average pooling.194 # This happens if you choose a dimensionality not equal 2048.195 if pred.shape[2] != 1 or pred.shape[3] != 1:196 pred = adaptive_avg_pool2d(pred, output_size=(1, 1))197 198 pred_arr[start:end] = pred.cpu().data.numpy().reshape(self.batch_size, -1)199 200 if verbose:201 print(' done')202 203 return pred_arr204 205 206 def calculate_frechet_distance(self, mu1, sigma1, mu2, sigma2, eps=1e-6):207 """Numpy implementation of the Frechet Distance.208 The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)209 and X_2 ~ N(mu_2, C_2) is210 d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).211 Stable version by Dougal J. Sutherland.212 Params:213 -- mu1 : Numpy array containing the activations of a layer of the214 inception net (like returned by the function 'get_predictions')215 for generated samples.216 -- mu2 : The sample mean over activations, precalculated on an 217 representive data set.218 -- sigma1: The covariance matrix over activations for generated samples.219 -- sigma2: The covariance matrix over activations, precalculated on an 220 representive data set.221 Returns:222 -- : The Frechet Distance.223 """224 225 mu1 = np.atleast_1d(mu1)226 mu2 = np.atleast_1d(mu2)227 228 sigma1 = np.atleast_2d(sigma1)229 sigma2 = np.atleast_2d(sigma2)230 231 assert mu1.shape == mu2.shape, \232 'Training and test mean vectors have different lengths'233 assert sigma1.shape == sigma2.shape, \234 'Training and test covariances have different dimensions'235 236 diff = mu1 - mu2237 238 # Product might be almost singular239 covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)240 if not np.isfinite(covmean).all():241 msg = ('fid calculation produces singular product; '242 'adding %s to diagonal of cov estimates') % eps243 print(msg)244 offset = np.eye(sigma1.shape[0]) * eps245 covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))246 247 # Numerical error might give slight imaginary component248 if np.iscomplexobj(covmean):249 if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):250 m = np.max(np.abs(covmean.imag))251 raise ValueError('Imaginary component {}'.format(m))252 covmean = covmean.real253 254 tr_covmean = np.trace(covmean)255 256 return (diff.dot(diff) + np.trace(sigma1) +257 np.trace(sigma2) - 2 * tr_covmean)258 259 260class Reconstruction_Metrics():261 def __init__(self, metric_list=['ssim', 'psnr', 'l1', 'mae'], data_range=1, win_size=51, multichannel=True):262 self.data_range = data_range263 self.win_size = win_size264 self.multichannel = multichannel265 for metric in metric_list:266 if metric in ['ssim', 'psnr', 'l1', 'mae']:267 setattr(self, metric, True)268 else:269 print('unsupport reconstruction metric: %s'%metric)270 271 272 def __call__(self, inputs, gts):273 """274 inputs: the generated image, size (b,c,w,h), data range(0, data_range)275 gts: the ground-truth image, size (b,c,w,h), data range(0, data_range)276 """277 result = dict() 278 [b,n,w,h] = inputs.size()279 inputs = inputs.view(b*n, w, h).detach().cpu().numpy().astype(np.float32).transpose(1,2,0)280 gts = gts.view(b*n, w, h).detach().cpu().numpy().astype(np.float32).transpose(1,2,0)281 282 if hasattr(self, 'ssim'):283 ssim_value = compare_ssim(inputs, gts, data_range=self.data_range, 284 win_size=self.win_size, multichannel=self.multichannel) 285 result['ssim'] = ssim_value286 287 288 if hasattr(self, 'psnr'):289 psnr_value = compare_psnr(inputs, gts, self.data_range)290 result['psnr'] = psnr_value291 292 if hasattr(self, 'l1'):293 l1_value = compare_l1(inputs, gts)294 result['l1'] = l1_value 295 296 if hasattr(self, 'mae'):297 mae_value = compare_mae(inputs, gts)298 result['mae'] = mae_value 299 return result300 301 302 def calculate_from_disk(self, inputs, gts, save_path=None, img_size=(176,256), sort=True, debug=0):303 """304 inputs: .txt files, floders, image files (string), image files (list)305 gts: .txt files, floders, image files (string), image files (list)306 """307 if sort:308 input_image_list = sorted(get_image_list(inputs))309 gt_image_list = sorted(get_image_list(gts))310 else:311 input_image_list = get_image_list(inputs)312 gt_image_list = get_image_list(gts)313 314 size_flag = '{}_{}'.format(img_size[0], img_size[1])315 npz_file = os.path.join(save_path, size_flag + '_metrics.npz')316 if os.path.exists(npz_file):317 f = np.load(npz_file)318 psnr,ssim,ssim_256,mae,l1=f['psnr'],f['ssim'],f['ssim_256'],f['mae'],f['l1']319 else:320 psnr = []321 ssim = []322 ssim_256 = []323 mae = []324 l1 = []325 names = []326 327 for index in range(len(input_image_list)):328 name = os.path.basename(input_image_list[index])329 names.append(name)330 331 332 img_gt = (cv2.resize(imread(str(gt_image_list[index])).astype(np.float32), img_size,interpolation=cv2.INTER_CUBIC)) /255.0333 img_pred = (cv2.resize(imread(str(input_image_list[index])).astype(np.float32), img_size,interpolation=cv2.INTER_CUBIC)) / 255.0334 335 336 if debug != 0:337 plt.subplot('121')338 plt.imshow(img_gt)339 plt.title('Groud truth')340 plt.subplot('122')341 plt.imshow(img_pred)342 plt.title('Output')343 plt.show()344 345 psnr.append(compare_psnr(img_gt, img_pred, data_range=self.data_range))346 ssim.append(compare_ssim(img_gt, img_pred, data_range=self.data_range,347 win_size=self.win_size,multichannel=self.multichannel, channel_axis=2))348 mae.append(compare_mae(img_gt, img_pred))349 l1.append(compare_l1(img_gt, img_pred))350 351 img_gt_256 = img_gt*255.0352 img_pred_256 = img_pred*255.0353 ssim_256.append(compare_ssim(img_gt_256, img_pred_256, gaussian_weights=True, sigma=1.2,354 use_sample_covariance=False, multichannel=True, channel_axis=2,355 data_range=img_pred_256.max() - img_pred_256.min()))356 357 if np.mod(index, 200) == 0:358 print(359 str(index) + ' images processed',360 "PSNR: %.4f" % round(np.mean(psnr), 4),361 "SSIM_256: %.4f" % round(np.mean(ssim_256), 4),362 "MAE: %.4f" % round(np.mean(mae), 4),363 "l1: %.4f" % round(np.mean(l1), 4),364 )365 366 if save_path:367 np.savez(save_path + '/' + size_flag + '_metrics.npz', psnr=psnr, ssim=ssim, ssim_256=ssim_256, mae=mae, l1=l1, names=names)368 369 print(370 "PSNR: %.4f" % round(np.mean(psnr), 4),371 "PSNR Variance: %.4f" % round(np.var(psnr), 4),372 "SSIM_256: %.4f" % round(np.mean(ssim_256), 4),373 "SSIM_256 Variance: %.4f" % round(np.var(ssim_256), 4), 374 "MAE: %.4f" % round(np.mean(mae), 4),375 "MAE Variance: %.4f" % round(np.var(mae), 4),376 "l1: %.4f" % round(np.mean(l1), 4),377 "l1 Variance: %.4f" % round(np.var(l1), 4) 378 ) 379 380 dic = {"psnr":[round(np.mean(psnr), 6)],381 "psnr_variance": [round(np.var(psnr), 6)],382 "ssim_256": [round(np.mean(ssim_256), 6)],383 "ssim_256_variance": [round(np.var(ssim_256), 6)],384 "mae": [round(np.mean(mae), 6)],385 "mae_variance": [round(np.var(mae), 6)],386 "l1": [round(np.mean(l1), 6)],387 "l1_variance": [round(np.var(l1), 6)] } 388 389 return dic390 391 392def get_image_list(flist):393 if isinstance(flist, list):394 return flist395 396 # flist: image file path, image directory path, text file flist path397 if isinstance(flist, str):398 if os.path.isdir(flist):399 flist = list(glob.glob(flist + '/*.jpg')) + list(glob.glob(flist + '/*.png'))400 flist.sort()401 return flist402 403 if os.path.isfile(flist):404 try:405 return np.genfromtxt(flist, dtype=np.str)406 except:407 return [flist]408 print('can not read files from %s return empty list'%flist)409 return []410 411def compare_l1(img_true, img_test):412 img_true = img_true.astype(np.float32)413 img_test = img_test.astype(np.float32)414 return np.mean(np.abs(img_true - img_test)) 415 416def compare_mae(img_true, img_test):417 img_true = img_true.astype(np.float32)418 img_test = img_test.astype(np.float32)419 return np.sum(np.abs(img_true - img_test)) / np.sum(img_true + img_test)420 421def preprocess_path_for_deform_task(gt_path, distorted_path):422 distorted_image_list = sorted(get_image_list(distorted_path))423 gt_list=[]424 distorated_list=[]425 426 for distorted_image in distorted_image_list:427 image = os.path.basename(distorted_image)[1:]428 image = image.split('_to_')[-1]429 gt_image = gt_path + '/' + image.replace('jpg', 'png')430 if not os.path.isfile(gt_image):431 print(distorted_image, gt_image)432 print('=====')433 continue434 gt_list.append(gt_image)435 distorated_list.append(distorted_image) 436 437 return gt_list, distorated_list438 439 440 441class LPIPS():442 def __init__(self, use_gpu=True):443 444 self.model = lpips.LPIPS(net='alex').eval().cuda()445 self.use_gpu=use_gpu446 447 def __call__(self, image_1, image_2):448 """449 image_1: images with size (n, 3, w, h) with value [-1, 1]450 image_2: images with size (n, 3, w, h) with value [-1, 1]451 """452 result = self.model.forward(image_1, image_2)453 return result454 455 def calculate_from_disk(self, path_1, path_2,img_size, batch_size=64, verbose=False, sort=True):456 457 if sort:458 files_1 = sorted(get_image_list(path_1))459 files_2 = sorted(get_image_list(path_2))460 else:461 files_1 = get_image_list(path_1)462 files_2 = get_image_list(path_2)463 464 465 results=[]466 467 468 d0 = len(files_1)469 if batch_size > d0:470 print(('Warning: batch size is bigger than the data size. '471 'Setting batch size to data size'))472 batch_size = d0473 474 n_batches = d0 // batch_size475 476 477 for i in range(n_batches):478 if verbose:479 print('\rPropagating batch %d/%d' % (i + 1, n_batches))480 # end='', flush=True)481 start = i * batch_size482 end = start + batch_size483 484 imgs_1 = np.array([cv2.resize(imread(str(fn)).astype(np.float32),img_size,interpolation=cv2.INTER_CUBIC)/255.0 for fn in files_1[start:end]])485 imgs_2 = np.array([cv2.resize(imread(str(fn)).astype(np.float32),img_size,interpolation=cv2.INTER_CUBIC)/255.0 for fn in files_2[start:end]])486 487 imgs_1 = imgs_1.transpose((0, 3, 1, 2))488 imgs_2 = imgs_2.transpose((0, 3, 1, 2))489 490 img_1_batch = torch.from_numpy(imgs_1).type(torch.FloatTensor)491 img_2_batch = torch.from_numpy(imgs_2).type(torch.FloatTensor)492 493 if self.use_gpu:494 img_1_batch = img_1_batch.cuda()495 img_2_batch = img_2_batch.cuda()496 497 with torch.no_grad():498 result = self.model.forward(img_1_batch, img_2_batch)499 500 results.append(result)501 502 503 distance = torch.cat(results,0)[:,0,0,0].mean()504 505 print('lpips: %.3f'%distance)506 return distance507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523