PSNR/SSIM/LPIPS图像质量评估三件套（含代码）

在图像质量评估上，有三个重要指标：PSNR，SSIM，LPIPS。本文提供简易脚本分别实现。

PSNR，峰值信噪比，是基于MSE的像素比较低质量评估，一般30dB以上质量就不错，到40dB以上肉眼就很难分别了。
SSIM，结构相似性，从分布上来比较相似性，量化到（0-1）之间，越接近1则证明图像质量越好。具体数学公式可以看我之前的博客《SSIM》。
LPIPS，利用AI模型来量化图像之间的相似性。取值范围也是[0,1]，与SSIM相反，LPIPS是越小则证明图像质量越好。

像这种常见的图像质量评价指标，都会收录到torchmetrics里面。只需安装：

pip install torchmetrics

实验脚本：

import torch
from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
from torchmetrics.image import StructuralSimilarityIndexMeasure
from torchmetrics.image import PeakSignalNoiseRatio

_ = torch.manual_seed(123)


def psnr_torch(img1, img2):
    mse = ((img1 - img2) ** 2).view(img1.shape[0], -1).mean(1, keepdim=True)
    return 20 * torch.log10(1.0 / torch.sqrt(mse))


def psnr(img1, img2):
    metric = PeakSignalNoiseRatio()
    return metric(img1, img2)


def ssim(img1, img2):
    metric = StructuralSimilarityIndexMeasure(data_range=1.0)
    return metric(img1, img2)


def lpips(img1, img2):
    metric = LearnedPerceptualImagePatchSimilarity(net_type='vgg')
    return metric(img1, img2)


def _main():
    img1 = torch.rand(1, 3, 100, 100)
    img2 = torch.rand(1, 3, 100, 100)

    # PSNR
    print("PNSR: ", psnr_torch(img1, img2))
    print("PNSR1: ", psnr(img1, img2))
    print("SSIM: ", ssim(img1, img2))
    print("LPIPS: ", lpips(img1, img2))


if __name__ == "__main__":
    _main()

代码里给了两种PSNR实现方法，计算结果差别不大。欢迎自取～