[FFmpeg-devel] [PATCH] VMAF is now propagated to the AVFrame coming from the graph

[ichlubna at fit.vutbr.cz]

From: ichlubna <ichlubna at fit.vutbr.cz>

Related to my ticket here: https://trac.ffmpeg.org/ticket/10586
VMAF score was not propagated to AVFormat like SSIM or PSNR in the result of the filter graph. I have fixed this to make the usage consistent and possible to get VMAF score per-frame in libavfilter.
The only dirty thing here is the added for loop to compute the score twice. This is done to get the score in real time. Otherwise the score is delayed by one frame. Computing the score twice should not affect the final averaged result as each frame is added twice so the average does not change.

---
 libavfilter/vf_libvmaf.c | 76 +++++++++++++++++++++++++---------------
 1 file changed, 47 insertions(+), 29 deletions(-)

diff --git a/libavfilter/vf_libvmaf.c b/libavfilter/vf_libvmaf.c
index 2726b061ac..4c84877812 100644
--- a/libavfilter/vf_libvmaf.c
+++ b/libavfilter/vf_libvmaf.c
@@ -139,12 +139,40 @@ static int copy_picture_data(AVFrame *src, VmafPicture *dst, unsigned bpc)
     return 0;
 }
 
+static enum VmafPoolingMethod pool_method_map(const char *pool_method)
+{
+    if (pool_method) {
+        if (av_stristr(pool_method, "min"))
+            return VMAF_POOL_METHOD_MIN;
+        if (av_stristr(pool_method, "mean"))
+            return VMAF_POOL_METHOD_MEAN;
+        if (av_stristr(pool_method, "harmonic_mean"))
+            return VMAF_POOL_METHOD_HARMONIC_MEAN;
+    }
+
+    return VMAF_POOL_METHOD_MEAN;
+}
+
+static void set_meta(AVDictionary **metadata, const char *key, char comp, float d)
+{
+    char value[128];
+    snprintf(value, sizeof(value), "%f", d);
+    if (comp) {
+        char key2[128];
+        snprintf(key2, sizeof(key2), "%s%c", key, comp);
+        av_dict_set(metadata, key2, value, 0);
+    } else {
+        av_dict_set(metadata, key, value, 0);
+    }
+}
+
 static int do_vmaf(FFFrameSync *fs)
 {
     AVFilterContext *ctx = fs->parent;
     LIBVMAFContext *s = ctx->priv;
     VmafPicture pic_ref, pic_dist;
     AVFrame *ref, *dist;
+    double vmaf_score;
     int err = 0;
 
     int ret = ff_framesync_dualinput_get(fs, &dist, &ref);
@@ -160,25 +188,29 @@ static int do_vmaf(FFFrameSync *fs)
                av_color_range_name(ref->color_range));
     }
 
-    err = copy_picture_data(ref, &pic_ref, s->bpc);
-    if (err) {
-        av_log(s, AV_LOG_ERROR, "problem during vmaf_picture_alloc.\n");
-        return AVERROR(ENOMEM);
-    }
+    for(int i=0; i<2; i++){
+        err = copy_picture_data(ref, &pic_ref, s->bpc);
+        if (err) {
+            av_log(s, AV_LOG_ERROR, "problem during vmaf_picture_alloc.\n");
+            return AVERROR(ENOMEM);
+        }
 
-    err = copy_picture_data(dist, &pic_dist, s->bpc);
-    if (err) {
-        av_log(s, AV_LOG_ERROR, "problem during vmaf_picture_alloc.\n");
-        vmaf_picture_unref(&pic_ref);
-        return AVERROR(ENOMEM);
-    }
+        err = copy_picture_data(dist, &pic_dist, s->bpc);
+        if (err) {
+            av_log(s, AV_LOG_ERROR, "problem during vmaf_picture_alloc.\n");
+            vmaf_picture_unref(&pic_ref);
+            return AVERROR(ENOMEM);
+        }
 
-    err = vmaf_read_pictures(s->vmaf, &pic_ref, &pic_dist, s->frame_cnt++);
-    if (err) {
-        av_log(s, AV_LOG_ERROR, "problem during vmaf_read_pictures.\n");
-        return AVERROR(EINVAL);
+        err = vmaf_read_pictures(s->vmaf, &pic_ref, &pic_dist, s->frame_cnt++);
+        if (err) {
+            av_log(s, AV_LOG_ERROR, "problem during vmaf_read_pictures.\n");
+            return AVERROR(EINVAL);
+        }
     }
 
+    vmaf_score_at_index(s->vmaf, s->model[0], &vmaf_score, s->frame_cnt - 2); 
+    set_meta(&dist->metadata, "lavfi.vmaf", 0, vmaf_score);
     return ff_filter_frame(ctx->outputs[0], dist);
 }
 
@@ -637,20 +669,6 @@ static enum VmafOutputFormat log_fmt_map(const char *log_fmt)
     return VMAF_OUTPUT_FORMAT_XML;
 }
 
-static enum VmafPoolingMethod pool_method_map(const char *pool_method)
-{
-    if (pool_method) {
-        if (av_stristr(pool_method, "min"))
-            return VMAF_POOL_METHOD_MIN;
-        if (av_stristr(pool_method, "mean"))
-            return VMAF_POOL_METHOD_MEAN;
-        if (av_stristr(pool_method, "harmonic_mean"))
-            return VMAF_POOL_METHOD_HARMONIC_MEAN;
-    }
-
-    return VMAF_POOL_METHOD_MEAN;
-}
-
 static av_cold void uninit(AVFilterContext *ctx)
 {
     LIBVMAFContext *s = ctx->priv;
-- 
2.42.0