使用 VMAF 测量图片质量
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
This comment has been minimized.
文件结构:
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