使用 VMAF 测量图片质量
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import os | |
import shutil | |
from PIL import Image | |
def check_size(image_a, image_b): | |
global i_size | |
w_a, h_a = Image.open(image_a).size | |
w_b, h_b = Image.open(image_b).size | |
if w_a == w_b and h_a == h_b: | |
i_size = [w_a, h_a] | |
return True, i_size | |
return False, [] | |
def convert_yuv(image_a, image_b, i_fmt): | |
os.system(f'ffmpeg -y -i "{image_a}" -pix_fmt {i_fmt} -vsync 0 temp_a.yuv') | |
os.system(f'ffmpeg -y -i "{image_b}" -pix_fmt {i_fmt} -vsync 0 temp_b.yuv') | |
def get_vmaf(i_size, i_fmt, log_fmt, log_name): | |
cmd = f'vmafossexec.exe {i_fmt} {i_size[0]} {i_size[1]} temp_a.yuv temp_b.yuv ' | |
cmd += f'model/vmaf_v0.6.1.pkl --log-fmt {log_fmt} --log {log_name}.{log_fmt} ' | |
cmd += '--psnr --ssim --ms-ssim --thread 2 --subsample 1' | |
print(cmd) | |
os.system(cmd) | |
os.remove('temp_a.yuv') | |
os.remove('temp_b.yuv') | |
def go(image_a, image_b, i_fmt='yuv444p10le', log_fmt='json', log_name='vmaf_output'): | |
is_same, i_size = check_size(image_a, image_b) | |
if not is_same: | |
print('需要两张大小完全相同的图片') | |
return | |
convert_yuv(image_a, image_b, i_fmt) | |
get_vmaf(i_size, i_fmt, log_fmt, log_name) | |
go( | |
'groundtruth.png', # 参考无噪图片 | |
'gaussian_bm3d.png', # 降噪结果 | |
log_name='gaussian_bm3d' # 输出结果文件名称 | |
) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
文件结构:
VMAF Windows Binary: https://ci.appveyor.com/project/EwoutH/vmaf/history
包含: vmafossexec.exe、libvmaf-0.dll
模型文件: https://github.com/Netflix/vmaf/tree/master/model
包含: vmaf_v0.6.1.pkl、vmaf_v0.6.1.pkl.model