rust-av/av-metrics

Implement VMAF support

shssoichiro opened this issue · 9 comments

The codebase for VMAF is quite large, but it would be nice to have a pure Rust implementation. This is by no means high priority, just if someone (probably me) gets bored.

Ahh yes!

I’m working on an image optimization service for the e-commerce world that uses VMAF internally.

I’ve been using it VIA the C API for some time, and while I love VMAF, I still haven’t yet figured out what temp_data does (it’s not used in the FFmpeg implementation):

int compute_vmaf(double* vmaf_score, char* fmt, int width, int height, 
int (*read_frame)(float *ref_data, float *main_data, float *temp_data, int stride, void *user_data), 
void *user_data, char *model_path, char *log_path, char *log_fmt, int disable_clip, 
int disable_avx, int enable_transform, int phone_model, int do_psnr, int do_ssim, 
int do_ms_ssim, char *pool_method, int thread, int subsample, int enable_conf_interval);

The callback argument is weird because I can’t fill the buffers directly with a single frame and then return the value 2 (what seems to signal to VMAF that it’s done).

Anyway what I wanted to mention. I still haven’t figured out why and perhaps it’s me, but for some reason I can’t run multiple instances of VMAF in a single process in parallel on different media. It’s significantly more likely to segfault. If it’s a bug with VMAF it’s also not deterministic for some reason; sometimes it runs just fine…

If the issue I’ve encountered is a hard limitation of VMAF. Perhaps a rust implementation could improve upon the original in this regard.

Also for anyone interested, I just publish a VMAF wrapper to crates here vmaf-sys.

It’s pretty direct. But it does the grunt work of downloading, building and statically linking libvmaf internally (so you don’t have to have it pre-installed on your system, like when building FFmpeg).

Also the default and 4K model is baked into the binary.

Thanks @colbyn! I think I will go ahead and use vmaf-sys to get basic support in this crate, then Rewrite it In Rust at a later time.

@shssoichiro Okay sounds good. Let me know if you have any problems.

Update: finally got the dumb thing to build with docs.rs.

I'm having some difficulties with the wrapper you created, and wondering if there's a good way to resolve this via Rust FFI. It looks like the run function takes a model_path which needs to point to a path representing one of the model files, but that's not possible the way the wrapper is currently set up--pulling the models in with include_str.

@shssoichiro Oh I was just working on this (and fixed a secondary issue described below).

Initially in my head I was thinking the scope of vmaf-sys would be direct bindings to libvmaf and the included models. But I think I may as well move handling of such to the 'sys' crate since it's directly related to the compute_vmaf function. Anyway here's the docs for it. Since compute_vmaf requires a file path, it just creates a temporary directory and automatically removes such when the program exits (as explained in the docs).

I never realized this until just now. The compute_vmaf function has an argument called model_path that also at runtime checks for a second file with the same name, but with a .model extension (I never realized this since before I just included the whole models dir from the VMAF git repo).

Looks like the docs hasn't built yet. But for the 0.0.8 version you can see a helper function here that should give you a valid file path reference that you can give to VMAF that also includes the secondary (*.pkl.model) file.

use std::path::PathBuf;
use std::ffi::CStr;
use std::ffi::CString;
use std::os::raw::c_char;
use std::os::raw::c_int;
use libc::{size_t, c_float};

// ...

// SETTINGS
let mut vmaf_score = 0.0;
let model_path = vmaf_sys::extras::get_def_model_path()
    .to_str()
    .expect("PathBuf to str failed")
    .to_owned();
let model_path = CString::new(model_path).expect("CString::new failed");
let mut fmt = CString::new(String::from("yuv420p")).expect("CString::new failed");
let width = source1.width;
let height = source1.height;
let log_path: *mut c_char = std::ptr::null_mut();
let log_fmt: *mut c_char = std::ptr::null_mut();
let disable_clip = 0;
let disable_avx = 0;
let enable_transform = 0;
let phone_model = 0;
let do_psnr = 0;
let do_ssim = 0;
let do_ms_ssim = 0;
let pool_method: *mut c_char = std::ptr::null_mut();
let n_thread = 1;
let n_subsample = 1;
let enable_conf_interval = 0;

// GO!
let compute_vmaf_res = vmaf_sys::compute_vmaf(
    &mut vmaf_score,
    fmt.as_ptr() as *mut c_char,
    width as c_int,
    height as c_int,
    Some(read_frame),
    vmaf_ctx as *mut libc::c_void,
    model_path.as_ptr() as *mut c_char,
    log_path,
    log_fmt,
    disable_clip,
    disable_avx,
    enable_transform,
    phone_model,
    do_psnr,
    do_ssim,
    do_ms_ssim,
    pool_method,
    n_thread,
    n_subsample,
    enable_conf_interval
);

@shssoichiro Did that solve your problem?

Update (if it interests anyone): the results from the "default" model (well the default in FFmpeg last time I checked - vmaf_v0.6.1.pkl) is weird. Given two identical images the 4K model (vmaf_4k_v0.6.1.pkl) usually returns 100, but vmaf_v0.6.1.pkl usually returns something like 97... IDK maybe it's nothing and I'm just being dumb.

Update: Oh yep. It should be fixed now in 0.0.9 (I just tested it).

Yes, thanks! Now I just have to figure out how this read_frame thing works... 🙂

Update:

Just FYI for anyone using vmaf-sys, I've recently open sourced a project called Imager (imager-io/imager) that uses it internally.

See here for an example of calling vmaf_sys::compute_vmaf (may be more useful than my old snippets). An example of the read_frame callback is here.

Also older versions (pre 0.0.10) won't build (details here).