Skip to content

Different SSIM metrics in CNN-based super resolution algorithms (e.g., EDSR CVPRW2017, RDN CVPR2018, MSRN ECCV2018).

License

Notifications You must be signed in to change notification settings

HolmesShuan/EDSR-ssim

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

13 Commits
 
 
 
 

Repository files navigation

EDSR-ssim

There are different SSIM metrics used in popular CNN-based super resolution algorithms, such as EDSR, RDN and MSRN. In this project, we re-implement those metrics, from MATLAB to Python.

How to use ?

  1. Please refer to the PSNR setting proposed in EDSR. Then, one may edit utility.py as follow:
'''
def calc_psnr(sr, hr, scale, rgb_range, dataset=None):
    if hr.nelement() == 1: return 0

    diff = (sr - hr) / rgb_range
    if dataset and dataset.dataset.benchmark:
        shave = scale
        if diff.size(1) > 1:
            gray_coeffs = [65.738, 129.057, 25.064]
            convert = diff.new_tensor(gray_coeffs).view(1, 3, 1, 1) / 256
            diff = diff.mul(convert).sum(dim=1)
    else:
        shave = scale + 6

    valid = diff[..., shave:-shave, shave:-shave]
    mse = valid.pow(2).mean()

    return -10 * math.log10(mse)
'''
import numpy as np
from scipy import signal
from skimage.measure import compare_ssim

def matlab_style_gauss2D(shape=(3,3),sigma=0.5):
  """
  2D gaussian mask - should give the same result as MATLAB's fspecial('gaussian',[shape],[sigma])
  Acknowledgement : https://stackoverflow.com/questions/17190649/how-to-obtain-a-gaussian-filter-in-python (Author@ali_m)
  """
  m,n = [(ss-1.)/2. for ss in shape]
  y,x = np.ogrid[-m:m+1,-n:n+1]
  h = np.exp( -(x*x + y*y) / (2.*sigma*sigma) )
  h[ h < np.finfo(h.dtype).eps*h.max() ] = 0
  sumh = h.sum()
  if sumh != 0:
    h /= sumh
  return h

def calc_ssim(X, Y, scale, rgb_range, dataset=None, sigma=1.5, K1=0.01, K2=0.03, R=255):
  '''
  X : y channel (i.e., luminance) of transformed YCbCr space of X
  Y : y channel (i.e., luminance) of transformed YCbCr space of Y
  Please follow the setting of psnr_ssim.m in EDSR (Enhanced Deep Residual Networks for Single Image Super-Resolution CVPRW2017).
  Official Link : https://github.com/LimBee/NTIRE2017/tree/db34606c2844e89317aac8728a2de562ef1f8aba
  The authors of EDSR use MATLAB's ssim as the evaluation tool, 
  thus this function is the same as ssim.m in MATLAB with C(3) == C(2)/2. 
  '''
  gaussian_filter = matlab_style_gauss2D((11, 11), sigma)

  if dataset and dataset.dataset.benchmark:
    shave = scale
    if X.size(1) > 1:
        gray_coeffs = [65.738, 129.057, 25.064]
        convert = X.new_tensor(gray_coeffs).view(1, 3, 1, 1) / 256
        X = X.mul(convert).sum(dim=1)
        Y = Y.mul(convert).sum(dim=1)
  else:
    shave = scale + 6

  X = X[..., shave:-shave, shave:-shave].squeeze().cpu().numpy().astype(np.float64) 
  Y = Y[..., shave:-shave, shave:-shave].squeeze().cpu().numpy().astype(np.float64)

  window = gaussian_filter

  ux = signal.convolve2d(X, window, mode='same', boundary='symm')
  uy = signal.convolve2d(Y, window, mode='same', boundary='symm')

  uxx = signal.convolve2d(X*X, window, mode='same', boundary='symm')
  uyy = signal.convolve2d(Y*Y, window, mode='same', boundary='symm')
  uxy = signal.convolve2d(X*Y, window, mode='same', boundary='symm')

  vx = uxx - ux * ux
  vy = uyy - uy * uy
  vxy = uxy - ux * uy

  C1 = (K1 * R) ** 2
  C2 = (K2 * R) ** 2

  A1, A2, B1, B2 = ((2 * ux * uy + C1, 2 * vxy + C2, ux ** 2 + uy ** 2 + C1, vx + vy + C2))
  D = B1 * B2
  S = (A1 * A2) / D
  mssim = S.mean()

  return mssim
  1. Edit Line 93 in trainer.py as follow:
'''
self.ckp.log[-1, idx_data, idx_scale] += utility.calc_psnr(
    sr, hr, scale, self.args.rgb_range, dataset=d
)
'''
self.ckp.log[-1, idx_data, idx_scale] += utility.calc_ssim(
    sr, hr, scale, self.args.rgb_range, dataset=d
)
  1. Edit Line 105 in trainer.py as follow:
