diff mbox series

[FFmpeg-devel,v1] avfilter/vf_libvmaf: Add metadata propagation support

Message ID 20240824012544.43085-2-yigithanyigitdevel@gmail.com
State New
Headers show
Series [FFmpeg-devel,v1] avfilter/vf_libvmaf: Add metadata propagation support | expand

Checks

Context Check Description
yinshiyou/make_loongarch64 success Make finished
yinshiyou/make_fate_loongarch64 success Make fate finished
andriy/make_x86 success Make finished
andriy/make_fate_x86 success Make fate finished

Commit Message

Yigithan Yigit Aug. 24, 2024, 1:25 a.m. UTC
---
 libavfilter/vf_libvmaf.c | 271 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 271 insertions(+)
diff mbox series

Patch

diff --git a/libavfilter/vf_libvmaf.c b/libavfilter/vf_libvmaf.c
index f655092b20..0b4a63cf92 100644
--- a/libavfilter/vf_libvmaf.c
+++ b/libavfilter/vf_libvmaf.c
@@ -27,8 +27,11 @@ 
 #include "config_components.h"
 
 #include <libvmaf.h>
+#include <libvmaf/version.h>
 
 #include "libavutil/avstring.h"
+#include "libavutil/dict.h"
+#include "libavutil/frame.h"
 #include "libavutil/mem.h"
 #include "libavutil/opt.h"
 #include "libavutil/pixdesc.h"
@@ -46,6 +49,29 @@ 
 #include "libavutil/hwcontext_cuda_internal.h"
 #endif
 
+#define VMAF_VERSION_INT_VER(major, minor, patch) \
+    ((major) * 10000 + (minor) * 100 + (patch))
+
+#if VMAF_VERSION_INT_VER(VMAF_API_VERSION_MAJOR, VMAF_API_VERSION_MINOR, VMAF_API_VERSION_PATCH) > VMAF_VERSION_INT_VER(3, 0, 0)
+#define CONFIG_LIBVMAF_METADATA_FILTER 1
+#else 
+#define CONFIG_LIBVMAF_METADATA_FILTER 0
+#endif
+
+#if CONFIG_LIBVMAF_METADATA_FILTER
+typedef struct FrameList {
+    AVFrame *frame;
+    unsigned frame_number;
+    unsigned propagated_handlers_cnt;
+    struct FrameList *next;
+} FrameList;
+
+typedef struct CallbackStruct {
+    struct LIBVMAFContext *s;
+    FrameList *frame_list;
+} CallbackStruct;
+#endif
+
 typedef struct LIBVMAFContext {
     const AVClass *class;
     FFFrameSync fs;
@@ -56,6 +82,14 @@  typedef struct LIBVMAFContext {
     int n_subsample;
     char *model_cfg;
     char *feature_cfg;
+#if CONFIG_LIBVMAF_METADATA_FILTER
+    char *metadata_feature_cfg;
+    struct {
+        VmafMetadataConfiguration *metadata_cfgs;
+        unsigned metadata_cfg_cnt;
+    } metadata_cfg_list;
+    CallbackStruct *cb;
+#endif
     VmafContext *vmaf;
     VmafModel **model;
     unsigned model_cnt;
@@ -77,6 +111,9 @@  static const AVOption libvmaf_options[] = {
     {"n_subsample", "Set interval for frame subsampling used when computing vmaf.",     OFFSET(n_subsample), AV_OPT_TYPE_INT, {.i64=1}, 1, UINT_MAX, FLAGS},
     {"model",  "Set the model to be used for computing vmaf.",                          OFFSET(model_cfg), AV_OPT_TYPE_STRING, {.str="version=vmaf_v0.6.1"}, 0, 1, FLAGS},
     {"feature",  "Set the feature to be used for computing vmaf.",                      OFFSET(feature_cfg), AV_OPT_TYPE_STRING, {.str=NULL}, 0, 1, FLAGS},
+#if CONFIG_LIBVMAF_METADATA_FILTER
+    {"metadata_handler",  "Set the feature to be propagated as metadata.",              OFFSET(metadata_feature_cfg), AV_OPT_TYPE_STRING, {.str="name=vmaf"}, 0, 1, FLAGS},
+#endif
     { NULL }
 };
 
@@ -105,6 +142,123 @@  static enum VmafPixelFormat pix_fmt_map(enum AVPixelFormat av_pix_fmt)
     }
 }
 
