NITIRE 2023官方的PSNR及SSIM计算代码
问题描述
做图像复原任务时,总避免不了计算PSNR
和SSIM
等图像质量评估指标,但是网上实在是太多计算这类指标的代码了,不同代码计算的结果还可能存在差异。有使用matlab计算SSIM
的,也有使用python计算SSIM
的。有使用第三方库skimage
计算SSIM
的,也有使用第三方库pyiqa
计算SSIM
的。博主曾经试过,对Urban100
测试集进行SSIM
的计算,即计算SR图像与HR图像的SSIM
时,数值差异较大,使用matlab计算的SSIM
为0.9644,而使用python计算的SSIM
为0.7980。因此,这让人一头雾水,不知道哪个才是所谓的标准。
NITIRE比赛
博主关注了今年的NITIRE 2023超分赛道的比赛,留意到轻量级超分赛道NTIRE 2023 Challenge on Efficient Super-Resolution中,提供了一个代码模板,让参赛者可以自己测试模型的显存占用,推理时间,PSNR和SSIM。报告原文有句话是这么说的:
A code example for calculating these metrics is available at https://github.com/ofsoundof/NTIRE2023_ESR. The code of the submitted solutions and the pre-trained weights are also available in this repository.
计算这些指标的代码例子可以在https://github.com/ofsoundof/NTIRE2023_ESR中找到。所提交的方案的代码和预训练权重都能在这个仓库中找到。
也就是我们可以在https://github.com/ofsoundof/NTIRE2023_ESR中找到NITIRE官方计算PSNR
和SSIM
的代码,由于NITRE比赛是CVPR(IEEE Conference on Computer Vision and Pattern Recognition)举办的极具影响力的计算机视觉底层任务比赛,十分具有公信力,因此建议使用此代码进行计算。
核心代码
计算PSNR和SSIM的核心代码
'''
# =======================================
# metric, PSNR and SSIM
# =======================================
'''
# ----------
# PSNR
# ----------
def calculate_psnr(img1, img2, border=0):
img1 = _bord_img(img1)
img2 = _bord_img(img2)
return _calculate_psnr(img1, img2)
# ----------
# SSIM
# ----------
def calculate_ssim(img1, img2, border=0):
img1 = _bord_img(img1)
img2 = _bord_img(img2)
return _calculate_ssim(img1, img2)
def _calculate_ssim(img, img2, test_y_channel=True):
if test_y_channel:
img = to_y_channel(img)
img2 = to_y_channel(img2)
ssims = []
for i in range(img.shape[2]):
ssims.append(_ssim(img[..., i], img2[..., i]))
return np.array(ssims).mean()
# 在y通道上计算
def _calculate_psnr(img, img2, test_y_channel=True,):
if test_y_channel:
img = to_y_channel(img)
img2 = to_y_channel(img2)
mse = np.mean((img - img2)**2)
if mse == 0:
return float('inf')
return 20. * np.log10(255. / np.sqrt(mse))
def _ssim(img, img2):
c1 = (0.01 * 255)**2
c2 = (0.03 * 255)**2
img = img.astype(np.float64)
img2 = img2.astype(np.float64)
kernel = cv2.getGaussianKernel(11, 1.5)
window = np.outer(kernel, kernel.transpose())
mu1 = cv2.filter2D(img, -1, window)[5:-5, 5:-5]
mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
mu1_sq = mu1**2
mu2_sq = mu2**2
mu1_mu2 = mu1 * mu2
sigma1_sq = cv2.filter2D(img**2, -1, window)[5:-5, 5:-5] - mu1_sq
sigma2_sq = cv2.filter2D(img2**2, -1, window)[5:-5, 5:-5] - mu2_sq
sigma12 = cv2.filter2D(img * img2, -1, window)[5:-5, 5:-5] - mu1_mu2
ssim_map = ((2 * mu1_mu2 + c1) * (2 * sigma12 + c2)) / ((mu1_sq + mu2_sq + c1) * (sigma1_sq + sigma2_sq + c2))
return ssim_map.mean()
最后感谢小伙伴们的学习噢~
![](https://i-blog.csdnimg.cn/blog_migrate/d9064939a483b6980139b9ad9eeb24f8.png#pic_center)
最后附上2023年高效超分的比赛报告链接,欢迎大家多多阅读分享:NTIRE 2023 Challenge on Efficient Super-Resolution: Methods and Results