文件准备:
创建两个文件夹,分别叫做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)