Add Feature propagation to metadata

Author: yigithanyigit

libvmaf/include/libvmaf/libvmaf.h

@@ -253,6 +253,51 @@ int vmaf_score_at_index_model_collection(VmafContext *vmaf,
 int vmaf_feature_score_at_index(VmafContext *vmaf, const char *feature_name,
                                 double *score, unsigned index);
 
+/**
+ * Metadata structure.
+ *
+ * @param feature_name   Name of the feature to fetch.
+ *
+ * @param picture_index  Picture index.
+ *
+ * @param score          Score.
+ *
+ * @note This structure is used to pass metadata to a callback function.
+ */
+typedef struct VmafMetadata {
+    char *feature_name;
+    unsigned picture_index;
+    double score;
+} VmafMetadata;
+
+/**
+ * Metadata configuration.
+ *
+ * @param feature_name Name of the feature to fetch.
+ *
+ * @param callback     Callback to receive metadata.
+ *
+ * @param data         User data to pass to the callback.
+ */
+typedef struct VmafMetadataConfiguration {
+    char *feature_name;
+    void (*callback)(void *data, VmafMetadata *metadata);
+    void *data;
+} VmafMetadataConfiguration;
+
+/**
+ * Register a callback to receive VMAF metadata.
+ *
+ * @param vmaf The VMAF context allocated with `vmaf_init()`.
+ *
+ * @param cfg  Metadata configuration.
+ *
+ *
+ * @return 0 on success, or < 0 (a negative errno code) on error.
+ */
+
+int vmaf_register_metadata_handler(VmafContext *vmaf, VmafMetadataConfiguration cfg);
+
 /**
  * Pooled VMAF score for a specific interval.
  *

libvmaf/src/feature/feature_collector.c

@@ -24,9 +24,12 @@
 #include <string.h>
 
 #include "dict.h"
+#include "metadata_handler.h"
 #include "feature_collector.h"
 #include "feature_name.h"
+#include "libvmaf/libvmaf.h"
 #include "log.h"
+#include "predict.h"
 
 static int aggregate_vector_init(AggregateVector *aggregate_vector)
 {
@@ -215,8 +218,12 @@ int vmaf_feature_collector_init(VmafFeatureCollector **const feature_collector)
     if (err) goto free_feature_vector;
     err = pthread_mutex_init(&(fc->lock), NULL);
     if (err) goto free_aggregate_vector;
+    err = vmaf_metadata_init(&(fc->metadata));
+    if (err) goto free_mutex;
     return 0;
 
+free_mutex:
+    pthread_mutex_destroy(&(fc->lock));
 free_aggregate_vector:
     aggregate_vector_destroy(&(fc->aggregate_vector));
 free_feature_vector:
@@ -227,6 +234,62 @@ int vmaf_feature_collector_init(VmafFeatureCollector **const feature_collector)
     return -ENOMEM;
 }
 
+int vmaf_feature_collector_mount_model(VmafFeatureCollector *feature_collector,
+                                       VmafModel *model)
+{
+    if (!feature_collector) return -EINVAL;
+    if (!model) return -EINVAL;
+
+    VmafPredictModel *m = malloc(sizeof(VmafPredictModel));
+    if (!m) return -ENOMEM;
+    m->model = model;
+    m->next = NULL;
+
+    VmafPredictModel **head = &feature_collector->models;
+    while (*head && (*head)->next != NULL)
+        *head = (*head)->next;
+
+    if (!(*head))
+        *head = m;
+    else
+        (*head)->next = m;
+
+    return 0;
+}
+
+int vmaf_feature_collector_unmount_model(VmafFeatureCollector *feature_collector,
+                                         VmafModel *model)
+{
+    if (!feature_collector) return -EINVAL;
+    if (!model) return -EINVAL;
+
+    VmafPredictModel **head = &feature_collector->models;
+    while (*head && (*head)->model != model)
+        head = &(*head)->next;
+
+    if (!(*head)) return -EINVAL;
+
+    VmafPredictModel *m = *head;
+    *head = m->next;
+    free(m);
+
+    return 0;
+}
+
+int vmaf_feature_collector_register_metadata(VmafFeatureCollector *feature_collector,
+                                             VmafMetadataConfiguration metadata_cfg)
+{
+    if (!feature_collector) return -EINVAL;
+    if (!metadata_cfg.feature_name) return -EINVAL;
+    if (!metadata_cfg.callback) return -EINVAL;
+
+    VmafCallbackList *metadata = feature_collector->metadata;
+    int err = vmaf_metadata_append(metadata, metadata_cfg);
+    if (err) return err;
+
+    return 0;
+}
+
 static FeatureVector *find_feature_vector(VmafFeatureCollector *fc,
                                           const char *feature_name)
 {
@@ -280,6 +343,51 @@ int vmaf_feature_collector_append(VmafFeatureCollector *feature_collector,
     }
 
     err = feature_vector_append(feature_vector, picture_index, score);
+    if (err) goto unlock;
+
+    int res = 0;
+
+    VmafCallbackItem *metadata_iter = feature_collector->metadata ?
+                                      feature_collector->metadata->head : NULL;
+    while (metadata_iter) {
+        // Check current feature name is the same as the metadata feature name
+        if (!strcmp(metadata_iter->metadata_cfg.feature_name, feature_name)) {
+
+            // Call the callback function with the metadata feature name
+            VmafMetadata data = {
+                .feature_name = metadata_iter->metadata_cfg.feature_name,
+                .picture_index = picture_index,
+                .score = score,
+            };
+            metadata_iter->metadata_cfg.callback(metadata_iter->metadata_cfg.data, &data);
+            // Move to the next metadata
+            goto next_metadata;
+        }
+
+        VmafPredictModel *model_iter = feature_collector->models;
+
+        // If metadata feature name is not the same as the current feature feature_name
+        // Check if metadata feature name is the predicted feature
+        while (model_iter) {
+            VmafModel *model = model_iter->model;
+
+            pthread_mutex_unlock(&(feature_collector->lock));
+            res = vmaf_feature_collector_get_score(feature_collector,
+                    model->name, &score, picture_index);
+            pthread_mutex_lock(&(feature_collector->lock));
+
+            if (res) {
+                pthread_mutex_unlock(&(feature_collector->lock));
+                res |= vmaf_predict_score_at_index(model, feature_collector,
+                        picture_index, &score, true, true, 0);
+                pthread_mutex_lock(&(feature_collector->lock));
+            }
+            model_iter = model_iter->next;
+        }
+
+    next_metadata:
+        metadata_iter = metadata_iter->next;
+    }
 
 unlock:
     feature_collector->timer.end = clock();
@@ -338,6 +446,10 @@ void vmaf_feature_collector_destroy(VmafFeatureCollector *feature_collector)
     aggregate_vector_destroy(&(feature_collector->aggregate_vector));
     for (unsigned i = 0; i < feature_collector->cnt; i++)
         feature_vector_destroy(feature_collector->feature_vector[i]);
+    while (feature_collector->models)
+        vmaf_feature_collector_unmount_model(feature_collector,
+                                             feature_collector->models->model);
+    vmaf_metadata_destroy(feature_collector->metadata);
     free(feature_collector->feature_vector);
     pthread_mutex_unlock(&(feature_collector->lock));
     pthread_mutex_destroy(&(feature_collector->lock));

libvmaf/src/feature/feature_collector.h

@@ -24,6 +24,8 @@
 #include <time.h>
 
 #include "dict.h"
+#include "model.h"
+#include "metadata_handler.h"
 
 typedef struct {
     char *name;
@@ -42,20 +44,32 @@ typedef struct {
     unsigned cnt, capacity;
 } AggregateVector;
 
+typedef struct VmafPredictModel {
+    VmafModel *model;
+    struct VmafPredictModel *next;
+} VmafPredictModel;
+
 typedef struct VmafFeatureCollector {
     FeatureVector **feature_vector;
     AggregateVector aggregate_vector;
+    VmafCallbackList *metadata;
+    VmafPredictModel *models;
     unsigned cnt, capacity;
     struct { clock_t begin, end; } timer;
     pthread_mutex_t lock;
 } VmafFeatureCollector;
 
 int vmaf_feature_collector_init(VmafFeatureCollector **const feature_collector);
 
+int vmaf_feature_collector_mount_model(VmafFeatureCollector *feature_collector, VmafModel *model);
+
 int vmaf_feature_collector_append(VmafFeatureCollector *feature_collector,
                                   const char *feature_name, double score,
                                   unsigned index);
 
+int vmaf_feature_collector_register_metadata(VmafFeatureCollector *feature_collector,
+                                             VmafMetadataConfiguration metadata_cfg);
+
 int vmaf_feature_collector_append_with_dict(VmafFeatureCollector *fc,
         VmafDictionary *dict, const char *feature_name, double score,
         unsigned index);

libvmaf/src/libvmaf.c

@@ -32,6 +32,7 @@
 #include "cpu.h"
 #include "feature/feature_extractor.h"
 #include "feature/feature_collector.h"
+#include "metadata_handler.h"
 #include "fex_ctx_vector.h"
 #include "log.h"
 #include "model.h"
@@ -362,6 +363,10 @@ int vmaf_use_features_from_model(VmafContext *vmaf, VmafModel *model)
             return err;
         }
     }
+
+    err = vmaf_feature_collector_mount_model(vmaf->feature_collector, model);
+    if (err) return err;
+
     return 0;
 }
 
@@ -739,6 +744,13 @@ int vmaf_read_pictures(VmafContext *vmaf, VmafPicture *ref, VmafPicture *dist,
     return err;
 }
 
+int vmaf_register_metadata_handler(VmafContext *vmaf, VmafMetadataConfiguration cfg)
+{
+    if (!vmaf) return -EINVAL;
+
+    return vmaf_feature_collector_register_metadata(vmaf->feature_collector, cfg);
+}
+
 int vmaf_feature_score_at_index(VmafContext *vmaf, const char *feature_name,
                                 double *score, unsigned index)
 {
@@ -763,7 +775,7 @@ int vmaf_score_at_index(VmafContext *vmaf, VmafModel *model, double *score,
                                          score, index);
     if (err) {
         err = vmaf_predict_score_at_index(model, vmaf->feature_collector, index,
-                                          score, true, 0);
+                                          score, true, false, 0);
     }
 
     return err;

libvmaf/src/meson.build

@@ -462,6 +462,7 @@ libvmaf_sources = [
     src_dir + 'pdjson.c',
     src_dir + 'log.c',
     src_dir + 'framesync.c',
+    src_dir + 'metadata_handler.c',
 ]
 
 if is_cuda_enabled

libvmaf/src/metadata_handler.c

@@ -0,0 +1,79 @@
+/**
+ *
+ *  Copyright 2016-2020 Netflix, Inc.
+ *
+ *     Licensed under the BSD+Patent License (the "License");
+ *     you may not use this file except in compliance with the License.
+ *     You may obtain a copy of the License at
+ *
+ *         https://opensource.org/licenses/BSDplusPatent
+ *
+ *     Unless required by applicable law or agreed to in writing, software
+ *     distributed under the License is distributed on an "AS IS" BASIS,
+ *     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *     See the License for the specific language governing permissions and
+ *     limitations under the License.
+ *
+ */
+
+#include <errno.h>
+#include <stdlib.h>
+#include <string.h>
+
+#include "metadata_handler.h"
+
+int vmaf_metadata_init(VmafCallbackList **const metadata)
+{
+    if (!metadata) return -EINVAL;
+
+    VmafCallbackList *const metadata_s = *metadata =
+        malloc(sizeof(*metadata_s));
+    if (!metadata_s) goto fail;
+
+    metadata_s->head = NULL;
+
+    return 0;
+
+fail:
+    return -ENOMEM;
+}
+
+int vmaf_metadata_append(VmafCallbackList *metadata, const VmafMetadataConfiguration metadata_cfg)
+{
+    if (!metadata) return -EINVAL;
+
+    VmafCallbackItem *node = malloc(sizeof(*node));
+    if (!node) goto fail;
+    memset(node, 0, sizeof(*node));
+
+    node->metadata_cfg = metadata_cfg;
+
+    if (!metadata->head) {
+        metadata->head = node;
+    } else {
+        VmafCallbackItem *iter = metadata->head;
+        while (iter->next) iter = iter->next;
+        iter->next = node;
+    }
+
+    return 0;
+
+fail:
+    return -ENOMEM;
+}
+
+int vmaf_metadata_destroy(VmafCallbackList *metadata)
+{
+    if (!metadata) return -EINVAL;
+
+    VmafCallbackItem *iter = metadata->head;
+    while (iter) {
+        VmafCallbackItem *next = iter->next;
+        free(iter);
+        iter = next;
+    }
+
+    free(metadata);
+
+    return 0;
+}