Added necessary code changes and merged with repo https://github.com/XPixelGroup/BasicSR

.eggs/README.txt

@@ -0,0 +1,6 @@
+This directory contains eggs that were downloaded by setuptools to build, test, and run plug-ins.
+
+This directory caches those eggs to prevent repeated downloads.
+
+However, it is safe to delete this directory.
+

.gitignore

@@ -6,4 +6,6 @@ logs/
 *__pycache__*
 *.sh
 datasets
-basicsr.egg-info
+basicsr.egg-info
+.eggs/*
+build/*

basicsr/metrics/__init__.py

@@ -1,10 +1,20 @@
-# ------------------------------------------------------------------------
-# Copyright (c) 2022 megvii-model. All Rights Reserved.
-# ------------------------------------------------------------------------
-# Modified from BasicSR (https://github.com/xinntao/BasicSR)
-# Copyright 2018-2020 BasicSR Authors
-# ------------------------------------------------------------------------
+from copy import deepcopy
+
+from basicsr.utils.registry import METRIC_REGISTRY
 from .niqe import calculate_niqe
-from .psnr_ssim import calculate_psnr, calculate_ssim, calculate_ssim_left, calculate_psnr_left, calculate_skimage_ssim, calculate_skimage_ssim_left
+from .psnr_ssim import calculate_psnr, calculate_ssim
+
+__all__ = ['calculate_psnr', 'calculate_ssim', 'calculate_niqe']
+
+
+def calculate_metric(data, opt):
+ """Calculate metric from data and options.
 
-__all__ = ['calculate_psnr', 'calculate_ssim', 'calculate_niqe', 'calculate_ssim_left', 'calculate_psnr_left', 'calculate_skimage_ssim', 'calculate_skimage_ssim_left']
+ Args:
+ opt (dict): Configuration. It must contain:
+ type (str): Model type.
+ """
+ opt = deepcopy(opt)
+ metric_type = opt.pop('type')
+ metric = METRIC_REGISTRY.get(metric_type)(**data, **opt)
+ return metric

basicsr/metrics/fid.py

@@ -1,35 +1,22 @@
-# ------------------------------------------------------------------------
-# Copyright (c) 2022 megvii-model. All Rights Reserved.
-# ------------------------------------------------------------------------
-# Modified from BasicSR (https://github.com/xinntao/BasicSR)
-# Copyright 2018-2020 BasicSR Authors
-# ------------------------------------------------------------------------
 import numpy as np
 import torch
 import torch.nn as nn
 from scipy import linalg
 from tqdm import tqdm
 
-from basicsr.models.archs.inception import InceptionV3
+from basicsr.archs.inception import InceptionV3
 
 
-def load_patched_inception_v3(device='cuda',
- resize_input=True,
- normalize_input=False):
+def load_patched_inception_v3(device='cuda', resize_input=True, normalize_input=False):
  # we may not resize the input, but in [rosinality/stylegan2-pytorch] it
  # does resize the input.
- inception = InceptionV3([3],
- resize_input=resize_input,
- normalize_input=normalize_input)
+ inception = InceptionV3([3], resize_input=resize_input, normalize_input=normalize_input)
  inception = nn.DataParallel(inception).eval().to(device)
  return inception
 
 
 @torch.no_grad()
-def extract_inception_features(data_generator,
- inception,
- len_generator=None,
- device='cuda'):
+def extract_inception_features(data_generator, inception, len_generator=None, device='cuda'):
  """Extract inception features.
 
  Args:
@@ -63,33 +50,27 @@ def extract_inception_features(data_generator,
 def calculate_fid(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)).
+ 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.
 
  Args:
  mu1 (np.array): The sample mean over activations.
- sigma1 (np.array): The covariance matrix over activations for
- generated samples.
- mu2 (np.array): The sample mean over activations, precalculated on an
- representative data set.
- sigma2 (np.array): The covariance matrix over activations,
- precalculated on an representative data set.
+ sigma1 (np.array): The covariance matrix over activations for generated samples.
+ mu2 (np.array): The sample mean over activations, precalculated on an representative data set.
+ sigma2 (np.array): The covariance matrix over activations, precalculated on an representative data set.
 
  Returns:
  float: The Frechet Distance.
  """
  assert mu1.shape == mu2.shape, 'Two mean vectors have different lengths'
- assert sigma1.shape == sigma2.shape, (
- 'Two covariances have different dimensions')
+ assert sigma1.shape == sigma2.shape, ('Two covariances have different dimensions')
 
  cov_sqrt, _ = linalg.sqrtm(sigma1 @ sigma2, disp=False)
 
  # Product might be almost singular
  if not np.isfinite(cov_sqrt).all():
- print('Product of cov matrices is singular. Adding {eps} to diagonal '
- 'of cov estimates')
+ print('Product of cov matrices is singular. Adding {eps} to diagonal of cov estimates')
  offset = np.eye(sigma1.shape[0]) * eps
  cov_sqrt = linalg.sqrtm((sigma1 + offset) @ (sigma2 + offset))
 

basicsr/metrics/metric_util.py

@@ -1,12 +1,6 @@
-# ------------------------------------------------------------------------
-# Copyright (c) 2022 megvii-model. All Rights Reserved.
-# ------------------------------------------------------------------------
-# Modified from BasicSR (https://github.com/xinntao/BasicSR)
-# Copyright 2018-2020 BasicSR Authors
-# ------------------------------------------------------------------------
 import numpy as np
 
-from basicsr.utils.matlab_functions import bgr2ycbcr
+from basicsr.utils import bgr2ycbcr
 
 
 def reorder_image(img, input_order='HWC'):
@@ -27,9 +21,7 @@ def reorder_image(img, input_order='HWC'):
  """
 
  if input_order not in ['HWC', 'CHW']:
- raise ValueError(
- f'Wrong input_order {input_order}. Supported input_orders are '
- "'HWC' and 'CHW'")
+ raise ValueError(f"Wrong input_order {input_order}. Supported input_orders are 'HWC' and 'CHW'")
  if len(img.shape) == 2:
  img = img[..., None]
  if input_order == 'CHW':

basicsr/metrics/niqe.py

@@ -1,20 +1,17 @@
-# ------------------------------------------------------------------------
-# Copyright (c) 2022 megvii-model. All Rights Reserved.
-# ------------------------------------------------------------------------
-# Modified from BasicSR (https://github.com/xinntao/BasicSR)
-# Copyright 2018-2020 BasicSR Authors
-# ------------------------------------------------------------------------
 import cv2
 import math
 import numpy as np
-from scipy.ndimage.filters import convolve
+import os
+from scipy.ndimage import convolve
 from scipy.special import gamma
 
 from basicsr.metrics.metric_util import reorder_image, to_y_channel
+from basicsr.utils.matlab_functions import imresize
+from basicsr.utils.registry import METRIC_REGISTRY
 
 
 def estimate_aggd_param(block):
- """Estimate AGGD (Asymmetric Generalized Gaussian Distribution) paramters.
+ """Estimate AGGD (Asymmetric Generalized Gaussian Distribution) parameters.
 
  Args:
  block (ndarray): 2D Image block.
@@ -26,15 +23,13 @@ def estimate_aggd_param(block):
  block = block.flatten()
  gam = np.arange(0.2, 10.001, 0.001) # len = 9801
  gam_reciprocal = np.reciprocal(gam)
- r_gam = np.square(gamma(gam_reciprocal * 2)) / (
- gamma(gam_reciprocal) * gamma(gam_reciprocal * 3))
+ r_gam = np.square(gamma(gam_reciprocal * 2)) / (gamma(gam_reciprocal) * gamma(gam_reciprocal * 3))
 
  left_std = np.sqrt(np.mean(block[block < 0]**2))
  right_std = np.sqrt(np.mean(block[block > 0]**2))
  gammahat = left_std / right_std
  rhat = (np.mean(np.abs(block)))**2 / np.mean(block**2)
- rhatnorm = (rhat * (gammahat**3 + 1) *
- (gammahat + 1)) / ((gammahat**2 + 1)**2)
+ rhatnorm = (rhat * (gammahat**3 + 1) * (gammahat + 1)) / ((gammahat**2 + 1)**2)
  array_position = np.argmin((r_gam - rhatnorm)**2)
 
  alpha = gam[array_position]
@@ -70,22 +65,18 @@ def compute_feature(block):
  return feat
 
 
-def niqe(img,
- mu_pris_param,
- cov_pris_param,
- gaussian_window,
- block_size_h=96,
- block_size_w=96):
+def niqe(img, mu_pris_param, cov_pris_param, gaussian_window, block_size_h=96, block_size_w=96):
  """Calculate NIQE (Natural Image Quality Evaluator) metric.
 
- Ref: Making a "Completely Blind" Image Quality Analyzer.
+ ``Paper: Making a "Completely Blind" Image Quality Analyzer``
+
  This implementation could produce almost the same results as the official
  MATLAB codes: https://live.ece.utexas.edu/research/quality/niqe_release.zip
 
  Note that we do not include block overlap height and width, since they are
  always 0 in the official implementation.
 
- For good performance, it is advisable by the official implemtation to
+ For good performance, it is advisable by the official implementation to
  divide the distorted image in to the same size patched as used for the
  construction of multivariate Gaussian model.
 
@@ -104,8 +95,7 @@ def niqe(img,
  block_size_w (int): Width of the blocks in to which image is divided.
  Default: 96 (the official recommended value).
  """
- assert img.ndim == 2, (
- 'Input image must be a gray or Y (of YCbCr) image with shape (h, w).')
+ assert img.ndim == 2, ('Input image must be a gray or Y (of YCbCr) image with shape (h, w).')
  # crop image
  h, w = img.shape
  num_block_h = math.floor(h / block_size_h)
@@ -115,32 +105,22 @@ def niqe(img,
  distparam = [] # dist param is actually the multiscale features
  for scale in (1, 2): # perform on two scales (1, 2)
  mu = convolve(img, gaussian_window, mode='nearest')
- sigma = np.sqrt(
- np.abs(
- convolve(np.square(img), gaussian_window, mode='nearest') -
- np.square(mu)))
+ sigma = np.sqrt(np.abs(convolve(np.square(img), gaussian_window, mode='nearest') - np.square(mu)))
  # normalize, as in Eq. 1 in the paper
  img_nomalized = (img - mu) / (sigma + 1)
 
  feat = []
  for idx_w in range(num_block_w):
  for idx_h in range(num_block_h):
  # process ecah block
- block = img_nomalized[idx_h * block_size_h //
- scale:(idx_h + 1) * block_size_h //
- scale, idx_w * block_size_w //
- scale:(idx_w + 1) * block_size_w //
- scale]
+ block = img_nomalized[idx_h * block_size_h // scale:(idx_h + 1) * block_size_h // scale,
+ idx_w * block_size_w // scale:(idx_w + 1) * block_size_w // scale]
  feat.append(compute_feature(block))
 
  distparam.append(np.array(feat))
- # TODO: matlab bicubic downsample with anti-aliasing
- # for simplicity, now we use opencv instead, which will result in
- # a slight difference.
+
  if scale == 1:
- h, w = img.shape
- img = cv2.resize(
- img / 255., (w // 2, h // 2), interpolation=cv2.INTER_LINEAR)
+ img = imresize(img / 255., scale=0.5, antialiasing=True)
  img = img * 255.
 
  distparam = np.concatenate(distparam, axis=1)
@@ -154,20 +134,25 @@ def niqe(img,
  # compute niqe quality, Eq. 10 in the paper
  invcov_param = np.linalg.pinv((cov_pris_param + cov_distparam) / 2)
  quality = np.matmul(
- np.matmul((mu_pris_param - mu_distparam), invcov_param),
- np.transpose((mu_pris_param - mu_distparam)))
- quality = np.sqrt(quality)
+ np.matmul((mu_pris_param - mu_distparam), invcov_param), np.transpose((mu_pris_param - mu_distparam)))
 
+ quality = np.sqrt(quality)
+ quality = float(np.squeeze(quality))
  return quality
 
 
-def calculate_niqe(img, crop_border, input_order='HWC', convert_to='y'):
+@METRIC_REGISTRY.register()
+def calculate_niqe(img, crop_border, input_order='HWC', convert_to='y', **kwargs):
  """Calculate NIQE (Natural Image Quality Evaluator) metric.
 
- Ref: Making a "Completely Blind" Image Quality Analyzer.
+ ``Paper: Making a "Completely Blind" Image Quality Analyzer``
+
  This implementation could produce almost the same results as the official
  MATLAB codes: https://live.ece.utexas.edu/research/quality/niqe_release.zip
 
+ > MATLAB R2021a result for tests/data/baboon.png: 5.72957338 (5.7296)
+ > Our re-implementation result for tests/data/baboon.png: 5.7295763 (5.7296)
+
  We use the official params estimated from the pristine dataset.
  We use the recommended block size (96, 96) without overlaps.
 
@@ -181,15 +166,15 @@ def calculate_niqe(img, crop_border, input_order='HWC', convert_to='y'):
  pixels are not involved in the metric calculation.
  input_order (str): Whether the input order is 'HW', 'HWC' or 'CHW'.
  Default: 'HWC'.
- convert_to (str): Whether coverted to 'y' (of MATLAB YCbCr) or 'gray'.
+ convert_to (str): Whether converted to 'y' (of MATLAB YCbCr) or 'gray'.
  Default: 'y'.
 
  Returns:
  float: NIQE result.
  """
-
+ ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
  # we use the official params estimated from the pristine dataset.
- niqe_pris_params = np.load('basicsr/metrics/niqe_pris_params.npz')
+ niqe_pris_params = np.load(os.path.join(ROOT_DIR, 'niqe_pris_params.npz'))
  mu_pris_param = niqe_pris_params['mu_pris_param']
  cov_pris_param = niqe_pris_params['cov_pris_param']
  gaussian_window = niqe_pris_params['gaussian_window']
@@ -206,6 +191,9 @@ def calculate_niqe(img, crop_border, input_order='HWC', convert_to='y'):
  if crop_border != 0:
  img = img[crop_border:-crop_border, crop_border:-crop_border]
 
+ # round is necessary for being consistent with MATLAB's result
+ img = img.round()
+
  niqe_result = niqe(img, mu_pris_param, cov_pris_param, gaussian_window)
 
  return niqe_result