Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
使用 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' # 输出结果文件名称
)
@Sg4Dylan

This comment has been minimized.

Copy link
Owner Author

@Sg4Dylan Sg4Dylan commented Jun 7, 2020

文件结构:

│  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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.