import os import pathlib import torch import numpy as np import skimage from imageio import imread from scipy import linalg from torch.nn.functional import adaptive_avg_pool2d from skimage.metrics import structural_similarity as compare_ssim from skimage.metrics import peak_signal_noise_ratio as compare_psnr import glob import argparse import matplotlib.pyplot as plt from inception import InceptionV3 #from scripts.PerceptualSimilarity.models import dist_model as dm import lpips import pandas as pd import json import imageio import cv2 print(skimage.__version__) class FID(): """docstring for FID Calculates the Frechet Inception Distance (FID) to evalulate GANs The FID metric calculates the distance between two distributions of images. Typically, we have summary statistics (mean & covariance matrix) of one of these distributions, while the 2nd distribution is given by a GAN. When run as a stand-alone program, it compares the distribution of images that are stored as PNG/JPEG at a specified location with a distribution given by summary statistics (in pickle format). The FID is calculated by assuming that X_1 and X_2 are the activations of the pool_3 layer of the inception net for generated samples and real world samples respectivly. See --help to see further details. Code apapted from https://github.com/bioinf-jku/TTUR to use PyTorch instead of Tensorflow Copyright 2018 Institute of Bioinformatics, JKU Linz Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. """ def __init__(self): self.dims = 2048 self.batch_size = 128 self.cuda = True self.verbose=False block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[self.dims] self.model = InceptionV3([block_idx]) if self.cuda: # TODO: put model into specific GPU self.model.cuda() def __call__(self, images, gt_path): """ images: list of the generated image. The values must lie between 0 and 1. gt_path: the path of the ground truth images. The values must lie between 0 and 1. """ if not os.path.exists(gt_path): raise RuntimeError('Invalid path: %s' % gt_path) print('calculate gt_path statistics...') m1, s1 = self.compute_statistics_of_path(gt_path, self.verbose) print('calculate generated_images statistics...') m2, s2 = self.calculate_activation_statistics(images, self.verbose) fid_value = self.calculate_frechet_distance(m1, s1, m2, s2) return fid_value def calculate_from_disk(self, generated_path, gt_path, img_size): """ """ if not os.path.exists(gt_path): raise RuntimeError('Invalid path: %s' % gt_path) if not os.path.exists(generated_path): raise RuntimeError('Invalid path: %s' % generated_path) print ('exp-path - '+generated_path) print('calculate gt_path statistics...') m1, s1 = self.compute_statistics_of_path(gt_path, self.verbose, img_size) print('calculate generated_path statistics...') m2, s2 = self.compute_statistics_of_path(generated_path, self.verbose, img_size) print('calculate frechet distance...') fid_value = self.calculate_frechet_distance(m1, s1, m2, s2) print('fid_distance %f' % (fid_value)) return fid_value def compute_statistics_of_path(self, path , verbose, img_size): size_flag = '{}_{}'.format(img_size[0], img_size[1]) npz_file = os.path.join(path, size_flag + '_statistics.npz') if os.path.exists(npz_file): f = np.load(npz_file) m, s = f['mu'][:], f['sigma'][:] f.close() else: path = pathlib.Path(path) files = list(path.glob('*.jpg')) + list(path.glob('*.png')) imgs = (np.array([(cv2.resize(imread(str(fn)).astype(np.float32),img_size,interpolation=cv2.INTER_CUBIC)) for fn in files]))/255.0 # Bring images to shape (B, 3, H, W) imgs = imgs.transpose((0, 3, 1, 2)) # Rescale images to be between 0 and 1 m, s = self.calculate_activation_statistics(imgs, verbose) np.savez(npz_file, mu=m, sigma=s) return m, s def calculate_activation_statistics(self, images, verbose): """Calculation of the statistics used by the FID. Params: -- images : Numpy array of dimension (n_images, 3, hi, wi). The values must lie between 0 and 1. -- model : Instance of inception model -- batch_size : The images numpy array is split into batches with batch size batch_size. A reasonable batch size depends on the hardware. -- dims : Dimensionality of features returned by Inception -- cuda : If set to True, use GPU -- verbose : If set to True and parameter out_step is given, the number of calculated batches is reported. Returns: -- mu : The mean over samples of the activations of the pool_3 layer of the inception model. -- sigma : The covariance matrix of the activations of the pool_3 layer of the inception model. """ act = self.get_activations(images, verbose) mu = np.mean(act, axis=0) sigma = np.cov(act, rowvar=False) return mu, sigma def get_activations(self, images, verbose=False): """Calculates the activations of the pool_3 layer for all images. Params: -- images : Numpy array of dimension (n_images, 3, hi, wi). The values must lie between 0 and 1. -- model : Instance of inception model -- batch_size : the images numpy array is split into batches with batch size batch_size. A reasonable batch size depends on the hardware. -- dims : Dimensionality of features returned by Inception -- cuda : If set to True, use GPU -- verbose : If set to True and parameter out_step is given, the number of calculated batches is reported. Returns: -- A numpy array of dimension (num images, dims) that contains the activations of the given tensor when feeding inception with the query tensor. """ self.model.eval() d0 = images.shape[0] if self.batch_size > d0: print(('Warning: batch size is bigger than the data size. ' 'Setting batch size to data size')) self.batch_size = d0 n_batches = d0 // self.batch_size n_used_imgs = n_batches * self.batch_size pred_arr = np.empty((n_used_imgs, self.dims)) for i in range(n_batches): if verbose: print('\rPropagating batch %d/%d' % (i + 1, n_batches)) # end='', flush=True) start = i * self.batch_size end = start + self.batch_size batch = torch.from_numpy(images[start:end]).type(torch.FloatTensor) # batch = Variable(batch, volatile=True) if self.cuda: batch = batch.cuda() pred = self.model(batch)[0] # If model output is not scalar, apply global spatial average pooling. # This happens if you choose a dimensionality not equal 2048. if pred.shape[2] != 1 or pred.shape[3] != 1: pred = adaptive_avg_pool2d(pred, output_size=(1, 1)) pred_arr[start:end] = pred.cpu().data.numpy().reshape(self.batch_size, -1) if verbose: print(' done') return pred_arr def calculate_frechet_distance(self, 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 a layer of the inception net (like returned by the function 'get_predictions') for generated samples. -- mu2 : The sample mean over activations, precalculated on an representive data set. -- sigma1: The covariance matrix over activations for generated samples. -- sigma2: The covariance matrix over activations, precalculated 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 print(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) class Reconstruction_Metrics(): def __init__(self, metric_list=['ssim', 'psnr', 'l1', 'mae'], data_range=1, win_size=51, multichannel=True): self.data_range = data_range self.win_size = win_size self.multichannel = multichannel for metric in metric_list: if metric in ['ssim', 'psnr', 'l1', 'mae']: setattr(self, metric, True) else: print('unsupport reconstruction metric: %s'%metric) def __call__(self, inputs, gts): """ inputs: the generated image, size (b,c,w,h), data range(0, data_range) gts: the ground-truth image, size (b,c,w,h), data range(0, data_range) """ result = dict() [b,n,w,h] = inputs.size() inputs = inputs.view(b*n, w, h).detach().cpu().numpy().astype(np.float32).transpose(1,2,0) gts = gts.view(b*n, w, h).detach().cpu().numpy().astype(np.float32).transpose(1,2,0) if hasattr(self, 'ssim'): ssim_value = compare_ssim(inputs, gts, data_range=self.data_range, win_size=self.win_size, multichannel=self.multichannel) result['ssim'] = ssim_value if hasattr(self, 'psnr'): psnr_value = compare_psnr(inputs, gts, self.data_range) result['psnr'] = psnr_value if hasattr(self, 'l1'): l1_value = compare_l1(inputs, gts) result['l1'] = l1_value if hasattr(self, 'mae'): mae_value = compare_mae(inputs, gts) result['mae'] = mae_value return result def calculate_from_disk(self, inputs, gts, save_path=None, img_size=(176,256), sort=True, debug=0): """ inputs: .txt files, floders, image files (string), image files (list) gts: .txt files, floders, image files (string), image files (list) """ if sort: input_image_list = sorted(get_image_list(inputs)) gt_image_list = sorted(get_image_list(gts)) else: input_image_list = get_image_list(inputs) gt_image_list = get_image_list(gts) size_flag = '{}_{}'.format(img_size[0], img_size[1]) npz_file = os.path.join(save_path, size_flag + '_metrics.npz') if os.path.exists(npz_file): f = np.load(npz_file) psnr,ssim,ssim_256,mae,l1=f['psnr'],f['ssim'],f['ssim_256'],f['mae'],f['l1'] else: psnr = [] ssim = [] ssim_256 = [] mae = [] l1 = [] names = [] for index in range(len(input_image_list)): name = os.path.basename(input_image_list[index]) names.append(name) img_gt = (cv2.resize(imread(str(gt_image_list[index])).astype(np.float32), img_size,interpolation=cv2.INTER_CUBIC)) /255.0 img_pred = (cv2.resize(imread(str(input_image_list[index])).astype(np.float32), img_size,interpolation=cv2.INTER_CUBIC)) / 255.0 if debug != 0: plt.subplot('121') plt.imshow(img_gt) plt.title('Groud truth') plt.subplot('122') plt.imshow(img_pred) plt.title('Output') plt.show() psnr.append(compare_psnr(img_gt, img_pred, data_range=self.data_range)) ssim.append(compare_ssim(img_gt, img_pred, data_range=self.data_range, win_size=self.win_size,multichannel=self.multichannel, channel_axis=2)) mae.append(compare_mae(img_gt, img_pred)) l1.append(compare_l1(img_gt, img_pred)) img_gt_256 = img_gt*255.0 img_pred_256 = img_pred*255.0 ssim_256.append(compare_ssim(img_gt_256, img_pred_256, gaussian_weights=True, sigma=1.2, use_sample_covariance=False, multichannel=True, channel_axis=2, data_range=img_pred_256.max() - img_pred_256.min())) if np.mod(index, 200) == 0: print( str(index) + ' images processed', "PSNR: %.4f" % round(np.mean(psnr), 4), "SSIM_256: %.4f" % round(np.mean(ssim_256), 4), "MAE: %.4f" % round(np.mean(mae), 4), "l1: %.4f" % round(np.mean(l1), 4), ) if save_path: np.savez(save_path + '/' + size_flag + '_metrics.npz', psnr=psnr, ssim=ssim, ssim_256=ssim_256, mae=mae, l1=l1, names=names) print( "PSNR: %.4f" % round(np.mean(psnr), 4), "PSNR Variance: %.4f" % round(np.var(psnr), 4), "SSIM_256: %.4f" % round(np.mean(ssim_256), 4), "SSIM_256 Variance: %.4f" % round(np.var(ssim_256), 4), "MAE: %.4f" % round(np.mean(mae), 4), "MAE Variance: %.4f" % round(np.var(mae), 4), "l1: %.4f" % round(np.mean(l1), 4), "l1 Variance: %.4f" % round(np.var(l1), 4) ) dic = {"psnr":[round(np.mean(psnr), 6)], "psnr_variance": [round(np.var(psnr), 6)], "ssim_256": [round(np.mean(ssim_256), 6)], "ssim_256_variance": [round(np.var(ssim_256), 6)], "mae": [round(np.mean(mae), 6)], "mae_variance": [round(np.var(mae), 6)], "l1": [round(np.mean(l1), 6)], "l1_variance": [round(np.var(l1), 6)] } return dic def get_image_list(flist): if isinstance(flist, list): return flist # flist: image file path, image directory path, text file flist path if isinstance(flist, str): if os.path.isdir(flist): flist = list(glob.glob(flist + '/*.jpg')) + list(glob.glob(flist + '/*.png')) flist.sort() return flist if os.path.isfile(flist): try: return np.genfromtxt(flist, dtype=np.str) except: return [flist] print('can not read files from %s return empty list'%flist) return [] def compare_l1(img_true, img_test): img_true = img_true.astype(np.float32) img_test = img_test.astype(np.float32) return np.mean(np.abs(img_true - img_test)) def compare_mae(img_true, img_test): img_true = img_true.astype(np.float32) img_test = img_test.astype(np.float32) return np.sum(np.abs(img_true - img_test)) / np.sum(img_true + img_test) def preprocess_path_for_deform_task(gt_path, distorted_path): distorted_image_list = sorted(get_image_list(distorted_path)) gt_list=[] distorated_list=[] for distorted_image in distorted_image_list: image = os.path.basename(distorted_image)[1:] image = image.split('_to_')[-1] gt_image = gt_path + '/' + image.replace('jpg', 'png') if not os.path.isfile(gt_image): print(distorted_image, gt_image) print('=====') continue gt_list.append(gt_image) distorated_list.append(distorted_image) return gt_list, distorated_list class LPIPS(): def __init__(self, use_gpu=True): self.model = lpips.LPIPS(net='alex').eval().cuda() self.use_gpu=use_gpu def __call__(self, image_1, image_2): """ image_1: images with size (n, 3, w, h) with value [-1, 1] image_2: images with size (n, 3, w, h) with value [-1, 1] """ result = self.model.forward(image_1, image_2) return result def calculate_from_disk(self, path_1, path_2,img_size, batch_size=64, verbose=False, sort=True): if sort: files_1 = sorted(get_image_list(path_1)) files_2 = sorted(get_image_list(path_2)) else: files_1 = get_image_list(path_1) files_2 = get_image_list(path_2) results=[] d0 = len(files_1) if batch_size > d0: print(('Warning: batch size is bigger than the data size. ' 'Setting batch size to data size')) batch_size = d0 n_batches = d0 // batch_size for i in range(n_batches): if verbose: print('\rPropagating batch %d/%d' % (i + 1, n_batches)) # end='', flush=True) start = i * batch_size end = start + batch_size 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]]) 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]]) imgs_1 = imgs_1.transpose((0, 3, 1, 2)) imgs_2 = imgs_2.transpose((0, 3, 1, 2)) img_1_batch = torch.from_numpy(imgs_1).type(torch.FloatTensor) img_2_batch = torch.from_numpy(imgs_2).type(torch.FloatTensor) if self.use_gpu: img_1_batch = img_1_batch.cuda() img_2_batch = img_2_batch.cuda() with torch.no_grad(): result = self.model.forward(img_1_batch, img_2_batch) results.append(result) distance = torch.cat(results,0)[:,0,0,0].mean() print('lpips: %.3f'%distance) return distance