ssim和psnr计算

pytorch , , 评论

SSIM和PSNR计算

在深度学习的test阶段，需要对最终的结果与GT计算结构相似度SSIM与峰值信噪比PSNR

这里有两种计算方法，第一种是用skimage库来计算，另一种手动计算，手动计算方法计算SSIM有点问题，这里建议采用第一种，毕竟是官方库，肯定是没有错的

一、skimage

首先用imageio.imread读入图片，这里读入的为HWC uint8 255类型，注意如果计算SSIM的话是需要提供multichannel=True的，这个参数要求最后一维为通道数，不然会报错，但是对HW的顺序无要求

multichannel : If True, treat the last dimension of the array as channels. Similarity calculations are done independently for each channel then averaged.

使用方法如下

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
from skimage.metrics import structural_similarity as SSIM
from skimage.metrics import peak_signal_noise_ratio as PSNR
from imageio import imread
import numpy as np

img1 = imread('1.jpg')
img2 = imread('2.jpg')

img2 = np.resize(img2, (img1.shape[0], img1.shape[1], img1.shape[2]))

ssim = SSIM(img1, img2, multichannel=True)
psnr = PSNR(img1, img2)
print(ssim)
print(psnr)
'''
输出ssim以及psnr的值
'''

二、非官方库

下面直接给出带代码，新建metrics.py文件，然后在需要测试的文件中from metrics import PSNR, SSIM即可

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
import torch
import torch.nn.functional as F
from torch.autograd import Variable
import numpy as np
import math


def gaussian(window_size, sigma):
    gauss = torch.Tensor([math.exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])
    return gauss/gauss.sum()

def create_window(window_size, channel):
    _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
    _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
    window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous())
    return window

def _ssim(img1, img2, window, window_size, channel, size_average = True):
    mu1 = F.conv2d(img1, window, padding = window_size//2, groups = channel)
    mu2 = F.conv2d(img2, window, padding = window_size//2, groups = channel)

    mu1_sq = mu1.pow(2)
    mu2_sq = mu2.pow(2)
    mu1_mu2 = mu1*mu2

    sigma1_sq = F.conv2d(img1*img1, window, padding = window_size//2, groups = channel) - mu1_sq
    sigma2_sq = F.conv2d(img2*img2, window, padding = window_size//2, groups = channel) - mu2_sq
    sigma12 = F.conv2d(img1*img2, window, padding = window_size//2, groups = channel) - mu1_mu2

    C1 = 0.01**2
    C2 = 0.03**2

    ssim_map = ((2*mu1_mu2 + C1)*(2*sigma12 + C2))/((mu1_sq + mu2_sq + C1)*(sigma1_sq + sigma2_sq + C2))

    if size_average:
        return ssim_map.mean()
    else:
        return ssim_map.mean(1).mean(1).mean(1)

class SSIM(torch.nn.Module):
    def __init__(self, window_size = 11, size_average = True):
        super(SSIM, self).__init__()
        self.window_size = window_size
        self.size_average = size_average
        self.channel = 1
        self.window = create_window(window_size, self.channel)

    def forward(self, img1, img2):
        (_, channel, _, _) = img1.size()

        if channel == self.channel and self.window.data.type() == img1.data.type():
            window = self.window
        else:
            window = create_window(self.window_size, channel)

            if img1.is_cuda:
                window = window.cuda(img1.get_device())
            window = window.type_as(img1)

            self.window = window
            self.channel = channel


        return _ssim(img1, img2, window, self.window_size, channel, self.size_average)

def ssim(img1, img2, window_size = 11, size_average = True):
    (_, channel, _, _) = img1.size()
    window = create_window(window_size, channel)

    if img1.is_cuda:
        window = window.cuda(img1.get_device())
    window = window.type_as(img1)

    return _ssim(img1, img2, window, window_size, channel, size_average)


def SSIM(img1,img2):
	# pdb.set_trace()
	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()
	ssim_value = float(ssim(img1, img2))
	print(ssim_value)
	return ssim_value

def PSNR(img1, img2):
    # pdb.set_trace()
    img1 = np.float64(img1)
    img2 = np.float64(img2)
    mse = np.mean( (img1 - img2) ** 2 )
    if mse == 0:
         return 100
    PIXEL_MAX = 255.0
    # PIXEL_MAX = 1.0
    psnr = 20 * math.log10(PIXEL_MAX / math.sqrt(mse))
    print(psnr)
    return psnr

harryteaustc

computer vision