blob: 23d247175a70ca76855af8ea09fd5dac9ea71828 [file] [log] [blame]
#!/usr/bin/env python
## Copyright (c) 2019, Alliance for Open Media. All rights reserved
##
## This source code is subject to the terms of the BSD 2 Clause License and
## the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
## was not distributed with this source code in the LICENSE file, you can
## obtain it at www.aomedia.org/license/software. If the Alliance for Open
## Media Patent License 1.0 was not distributed with this source code in the
## PATENTS file, you can obtain it at www.aomedia.org/license/patent.
##
__author__ = "maggie.sun@intel.com, ryan.lei@intel.com"
import os
import re
import logging
from Config import BinPath, LoggerName, VMAF
from Utils import GetShortContentName, ExecuteCmd
subloggername = "CalcQtyMetrics_VMAFTool"
loggername = LoggerName + '.' + '%s' % subloggername
logger = logging.getLogger(loggername)
Model_Pkg_File = os.path.join(BinPath, 'vmaf_v0.6.1.pkl')
VMAFMetricsFullList = ['VMAF_Y','VMAF_Y-NEG','PSNR_Y','PSNR_U','PSNR_V','SSIM_Y(dB)',
'MS-SSIM_Y(dB)','PSNR-HVS','CIEDE2000','APSNR_Y',
'APSNR_U','APSNR_V']
def ParseVMAFLogFile(vmaf_log):
floats = len(VMAFMetricsFullList) * [0.0]
flog = open(vmaf_log, 'r')
for line in flog:
m = re.search(r"\"vmaf\".*\s+mean=\"(\d+\.?\d*)\"\s+",line)
if m:
floats[0] = m.group(1)
m = re.search(r"\"vmaf_neg\".*\s+mean=\"(\d+\.?\d*)\"\s+", line)
if m:
floats[1] = m.group(1)
m = re.search(r"\"psnr_y\".*\s+mean=\"(\d+\.?\d*)\"\s+", line)
if m:
floats[2] = m.group(1)
m = re.search(r"\"psnr_cb\".*\s+mean=\"(\d+\.?\d*)\"\s+", line)
if m:
floats[3] = m.group(1)
m = re.search(r"\"psnr_cr\".*\s+mean=\"(\d+\.?\d*)\"\s+", line)
if m:
floats[4] = m.group(1)
m = re.search(r"\"float_ssim\".*\s+mean=\"(\d+\.?\d*)\"\s+", line)
if m:
floats[5] = m.group(1)
m = re.search(r"\"float_ms_ssim\".*\s+mean=\"(\d+\.?\d*)\"\s+", line)
if m:
floats[6] = m.group(1)
m = re.search(r"\"psnr_hvs\".*\s+mean=\"(\d+\.?\d*)\"\s+", line)
if m:
floats[7] = m.group(1)
m = re.search(r"\"ciede2000\".*\s+mean=\"(\d+\.?\d*)\"\s+", line)
if m:
floats[8] = m.group(1)
#<aggregate_metrics apsnr_y="46.817276" apsnr_cb="49.092538" apsnr_cr="50.014785" />
m = re.search(r"aggregate_metrics\s+apsnr_y=\"(\d+\.?\d*)\"\s+apsnr_cb=\"(\d+\.?\d*)\"\s+apsnr_cr=\"(\d+\.?\d*)\"", line)
if m:
floats[9] = m.group(1)
floats[10] = m.group(2)
floats[11] = m.group(3)
flog.close()
floats = [float(i) for i in floats]
print_str = "VMAF quality metrics: "
for metrics, idx in zip(VMAFMetricsFullList, range(len(VMAFMetricsFullList))):
print_str += "%s = %2.5f, " % (metrics, floats[idx])
logger.info(print_str)
return floats[0:len(VMAFMetricsFullList)]
def GetVMAFLogFile(recfile, path):
filename = GetShortContentName(recfile, False) + '_vmaf.log'
file = os.path.join(path, filename)
return file
################################################################################
##################### Exposed Functions ########################################
def VMAF_CalQualityMetrics(origfile, recfile, fmt, num, w, h, bit_depth,
logfilePath, LogCmdOnly=False):
vmaf_log = GetVMAFLogFile(recfile, logfilePath)
args = " -r %s -d %s --aom_ctc v1.0 -q -o %s" \
% (origfile, recfile, vmaf_log)
cmd = VMAF + args
ExecuteCmd(cmd, LogCmdOnly)
def VMAF_GatherQualityMetrics(recfile, logfilePath):
vmaf_log = GetVMAFLogFile(recfile, logfilePath)
results = ParseVMAFLogFile(vmaf_log)
return results