import os
import stat
import subprocess
from pathlib import Path
from typing import Optional, Union
import requests
from lfc_toolkit.src.data_handlers.formatters import \
get_formatted_filename_for_lf
from lfc_toolkit.src.data_handlers.lightfield import RAWLightFieldData
from .quality_tool_wrapper import QualityToolWrapper
VMAF_METRICS_COMMAND_TEMPLATE = """
{vmaf_exec}
-r {original_yuv}
-d {decoded_yuv}
-w {view_width}
-h {view_height}
-b {bit_depth}
-p 444
-o {output_report}
--feature psnr
--feature float_ssim
--json
"""
[docs]
class VMAFMetrics(QualityToolWrapper):
"""VMAF quality metrics wrapper."""
[docs]
def __init__(
self,
reports_path: str = "./reports",
vmaf_executable_path: str = "./vmaf",
version: str = "2.2.1",
force_download: bool = False,
) -> None:
"""Initializes VMAFMetrics, downloading the VMAF executable if needed.
:param reports_path: Path to store metric reports, defaults to "./reports"
:type reports_path: str
:param vmaf_executable_path: Path for the VMAF executable, defaults to "./vmaf"
:type vmaf_executable_path: str
:param version: VMAF version to use, defaults to "2.2.1"
:type version: str
:param force_download: Whether to re-download even if executable exists, defaults to False
:type force_download: bool
:return: None
:rtype: None
"""
QualityToolWrapper.__init__(self, reports_path=reports_path, required_format="yuv")
print(vmaf_executable_path)
self.__vmaf_executable_path = vmaf_executable_path
self.__vmaf_executable = f"{self.__vmaf_executable_path}/vmaf"
self.__vmaf_version = version
if force_download and os.path.exists(self.__vmaf_executable):
os.remove(self.__vmaf_executable)
if not os.path.isfile(self.__vmaf_executable):
self.__download_vmaf()
def __download_vmaf(self) -> None:
"""Downloads the VMAF executable from the official GitHub releases.
:return: None
:rtype: None
"""
vmaf_url = f"https://github.com/Netflix/vmaf/releases/download/v{self.__vmaf_version}/vmaf"
os.makedirs(self.__vmaf_executable_path, exist_ok=True)
r = requests.get(vmaf_url, stream=True) # , verify=False)
if r.status_code == 200:
with open(self.__vmaf_executable, "wb") as f:
for chunk in r:
f.write(chunk)
os.chmod(self.__vmaf_executable, stat.S_IEXEC)
def _call_metrics(
self,
original_lightfield: RAWLightFieldData,
decoded_lightfield: RAWLightFieldData,
) -> Optional[Union[str, Path]]:
"""Runs VMAF to compute metrics between original and decoded lightfields.
:param original_lightfield: Original (reference) lightfield
:type original_lightfield: RAWLightFieldData
:param decoded_lightfield: Decoded lightfield to evaluate
:type decoded_lightfield: RAWLightFieldData
:return: Path to the output report JSON, or None on failure
:rtype: Optional[Union[str, Path]]
"""
output_vmaf_report = get_formatted_filename_for_lf(
self.reports_path,
decoded_lightfield,
bpp=decoded_lightfield.bpp_for_naming,
file_extension="json",
extra="vmaf",
)
command = VMAF_METRICS_COMMAND_TEMPLATE.format(
vmaf_exec=self.__vmaf_executable,
view_width=decoded_lightfield.view_width,
view_height=decoded_lightfield.view_height,
bit_depth=10,
decoded_yuv=decoded_lightfield.raw_path,
original_yuv=original_lightfield.raw_path,
output_report=output_vmaf_report,
).split()
res = subprocess.run(command, capture_output=True)
if res.returncode != 0:
print("Got error running vmaf tool")
print(res.stderr)
return None
# raise Exception(res.stderr)
return output_vmaf_report