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).