VMAF PyTorch

Kirill Aistov, Maxim Koroteev

This is an (unofficial) PyTorch implementation of VMAF (Video Multi-Method Assessment Fusion) video quality metric based on the official C implementation.

Installation

From pip

pip install vmaf-torch

From source

git clone https://github.com/alvitrioliks/VMAF-torch.git
cd vmaf-torch
pip install -e .

Usage

Initialize

from vmaf_torch import VMAF
vmaf = VMAF()                            # init vmaf class
vmaf = vmaf.to("cuda")                   # move internal tensors to GPU
vmaf = vmaf.compile()                    # optionally use .compile() for up to 30% inference speed up, requires torch>=2.0, first pass might take ~30 seconds

Load yuv files

from vmaf_torch.utils import yuv_to_tensor  

# load reference video Y channel as [num_frames,1,h,w] tensor in [0,255] range                  
ref = yuv_to_tensor('/<your path>/BigBuckBunny_25fps.yuv', 1920, 1080, num_frames=50, channel='y')  
ref = ref.to(device)

# load distorted video Y channel as [num_frames,1,h,w] tensor in [0,255] range
dist = yuv_to_tensor('/<your path>/BigBuckBunny_20_288_375.yuv', 1920, 1080, num_frames=50, channel='y')  
dist = dist.to(device)

Compute VMAF score

vmaf(ref, dist)       # order of arguments matters

Compute elementary features

vmaf.compute_adm_features(ref, dist)
vmaf.compute_vif_features(ref, dist)
vmaf.compute_motion(ref)

Difference with the official implementation

We tried to reproduce all details of the original implementation, however some difference in outputs is possible due to the fact that the original uses quantized integer values while this implementation uses floats.

Comparing with the official Netflix implementation on Netflix Public Dataset that includes 79 video streams (using first 50 frames from each video), so the plot below contains 79 points, each obtained by averaging over 50 numbers:

Comparison
Difference

Absolute differences' statistics are as follows:

min: 8.7e-07
max: 0.0272
mean: 0.0063
std: 0.0079

To reproduce these numbers see vmaf_compare_implementations.ipynb

Notes

To ensure that gradients are computed correctly we perform gradient checking, see vmaf_gradient_checking.ipynb. For details refer to our arxiv paper.
Note that with clip_score=True (which is the default behaviour of the official VMAF implementation) it is possible to get gradients equal to 0 in some cases since clipping operation has derivative equal to 0 on parts of its domain. We recommend setting clip_score=False for tasks requiring gradient computation.
This implementation also includes VMAF NEG version of the algorithm, use NEG=True to enable it. We notice that with NEG=True it is also possible to get gradients equal to 0 in some cases due to the nature of this algorithm. We plan to investigate this in more detail in the future.

Citation

If you use this implementation in your project, please, use the following reference:

@misc{vmaf-torch,
      title={VMAF Re-implementation on PyTorch: Some Experimental Results}, 
      author={Kirill Aistov and Maxim Koroteev},
      year={2023},
      eprint={2310.15578},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}

Dependencies of vmaf-torch (4).
Dependency	Optional
numpy	false
pandas	false
torch	false
yuvio	false

Package profile

vmaf-torch

Releases

Release table

PyPI Downloads

Dependencies

Details

Project URLs

Dependencies

Package description