VMAF Output Python Parser and Class Wrapper

VMAFOutput.py

import os
import json
from typing import List


class AgregateMetrics:
    """
    Stores the attributes of the agregate metrics.
    """
    integer_adm2: float = None
    integer_adm_scale0: float = None
    integer_adm_scale1: float = None
    integer_adm_scale2: float = None
    integer_adm_scale3: float = None
    integer_motion2: float = None
    integer_motion: float = None
    integer_vif_scale0: float = None
    integer_vif_scale1: float = None
    integer_vif_scale2: float = None
    integer_vif_scale3: float = None
    psnr_y: float = None
    psnr_cb: float = None
    psnr_cr: float = None
    psnr_hvs_y: float = None
    psnr_hvs_cb: float = None
    psnr_hvs_cr: float = None
    psnr_hvs: float = None
    float_ssim: float = None
    float_ms_ssim: float = None
    ciede2000: float = None
    cambi: float = None
    vmaf: float = None

    def __init__(self, **kwargs):
        # Updated the dict overwriting its values or setting new attributes.
        self.__dict__.update(kwargs)

    def __str__(self):
        return str(self.__dict__)


class FrameMetrics(AgregateMetrics):
    """
    Stores the attributes of a single frame.
    """
    frameNum: int = None

    def __init__(self, **kwargs):
        super().__init__(**kwargs["metrics"])
        self.frameNum = kwargs["frameNum"]

    def __str__(self):
        return str(self.__dict__)


class PooledMetric:
    """
    Stores the atributes of a single pooled metric.
    """
    name: str = None
    min: float = None
    max: float = None
    mean: float = None
    harmonic_mean: float = None

    def __init__(self, **kwargs):
        # Updated the dict overwriting its values or setting new attributes.
        self.__dict__.update(kwargs)

    def __str__(self):
        return str(self.__dict__)


class PooledMetrics:
    """
    Wraps multiple PooledMetric together in a sort of list. This class allows to access the metrics by dot notation
    (e.g. `pooled_metrics.psnr_y`) or by dict notation (e.g. `pooled_metrics["psnr_y"]`)
    """
    def add(self, pooled_metric: PooledMetric):
        self.__setattr__(pooled_metric.name, pooled_metric)

    def remove(self, pooled_metric: PooledMetric):
        self.__delattr__(pooled_metric.name)

    def __getitem__(self, item):
        return self.__getattribute__(item)

    def __setitem__(self, key, value):
        self.__setattr__(key, value)

    def __str__(self):
        return str(self.__dict__)


class VMAFOutput:
    """
    Wraps the result of a vmaf output file.
    """
    version: str
    sequence_width: int
    sequence_height: int
    sequence_fps: float

    frames: List[FrameMetrics]
    pooled_metrics: PooledMetrics
    aggregate_metrics: AgregateMetrics

    def __init__(self, frames: List[FrameMetrics], pooled_metrics: PooledMetrics, aggregate_metrics: AgregateMetrics,
                 version: str = None, sequence_width: int = None, sequence_height: int = None,
                 sequence_fps: float = None):
        self.version = version
        self.sequence_width = sequence_width
        self.sequence_height = sequence_height
        self.sequence_fps = sequence_fps
        self.frames = frames
        self.pooled_metrics = pooled_metrics
        self.aggregate_metrics = aggregate_metrics

    def __str__(self):
        return str(self.__dict__)

    @staticmethod
    def parse_vmaf_json(vmaf_output: str):
        if not os.path.splitext(vmaf_output)[-1] == ".json":
            raise ValueError(f"Expecting .json file, got {os.path.splitext(vmaf_output)[-1]}")

        with open(vmaf_output) as freader:
            json_dict = json.load(fp=freader)  # noqa

        frames = []
        for frame in json_dict["frames"]:
            frames.append(FrameMetrics(**frame))

        pooled_metrics = PooledMetrics()
        for pooled_metric_name, pooled_metric_metrics in json_dict["pooled_metrics"].items():
            pooled_metrics.add(PooledMetric(name=pooled_metric_name, **pooled_metric_metrics))

        aggregate_metrics = AgregateMetrics(**json_dict["aggregate_metrics"])

        json_dict["version"] = json_dict["version"] if "version" in json_dict.keys() else None
        json_dict["fps"] = json_dict["fps"] if "fps" in json_dict.keys() else None

        return VMAFOutput(version=json_dict["version"], sequence_fps=json_dict["fps"], frames=frames,
                          pooled_metrics=pooled_metrics, aggregate_metrics=aggregate_metrics)