视频质量评价python，包括psnr，ssim，LPIPS

文件准备：

创建两个文件夹，分别叫做video_gen，video_real。

video_gen下放置生成的视频的帧，以1.jpg样式命名。

video_real下放置真值，命名方式同上。

评价代码：

"""
Video Quality Metrics
Copyright (c) 2014 Alex Izvorski <aizvorski@gmail.com>

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""
import os
# time:2023.10.14 9:12
# author: Yuanping
# 注意修改main中的图片格式，如jpg还是png等
# 评估很多帧图片psnr并计算平局值的函数，只需改变gen_path和real_path就行，或者将相应jpg格式图片放入相应文件夹即可
# 最好都是相同命名，且gen的数量得小于等于real
# 会自动将评估得到的值存入csv，包括每一帧的和平均值
# 不需要计算某个值可以在video_quality函数中把计算某个值的过程注释掉
from glob import glob

import lpips
from skimage.metrics import structural_similarity as ssim
import cv2
import numpy
import math

import numpy as np
import pandas as pd


def psnr(img1, img2):
    mse = numpy.mean((img1 - img2) ** 2)
    if mse == 0:
        print("error:mse=0")
        exit(1)
        # return 0
    PIXEL_MAX = 255.0
    return 20 * math.log10(PIXEL_MAX / math.sqrt(mse))


def cal_psnr(gen_path, real_path):
    psnr_list = []
    for i in range(len(gen_path)):
        gen_pic = cv2.imread(gen_path[i])
        real_pic = cv2.imread(real_path[i])
        print("frame picture:")
        psnr_value = psnr(gen_pic, real_pic)
        print(i, " frame psnr: ", psnr_value)
        psnr_list.append([i, psnr_value])
    psnr_csv = pd.DataFrame(psnr_list, columns=['frame', 'psnr_value'])
    # psnr_csv = psnr_csv._append({'frame': 'average', 'psnr_value': psnr_csv["psnr_value"].mean()}, ignore_index=True)
    # psnr_csv.to_csv('video_psnr.csv', index=False)
    return psnr_csv


def cal_ssim(gen_path, real_path, quality_df):
    ssim_list = []
    for i in range(len(gen_path)):
        gen_pic = cv2.imread(gen_path[i])
        real_pic = cv2.imread(real_path[i])
        ssim_value = ssim(gen_pic, real_pic, channel_axis=-1)
        print(i, " frame ssim: ", ssim_value)
        ssim_list.append(ssim_value)
    quality_df['ssim_value'] = ssim_list
    return quality_df


def cal_LPIPS(gen_path, real_path, quality_df):
    """
    参考代码：https://blog.csdn.net/weixin_43135178/article/details/127664187
    """
    ## Initializing the model
    loss_fn = lpips.LPIPS(net='alex', version=0.1)  # pip install lpips
    lpips_list = []
    for i in range(len(gen_path)):
        try:
            # Load images
            img0 = lpips.im2tensor(lpips.load_image(gen_path[i]))
            img1 = lpips.im2tensor(lpips.load_image(real_path[i]))
            # Compute distance,之所以会有后面detach和numpy是因为要把梯度去掉，这样才能加到dataframe里面
            # 之所以还有个mean函数，是因为不加mean的话得到的是[[[value]]]，加了mean可以直接得到值
            current_lpips_distance = loss_fn.forward(img0, img1).detach().numpy().mean()
            print(i, " frame lpips: ", current_lpips_distance)
            lpips_list.append(current_lpips_distance)
        except Exception as e:
            print(e)
    quality_df['lpips_value'] = lpips_list
    return quality_df


def video_quality(gen_path, real_path):
    quality_df = cal_psnr(gen_path, real_path)
    quality_df = cal_ssim(gen_path, real_path, quality_df)
    quality_df = cal_LPIPS(gen_path, real_path, quality_df)
    quality_df = quality_df._append({'frame': 'average',
                                     'ssim_value': quality_df["ssim_value"].mean(),
                                    'psnr_value': quality_df["psnr_value"].mean(),
                                     'lpips_value': quality_df["lpips_value"].mean()},
                                    ignore_index=True)
    quality_df.to_csv('video_quality.csv', index=False)


if __name__ == '__main__':
    gen_path = glob("video_gen/*.jpg")
    real_path = glob("video_real/*.jpg")
    video_quality(gen_path, real_path)