使用 VMAF 测量图片质量

vmaf_for_image.py

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'        # 输出结果文件名称
)

文件结构：

│  libvmaf-0.dll
│  vmafossexec.exe
│  vmaf_for_image.py
│
└─model
  ├─  vmaf_v0.6.1.pkl
  └─  vmaf_v0.6.1.pkl.model

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