欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页

Python计算图片SSIM和PSNR

程序员文章站 2023-12-31 13:28:58
...

分两种情况:

1. 在网络训练过程中计算Output和Groundtruth之间的SSIM,作为损失函数;

2. 直接计算两张图片之间的SSIM。

 

情况1

https://github.com/congyucn/pytorch-ssim

可直接使用上述代码。

 

情况2

参考上述代码,可将上述代码更改如下:

def ssim(img1,img2):
    img1 = torch.from_numpy(np.rollaxis(img1, 2)).float().unsqueeze(0)/255.0
    img2 = torch.from_numpy(np.rollaxis(img2, 2)).float().unsqueeze(0)/255.0   
    img1 = Variable( img1,  requires_grad=False)    # torch.Size([256, 256, 3])
    img2 = Variable( img2, requires_grad = False)
    ssim_value = pytorch_ssim.ssim(img1, img2).item()
    return ssim_value

而计算psnr的代码见:https://blog.csdn.net/qazwsxrx/article/details/104550550

总体而言,计算两张图片的psnr和ssim的代码如下所示:

import numpy 
import numpy as np
import math
import cv2
import torch
import pytorch_ssim
from torch.autograd import Variable

original = cv2.imread("1.png")      # numpy.adarray
contrast = cv2.imread("2.png",1)

def psnr(img1, img2):
    mse = numpy.mean( (img1 - img2) ** 2 )
    if mse == 0:
        return 100
    PIXEL_MAX = 255.0
    return 20 * math.log10(PIXEL_MAX / math.sqrt(mse))

def ssim(img1,img2):
    img1 = torch.from_numpy(np.rollaxis(img1, 2)).float().unsqueeze(0)/255.0
    img2 = torch.from_numpy(np.rollaxis(img2, 2)).float().unsqueeze(0)/255.0   
    img1 = Variable( img1,  requires_grad=False)    # torch.Size([256, 256, 3])
    img2 = Variable( img2, requires_grad = False)
    ssim_value = pytorch_ssim.ssim(img1, img2).item()
    return ssim_value

psnrValue = psnr(original,contrast)
ssimValue = ssim(original,contrast)
print(psnrValue)
print(ssimValue)

 

上一篇:

下一篇: