在图像质量评估上,有三个重要指标: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实现方法,计算结果差别不大。欢迎自取~

Logo

GitCode 天启AI是一款由 GitCode 团队打造的智能助手,基于先进的LLM(大语言模型)与多智能体 Agent 技术构建,致力于为用户提供高效、智能、多模态的创作与开发支持。它不仅支持自然语言对话,还具备处理文件、生成 PPT、撰写分析报告、开发 Web 应用等多项能力,真正做到“一句话,让 Al帮你完成复杂任务”。

更多推荐