+#if CONFIG_LIBVMAF_METADATA_FILTER
+static int add_to_frame_list(FrameList **head, AVFrame *frame, unsigned frame_number)
+{
+    FrameList *new_frame = av_malloc(sizeof(FrameList));
+    if (!new_frame)
+        return AVERROR(ENOMEM);
+
+    new_frame->frame = frame;
+    new_frame->frame_number = frame_number;
+    new_frame->propagated_handlers_cnt = 0;
+    new_frame->next = NULL;
+
+    if (*head == NULL) {
+        *head = new_frame;
+    } else {
+        FrameList *current = *head;
+        while (current->next != NULL) {
+            current = current->next;
+        }
+        current->next = new_frame;
+    }
+
+    return 0;
+}
+
+static int remove_from_frame_list(FrameList **frame_list, unsigned frame_number)
+{
+    FrameList *cur = *frame_list;
+    FrameList *prev = NULL;
+
+    while (cur) {
+        if (cur->frame_number == frame_number) {
+            if (prev)
+                prev->next = cur->next;
+            else
+                *frame_list = cur->next;
+            av_free(cur);
+            return 0;
+        }
+        prev = cur;
+        cur = cur->next;
+    }
+
+    return AVERROR(EINVAL);
+}
+
+static int free_frame_list(FrameList **frame_list)
+{
+    FrameList *cur = *frame_list;
+    while (cur) {
+        FrameList *next = cur->next;
+        av_frame_free(&cur->frame);
+        av_free(cur);
+        cur = next;
+    }
+    *frame_list = NULL;
+    return 0;
+}
+
+static FrameList* get_frame_from_frame_list(FrameList *frame_list,
+                                          unsigned frame_number)
+{
+    FrameList *cur = frame_list;
+    while (cur) {
+        if (cur->frame_number == frame_number)
+            return cur;
+        cur = cur->next;
+    }
+    return NULL;
+}
+
+static void set_meta(void *data, VmafMetadata *metadata)
+{
+    int err = 0;
+    FrameList *current_frame = NULL;
+    CallbackStruct *cb = data;
+    char value[128], key[128];
+    snprintf(value, sizeof(value), "%0.2f", metadata->score);
+    snprintf(key, sizeof(key), "%s.%d", metadata->feature_name, metadata->picture_index);
+
+    current_frame = get_frame_from_frame_list(cb->frame_list, metadata->picture_index);
+    if (!current_frame) {
+        av_log(NULL, AV_LOG_ERROR, "could not find frame with index: %d\n",
+               metadata->picture_index);
+        return;
+    }
+
+    err = av_dict_set(&current_frame->frame->metadata, key, value, 0);
+    if (err < 0)
+        av_log(NULL, AV_LOG_ERROR, "could not set metadata: %s\n", key);
+
+    current_frame->propagated_handlers_cnt++;
+
+    if (current_frame->propagated_handlers_cnt == cb->s->metadata_cfg_list.metadata_cfg_cnt) {
+        FrameList *cur = cb->frame_list;
+        // This code block allows to send frames monotonically
+        while(cur && cur->frame_number <= metadata->picture_index) {
+            if (cur->propagated_handlers_cnt == cb->s->metadata_cfg_list.metadata_cfg_cnt) {
+                FrameList *next;
+                // This necessary to avoid segfaults in filtergraphs
+                if (cb->s->fs.parent->outputs[0]) {
+                    ff_filter_frame(cb->s->fs.parent->outputs[0], cur->frame);
+                }
+                else
+                    av_frame_free(&cur->frame);
+                next = cur->next;
+                remove_from_frame_list(&cb->frame_list, cur->frame_number);
+                cur = next;
+            }
+            else
+                break;
+        }
+    }
+    av_log(cb->s->fs.parent, AV_LOG_DEBUG, "VMAF feature: %s, score: %f\n", key, metadata->score);
+}
+#endif
+
 static int copy_picture_data(AVFrame *src, VmafPicture *dst, unsigned bpc)
 {
     const int bytes_per_value = bpc > 8 ? 2 : 1;
@@ -160,13 +314,28 @@  static int do_vmaf(FFFrameSync *fs)
         return AVERROR(ENOMEM);
     }
 
+#if CONFIG_LIBVMAF_METADATA_FILTER
+    err = add_to_frame_list(&s->cb->frame_list, dist, s->frame_cnt);
+    if (err) {
+        av_log(s, AV_LOG_ERROR, "problem during add_to_frame_list.\n");
+        return AVERROR(ENOMEM);
+    }
+#endif
+
     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);
     }
 
+#if CONFIG_LIBVMAF_METADATA_FILTER
+    if (!s->metadata_cfg_list.metadata_cfg_cnt)
+        return ff_filter_frame(ctx->outputs[0], dist);
+    else
+        return 0;
+#else
     return ff_filter_frame(ctx->outputs[0], dist);