# '[{} x{}]\tPSNR: {:.3f} (Best: {:.3f} @epoch {})'.format(
'[{} x{}]\SSIM: {:.3f} (Best: {:.3f} @epoch {})'.format(

Other SSIMs:

Different SSIM metrics used in Super-Resolution papers:

import numpy as np
from scipy import signal
from skimage.measure import compare_ssim

def matlab_style_gauss2D(shape=(3,3),sigma=0.5):
  """
  2D gaussian mask - should give the same result as MATLAB's fspecial('gaussian',[shape],[sigma])
  Acknowledgement : https://stackoverflow.com/questions/17190649/how-to-obtain-a-gaussian-filter-in-python (Author@ali_m)
  """
  m,n = [(ss-1.)/2. for ss in shape]
  y,x = np.ogrid[-m:m+1,-n:n+1]
  h = np.exp( -(x*x + y*y) / (2.*sigma*sigma) )
  h[ h < np.finfo(h.dtype).eps*h.max() ] = 0
  sumh = h.sum()
  if sumh != 0:
    h /= sumh
  return h

def calc_ssim(X, Y, sigma=1.5, K1=0.01, K2=0.03, R = 255):
  '''
  X : y channel (i.e., luminance) of transformed YCbCr space of X
  Y : y channel (i.e., luminance) of transformed YCbCr space of Y
  Please follow the setting of Evaluate_PSNR_SSIM.m in RDN (Residual Dense Network for Image Super-Resolution CVPR2018).
  Official Link : https://github.com/yulunzhang/RDN
  '''
  gaussian_filter = matlab_style_gauss2D((11, 11), sigma)

  X = X.astype(np.float64)
  Y = Y.astype(np.float64)

  # Since matlab_style_gauss2D() yields normalized filter, this operation can be deprecated.
  window = gaussian_filter / np.sum(np.sum(gaussian_filter))
  
  window = np.fliplr(window)
  window = np.flipud(window)

  ux = signal.convolve2d(X, window, mode='valid', boundary='fill', fillvalue=0)
  uy = signal.convolve2d(Y, window, mode='valid', boundary='fill', fillvalue=0)

  uxx = signal.convolve2d(X*X, window, mode='valid', boundary='fill', fillvalue=0)
  uyy = signal.convolve2d(Y*Y, window, mode='valid', boundary='fill', fillvalue=0)
  uxy = signal.convolve2d(X*Y, window, mode='valid', boundary='fill', fillvalue=0)

  vx = uxx - ux * ux
  vy = uyy - uy * uy
  vxy = uxy - ux * uy

  C1 = (K1 * R) ** 2
  C2 = (K2 * R) ** 2

  A1, A2, B1, B2 = ((2 * ux * uy + C1, 2 * vxy + C2, ux ** 2 + uy ** 2 + C1, vx + vy + C2))
  D = B1 * B2
  S = (A1 * A2) / D
  mssim = S.mean()

  return mssim

def calc_ssim(X, Y, sigma=1.5, K1=0.01, K2=0.03, R=255):
  '''
  X : y channel (i.e., luminance) of transformed YCbCr space of X
  Y : y channel (i.e., luminance) of transformed YCbCr space of Y
  Please follow the setting of psnr_ssim.m in EDSR (Enhanced Deep Residual Networks for Single Image Super-Resolution CVPRW2017).
  Official Link : https://github.com/LimBee/NTIRE2017/tree/db34606c2844e89317aac8728a2de562ef1f8aba
  The authors of EDSR use MATLAB's ssim as the evaluation tool, 
  thus this function is the same as ssim.m in MATLAB with C(3) == C(2)/2. 
  '''
  gaussian_filter = matlab_style_gauss2D((11, 11), sigma)

  X = X.astype(np.float64)
  Y = Y.astype(np.float64)

  window = gaussian_filter

  ux = signal.convolve2d(X, window, mode='same', boundary='symm')
  uy = signal.convolve2d(Y, window, mode='same', boundary='symm')

  uxx = signal.convolve2d(X*X, window, mode='same', boundary='symm')
  uyy = signal.convolve2d(Y*Y, window, mode='same', boundary='symm')
  uxy = signal.convolve2d(X*Y, window, mode='same', boundary='symm')

  vx = uxx - ux * ux
  vy = uyy - uy * uy
  vxy = uxy - ux * uy

  C1 = (K1 * R) ** 2
  C2 = (K2 * R) ** 2

  A1, A2, B1, B2 = ((2 * ux * uy + C1, 2 * vxy + C2, ux ** 2 + uy ** 2 + C1, vx + vy + C2))
  D = B1 * B2
  S = (A1 * A2) / D
  mssim = S.mean()

  return mssim

def calc_ssim(X, Y):
  '''
  X (groundtruth): y channel (i.e., luminance) of transformed YCbCr space of X
  Y (prediction): y channel (i.e., luminance) of transformed YCbCr space of Y
  Please follow the setting of test.py in MSRN (Multi-scale Residual Network for Image Super-Resolution ECCV2018).
  Official Link : https://github.com/MIVRC/MSRN-PyTorch
  The authors of MSRN use scikit-image's compare_ssim as the evaluation tool, 
  note that this function is quite sensitive to the argument "data_range", emprically, the larger the higher output.
  '''
  ssim = compare_ssim(X, Y, data_range=max(Y.max(),X.max()) - min(X.min(),Y.min()) # one may obtain a slightly higher output than original setting
  # ssim = compare_ssim(X, Y, data_range=Y.max() - X.min())
  return ssim

Note that we omit the crop preprocess, which can be a game changer for super resolution problems. Please follow the exact crop setting in those papers before calling SSIM functions.

Disclaimer:

We are not responsible for better or worse performances than original results reported in those papers. Hopefully, you may find these codes helpful in your research or work.

About

Different SSIM metrics in CNN-based super resolution algorithms (e.g., EDSR CVPRW2017, RDN CVPR2018, MSRN ECCV2018).

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published