基于pytorch 的psnr和ssim计算_torch_msssim

基于pytorch 的psnr和ssim计算：

sr，hr是 NCHW格式的tensor, uint8 0-255

import math
import numpy as np
from pytorch_msssim import ssim, ms_ssim, SSIM, MS_SSIM
import datetime
import os
import sys

def calc_psnr(sr, hr):
    diff = (sr - hr) / 255.00
    mse  = diff.pow(2).mean()
    psnr = -10 * math.log10(mse)                    
    return float(psnr)

def calc_ssim(sr, hr):
    # def ssim(
    #     X,
    #     Y,
    #     data_range=255,
    #     size_average=True,
    #     win_size=11,
    #     win_sigma=1.5,
    #     win=None,
    #     K=(0.01, 0.03),
    #     nonnegative_ssim=False,
    # )
    ssim_val = ssim(sr, hr, data_range=255, size_average=True)
    return float(ssim_val)

def cur_timestamp_str():#2024-0103-1121
    now = datetime.datetime.now()
    year = str(now.year)
    month = str(now.month).zfill(2)
    day = str(now.day).zfill(2)
    hour = str(now.hour).zfill(2)
    minute = str(now.minute).zfill(2)

    content = "{}-{}{}-{}{}".format(year, month, day, hour, minute)
    return content
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

关于ssim的计算，以及作为损失函数介绍，参考：
https://pypi.org/project/pytorch-msssim/#description