+#endif
 }
 
 static AVDictionary **delimited_dict_parse(char *str, unsigned *cnt)
@@ -408,6 +577,83 @@  exit:
     return err;
 }
 
+#if CONFIG_LIBVMAF_METADATA_FILTER
+static int parse_metadata_handlers(AVFilterContext *ctx)
+{
+    LIBVMAFContext *s = ctx->priv;
+    AVDictionary **dict;
+    unsigned dict_cnt;
+    int err = 0;
+
+    if (!s->metadata_feature_cfg)
+        return 0;
+
+    dict_cnt = 0;
+    dict = delimited_dict_parse(s->metadata_feature_cfg, &dict_cnt);
+    if (!dict) {
+        av_log(ctx, AV_LOG_ERROR,
+               "could not parse metadata feature config: %s\n",
+               s->metadata_feature_cfg);
+        return AVERROR(EINVAL);
+    }
+
+    for (unsigned i = 0; i < dict_cnt; i++) {
+        VmafMetadataConfiguration *metadata_cfg = av_calloc(1, sizeof(*metadata_cfg));
+        const AVDictionaryEntry *e = NULL;
+        char *feature_name = NULL;
+
+        while (e = av_dict_iterate(dict[i], e)) {
+            if (!strcmp(e->key, "name")) {
+                metadata_cfg->feature_name = av_strdup(e->value);
+                continue;
+            }
+        }
+
+        metadata_cfg->data = s->cb;
+        metadata_cfg->callback = &set_meta;
+
+        err = vmaf_register_metadata_handler(s->vmaf, *metadata_cfg);
+        if (err) {
+            av_log(ctx, AV_LOG_ERROR,
+                   "problem during vmaf_register_metadata_handler: %s\n",
+                   feature_name);
+            goto exit;
+        }
+
+        s->metadata_cfg_list.metadata_cfgs = av_realloc(s->metadata_cfg_list.metadata_cfgs,
+                                             (s->metadata_cfg_list.metadata_cfg_cnt + 1) *
+                                             sizeof(*s->metadata_cfg_list.metadata_cfgs));
+        if (!s->metadata_cfg_list.metadata_cfgs) {
+            err = AVERROR(ENOMEM);
+            goto exit;
+        }
+
+        s->metadata_cfg_list.metadata_cfgs[s->metadata_cfg_list.metadata_cfg_cnt++] = *metadata_cfg;
+    }
+
+exit:
+    for (unsigned i = 0; i < dict_cnt; i++) {
+        if (dict[i])
+            av_dict_free(&dict[i]);
+    }
+    av_free(dict);
+    return err;
+}
+
+static int init_metadata(AVFilterContext *ctx)
+{
+    LIBVMAFContext *s = ctx->priv;
+
+    s->cb = av_calloc(1, sizeof(CallbackStruct));
+    if (!s->cb)
+        return AVERROR(ENOMEM);
+
+    s->cb->s = s;
+
+    return 0;
+}
+#endif
+
 static enum VmafLogLevel log_level_map(int log_level)
 {
     switch (log_level) {
@@ -441,6 +687,16 @@  static av_cold int init(AVFilterContext *ctx)
     if (err)
         return AVERROR(EINVAL);
 
+#if CONFIG_LIBVMAF_METADATA_FILTER
+    err = init_metadata(ctx);
+    if (err)
+        return err;
+
+    err = parse_metadata_handlers(ctx);
+    if (err)
+        return err;
+#endif
+
     err = parse_models(ctx);
     if (err)
         return err;
@@ -567,6 +823,21 @@  static av_cold void uninit(AVFilterContext *ctx)
                "problem flushing libvmaf context.\n");
     }
 
+#if CONFIG_LIBVMAF_METADATA_FILTER
+    if (s->metadata_cfg_list.metadata_cfgs) {
+        for (unsigned i = 0; i < s->metadata_cfg_list.metadata_cfg_cnt; i++) {
+            av_free(s->metadata_cfg_list.metadata_cfgs[i].feature_name);
+        }
+        av_free(s->metadata_cfg_list.metadata_cfgs);
+    }
+
+    err = free_frame_list(&s->cb->frame_list);
+    if (err) {
+        av_log(ctx, AV_LOG_ERROR,
+               "problem freeing frame list.\n");
+    }
+#endif
+
     for (unsigned i = 0; i < s->model_cnt; i++) {
         double vmaf_score;
         err = vmaf_score_pooled(s->vmaf, s->model[i], pool_method_map(s->pool),