计算YUV文件的PSNR与SSIM

import numpy as np
from PIL import Image
import glob
import math
import cv2
import os
raw_yuv_path = r'C:\Users\weiwei\Desktop\python_ssim 计算\raw_832x480_50.yuv'
com_yuv_path = r'C:\Users\weiwei\Desktop\python_ssim 计算\com_832x480_50.yuv'
raw_write_path = r'C:\Users\weiwei\Desktop\python_ssim 计算\1.yuv'
h, w, nfs = 832, 480, 50
def calculate_psnr(img1, img2):
    # img1 and img2 have range [0, 255]
    img1 = img1.astype(np.float64)
    img2 = img2.astype(np.float64)
    mse = np.mean((img1 - img2)**2)
    if mse == 0:
        return float('inf')
    return 20 * math.log10(255.0 / math.sqrt(mse))

def ssim(img1, img2):
    C1 = (0.01 * 255)**2
    C2 = (0.03 * 255)**2

    img1 = img1.astype(np.float64)
    img2 = img2.astype(np.float64)
    kernel = cv2.getGaussianKernel(11, 1.5)
    window = np.outer(kernel, kernel.transpose())

    mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5]  # valid
    mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
    mu1_sq = mu1**2
    mu2_sq = mu2**2
    mu1_mu2 = mu1 * mu2
    sigma1_sq = cv2.filter2D(img1**2, -1, window)[5:-5, 5:-5] - mu1_sq
    sigma2_sq = cv2.filter2D(img2**2, -1, window)[5:-5, 5:-5] - mu2_sq
    sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2

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


def calculate_ssim(img1, img2):
    '''calculate SSIM
    the same outputs as MATLAB's
    img1, img2: [0, 255]
    '''
    if not img1.shape == img2.shape:
        raise ValueError('Input images must have the same dimensions.')
    if img1.ndim == 2:
        return ssim(img1, img2)
    elif img1.ndim == 3:
        if img1.shape[2] == 3:
            ssims = []
            for i in range(3):
                ssims.append(ssim(img1, img2))
            return np.array(ssims).mean()
        elif img1.shape[2] == 1:
            return ssim(np.squeeze(img1), np.squeeze(img2))
    else:
        raise ValueError('Wrong input image dimensions.')


def import_yuv(seq_path, h, w, tot_frm, yuv_type='420p', start_frm=0, only_y=False):
    if yuv_type == '420p':
        hh, ww = h // 2, w // 2
    elif yuv_type == '444p':
        hh, ww = h, w
    else:
        raise Exception('yuv_type not supported.')

    y_size, u_size, v_size = h * w, hh * ww, hh * ww
    blk_size = y_size + u_size + v_size

    # init
    y_seq = np.zeros((tot_frm, h, w), dtype=np.uint8)
    if not only_y:
        u_seq = np.zeros((tot_frm, hh, ww), dtype=np.uint8)
        v_seq = np.zeros((tot_frm, hh, ww), dtype=np.uint8)

    # read data
    with open(seq_path, 'rb') as fp:
        for i in range(tot_frm):
            fp.seek(int(blk_size * (start_frm + i)), 0)  # skip frames
            y_frm = np.fromfile(fp, dtype=np.uint8, count=y_size).reshape(h, w)
            #y_frm = np.fromfile(fp, dtype=np.uint8, count=y_size)
            #v_frm = np.fromfile(fp, dtype=np.uint8, count=v_size)
            if only_y:
                y_seq[i, ...] = y_frm
            else:
                u_frm = np.fromfile(fp, dtype=np.uint8, \
                                    count=u_size).reshape(hh, ww)
                v_frm = np.fromfile(fp, dtype=np.uint8, \
                                    count=v_size).reshape(hh, ww)
                y_seq[i, ...], u_seq[i, ...], v_seq[i, ...] = y_frm, u_frm, v_frm

    if only_y:
        return y_seq
    else:
        return y_seq, u_seq, v_seq


# Press the green button in the gutter to run the script.
if __name__ == '__main__':

    folder_GT = r'C:\Users\weiwei\Desktop\python_YUV\data\raw'
    folder_Gen = r'C:\Users\weiwei\Desktop\python_YUV\data\compressed'

    crop_border = 4
    suffix = ''  # suffix for Gen images
    #test_Y = False  # True: test Y channel only; False: test RGB channels
    test_Y = True  # True: test Y channel only; False: test RGB channels

    PSNR_all = []
    SSIM_all = []
    img_list = sorted(glob.glob(folder_GT + '/*'))


    raw_y = import_yuv(
        seq_path=raw_yuv_path, h=h, w=w, tot_frm=nfs, start_frm=0, only_y=True
        )
    #write_yuv(raw_write_path=raw_write_path, y_seq=raw_y, u_seq=raw_u, v_seq=raw_v, h=h, w=w, tot_frm=nfs, start_frm=0, only_y=False)
    com_y = import_yuv(
        seq_path=com_yuv_path, h=h, w=w, tot_frm=nfs, start_frm=0, only_y=True
        )
    for i in range(len(raw_y)):
        im_raw_y = raw_y[i]
        im_com_y = com_y[i]


        # calculate PSNR and SSIM
        PSNR = calculate_psnr(im_raw_y, im_com_y)
        SSIM = calculate_ssim(im_raw_y, im_com_y)

        print('{:3d} 帧--- \tPSNR: {:.6f} dB, \tSSIM: {:.6f}'.format(
            i + 1, PSNR, SSIM))
        PSNR_all.append(PSNR)
        SSIM_all.append(SSIM)
    print('Average: PSNR: {:.6f} dB, SSIM: {:.6f}'.format(
        sum(PSNR_all) / len(PSNR_all),
        sum(SSIM_all) / len(SSIM_all)))

    # 输出第一帧并保存
    gt_image = raw_y[[0], :, :]
    gt_image = np.reshape(gt_image, (480, 832))
    im = Image.fromarray(gt_image)
    im.save("out.jpeg")

# 如果数据在GPU上，转移到CPU上计算
ssim_gt_data = 255 * (gt_data.squeeze(0).squeeze(0)).cpu().numpy()
ssim_lq_data = 255 * ((lq_data[0, radius, ...]).squeeze(0).cpu().numpy())
sim_en_data = 255 * (enhanced_data.squeeze(0).squeeze(0)).cpu().numpy()

# 使用FFmpeg计算
ffmpeg -s 832x480 -i G:\datas\test_18\HM16.5_LDP\QP37\BasketballDrill_832x480_50.yuv -s 832x480 -i G:\datas\test_18\raw\BasketballDrill_832x480_50.yuv -lavfi ssim="stats_file=ssim.log" -f null -