]> git.sesse.net Git - ffmpeg/blobdiff - libavfilter/vf_libvmaf.c
hwcontext_vulkan: reorder structure fields and add spaces in between
[ffmpeg] / libavfilter / vf_libvmaf.c
index dfe474c40cb2ad8a46f8407b468649a77548cb30..d2621c7a9f1c0a7e56304e3a46bf864bda1fbd69 100644 (file)
@@ -43,6 +43,7 @@ typedef struct LIBVMAFContext {
     int width;
     int height;
     double vmaf_score;
+    int vmaf_thread_created;
     pthread_t vmaf_thread;
     pthread_mutex_t lock;
     pthread_cond_t cond;
@@ -61,6 +62,9 @@ typedef struct LIBVMAFContext {
     int ssim;
     int ms_ssim;
     char *pool;
+    int n_threads;
+    int n_subsample;
+    int enable_conf_interval;
     int error;
 } LIBVMAFContext;
 
@@ -70,13 +74,16 @@ typedef struct LIBVMAFContext {
 static const AVOption libvmaf_options[] = {
     {"model_path",  "Set the model to be used for computing vmaf.",                     OFFSET(model_path), AV_OPT_TYPE_STRING, {.str="/usr/local/share/model/vmaf_v0.6.1.pkl"}, 0, 1, FLAGS},
     {"log_path",  "Set the file path to be used to store logs.",                        OFFSET(log_path), AV_OPT_TYPE_STRING, {.str=NULL}, 0, 1, FLAGS},
-    {"log_fmt",  "Set the format of the log (xml or json).",                            OFFSET(log_fmt), AV_OPT_TYPE_STRING, {.str=NULL}, 0, 1, FLAGS},
+    {"log_fmt",  "Set the format of the log (csv, json or xml).",                       OFFSET(log_fmt), AV_OPT_TYPE_STRING, {.str=NULL}, 0, 1, FLAGS},
     {"enable_transform",  "Enables transform for computing vmaf.",                      OFFSET(enable_transform), AV_OPT_TYPE_BOOL, {.i64=0}, 0, 1, FLAGS},
     {"phone_model",  "Invokes the phone model that will generate higher VMAF scores.",  OFFSET(phone_model), AV_OPT_TYPE_BOOL, {.i64=0}, 0, 1, FLAGS},
     {"psnr",  "Enables computing psnr along with vmaf.",                                OFFSET(psnr), AV_OPT_TYPE_BOOL, {.i64=0}, 0, 1, FLAGS},
     {"ssim",  "Enables computing ssim along with vmaf.",                                OFFSET(ssim), AV_OPT_TYPE_BOOL, {.i64=0}, 0, 1, FLAGS},
     {"ms_ssim",  "Enables computing ms-ssim along with vmaf.",                          OFFSET(ms_ssim), AV_OPT_TYPE_BOOL, {.i64=0}, 0, 1, FLAGS},
     {"pool",  "Set the pool method to be used for computing vmaf.",                     OFFSET(pool), AV_OPT_TYPE_STRING, {.str=NULL}, 0, 1, FLAGS},
+    {"n_threads", "Set number of threads to be used when computing vmaf.",              OFFSET(n_threads), AV_OPT_TYPE_INT, {.i64=0}, 0, UINT_MAX, FLAGS},
+    {"n_subsample", "Set interval for frame subsampling used when computing vmaf.",     OFFSET(n_subsample), AV_OPT_TYPE_INT, {.i64=1}, 1, UINT_MAX, FLAGS},
+    {"enable_conf_interval",  "Enables confidence interval.",                           OFFSET(enable_conf_interval), AV_OPT_TYPE_BOOL, {.i64=0}, 0, 1, FLAGS},
     { NULL }
 };
 
@@ -103,6 +110,7 @@ FRAMESYNC_DEFINE_CLASS(libvmaf, LIBVMAFContext, fs);
         const type *main_ptr = (const type *) s->gmain->data[0];                \
         \
         float *ptr = ref_data;                                                  \
+        float factor = 1.f / (1 << (bits - 8));                                 \
         \
         int h = s->height;                                                      \
         int w = s->width;                                                       \
@@ -111,7 +119,7 @@ FRAMESYNC_DEFINE_CLASS(libvmaf, LIBVMAFContext, fs);
         \
         for (i = 0; i < h; i++) {                                               \
             for ( j = 0; j < w; j++) {                                          \
-                ptr[j] = (float)ref_ptr[j];                                     \
+                ptr[j] = ref_ptr[j] * factor;                                   \
             }                                                                   \
             ref_ptr += ref_stride / sizeof(*ref_ptr);                           \
             ptr += stride / sizeof(*ptr);                                       \
@@ -121,7 +129,7 @@ FRAMESYNC_DEFINE_CLASS(libvmaf, LIBVMAFContext, fs);
         \
         for (i = 0; i < h; i++) {                                               \
             for (j = 0; j < w; j++) {                                           \
-                ptr[j] = (float)main_ptr[j];                                    \
+                ptr[j] = main_ptr[j] * factor;                                  \
             }                                                                   \
             main_ptr += main_stride / sizeof(*main_ptr);                        \
             ptr += stride / sizeof(*ptr);                                       \
@@ -130,6 +138,8 @@ FRAMESYNC_DEFINE_CLASS(libvmaf, LIBVMAFContext, fs);
     \
     ret = !s->frame_set;                                                        \
     \
+    av_frame_unref(s->gref);                                                    \
+    av_frame_unref(s->gmain);                                                   \
     s->frame_set = 0;                                                           \
     \
     pthread_cond_signal(&s->cond);                                              \
@@ -163,7 +173,8 @@ static void compute_vmaf_score(LIBVMAFContext *s)
                             read_frame, s, s->model_path, s->log_path,
                             s->log_fmt, 0, 0, s->enable_transform,
                             s->phone_model, s->psnr, s->ssim,
-                            s->ms_ssim, s->pool);
+                            s->ms_ssim, s->pool,
+                            s->n_threads, s->n_subsample, s->enable_conf_interval);
 }
 
 static void *call_vmaf(void *ctx)
@@ -224,8 +235,12 @@ static av_cold int init(AVFilterContext *ctx)
 
     s->gref = av_frame_alloc();
     s->gmain = av_frame_alloc();
+    if (!s->gref || !s->gmain)
+        return AVERROR(ENOMEM);
+
     s->error = 0;
 
+    s->vmaf_thread_created = 0;
     pthread_mutex_init(&s->lock, NULL);
     pthread_cond_init (&s->cond, NULL);
 
@@ -273,6 +288,7 @@ static int config_input_ref(AVFilterLink *inlink)
         av_log(ctx, AV_LOG_ERROR, "Thread creation failed.\n");
         return AVERROR(EINVAL);
     }
+    s->vmaf_thread_created = 1;
 
     return 0;
 }
@@ -315,7 +331,11 @@ static av_cold void uninit(AVFilterContext *ctx)
     pthread_cond_signal(&s->cond);
     pthread_mutex_unlock(&s->lock);
 
-    pthread_join(s->vmaf_thread, NULL);
+    if (s->vmaf_thread_created)
+    {
+        pthread_join(s->vmaf_thread, NULL);
+        s->vmaf_thread_created = 0;
+    }
 
     av_frame_free(&s->gref);
     av_frame_free(&s->gmain);
@@ -345,7 +365,7 @@ static const AVFilterPad libvmaf_outputs[] = {
     { NULL }
 };
 
-AVFilter ff_vf_libvmaf = {
+const AVFilter ff_vf_libvmaf = {
     .name          = "libvmaf",
     .description   = NULL_IF_CONFIG_SMALL("Calculate the VMAF between two video streams."),
     .preinit       = libvmaf_framesync_preinit,