virtualenv/FFmpeg
1
0
Fork 0
mirror of https://github.com/FFmpeg/FFmpeg.git synced 2024-12-12 19:18:44 +02:00
Code Issues Releases Activity
cd8515eb70
FFmpeg/libavfilter/vf_libvmaf.c

712 lines
21 KiB
C
Raw Blame History

/*
* Copyright (c) 2017 Ronald S. Bultje <rsbultje@gmail.com>
* Copyright (c) 2017 Ashish Pratap Singh <ashk43712@gmail.com>
*
* This file is part of FFmpeg.
*
* FFmpeg is free software; you can redistribute it and/or
* modify it under the terms of the GNU Lesser General Public
* License as published by the Free Software Foundation; either
* version 2.1 of the License, or (at your option) any later version.
*
* FFmpeg is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
* Lesser General Public License for more details.
*
* You should have received a copy of the GNU Lesser General Public
* License along with FFmpeg; if not, write to the Free Software
* Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
*/
/**
* @file
* Calculate the VMAF between two input videos.
*/
#include <libvmaf.h>
#include "libavutil/avstring.h"
#include "libavutil/opt.h"
#include "libavutil/pixdesc.h"
#include "avfilter.h"
#include "drawutils.h"
#include "formats.h"
#include "framesync.h"
#include "internal.h"
#include "video.h"
typedef struct LIBVMAFContext {
const AVClass *class;
FFFrameSync fs;
char *model_path;
char *log_path;
char *log_fmt;
int enable_transform;
int phone_model;
int psnr;
int ssim;
int ms_ssim;
char *pool;
int n_threads;
int n_subsample;
int enable_conf_interval;
char *model_cfg;
char *feature_cfg;
VmafContext *vmaf;
VmafModel **model;
unsigned model_cnt;
unsigned frame_cnt;
unsigned bpc;
} LIBVMAFContext;
#define OFFSET(x) offsetof(LIBVMAFContext, x)
#define FLAGS AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_VIDEO_PARAM
static const AVOption libvmaf_options[] = {
{"model_path", "use model='path=...'.", OFFSET(model_path), AV_OPT_TYPE_STRING, {.str=NULL}, 0, 1, FLAGS|AV_OPT_FLAG_DEPRECATED},
{"log_path", "Set the file path to be used to write log.", OFFSET(log_path), AV_OPT_TYPE_STRING, {.str=NULL}, 0, 1, FLAGS},
{"log_fmt", "Set the format of the log (csv, json, xml, or sub).", OFFSET(log_fmt), AV_OPT_TYPE_STRING, {.str="xml"}, 0, 1, FLAGS},
{"enable_transform", "use model='enable_transform=true'.", OFFSET(enable_transform), AV_OPT_TYPE_BOOL, {.i64=0}, 0, 1, FLAGS|AV_OPT_FLAG_DEPRECATED},
{"phone_model", "use model='enable_transform=true'.", OFFSET(phone_model), AV_OPT_TYPE_BOOL, {.i64=0}, 0, 1, FLAGS|AV_OPT_FLAG_DEPRECATED},
{"psnr", "use feature='name=psnr'.", OFFSET(psnr), AV_OPT_TYPE_BOOL, {.i64=0}, 0, 1, FLAGS|AV_OPT_FLAG_DEPRECATED},
{"ssim", "use feature='name=ssim'.", OFFSET(ssim), AV_OPT_TYPE_BOOL, {.i64=0}, 0, 1, FLAGS|AV_OPT_FLAG_DEPRECATED},
{"ms_ssim", "use feature='name=ms_ssim'.", OFFSET(ms_ssim), AV_OPT_TYPE_BOOL, {.i64=0}, 0, 1, FLAGS|AV_OPT_FLAG_DEPRECATED},
{"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", "model='enable_conf_interval=true'.", OFFSET(enable_conf_interval), AV_OPT_TYPE_BOOL, {.i64=0}, 0, 1, FLAGS|AV_OPT_FLAG_DEPRECATED},
{"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},
{ NULL }
};
FRAMESYNC_DEFINE_CLASS(libvmaf, LIBVMAFContext, fs);
static enum VmafPixelFormat pix_fmt_map(enum AVPixelFormat av_pix_fmt)
{
switch (av_pix_fmt) {
case AV_PIX_FMT_YUV420P:
case AV_PIX_FMT_YUV420P10LE:
case AV_PIX_FMT_YUV420P12LE:
case AV_PIX_FMT_YUV420P16LE:
return VMAF_PIX_FMT_YUV420P;
case AV_PIX_FMT_YUV422P:
case AV_PIX_FMT_YUV422P10LE:
case AV_PIX_FMT_YUV422P12LE:
case AV_PIX_FMT_YUV422P16LE:
return VMAF_PIX_FMT_YUV422P;
case AV_PIX_FMT_YUV444P:
case AV_PIX_FMT_YUV444P10LE:
case AV_PIX_FMT_YUV444P12LE:
case AV_PIX_FMT_YUV444P16LE:
return VMAF_PIX_FMT_YUV444P;
default:
return VMAF_PIX_FMT_UNKNOWN;
}
}
static int copy_picture_data(AVFrame *src, VmafPicture *dst, unsigned bpc)
{
int err = vmaf_picture_alloc(dst, pix_fmt_map(src->format), bpc,
src->width, src->height);
if (err)
return AVERROR(ENOMEM);
for (unsigned i = 0; i < 3; i++) {
uint8_t *src_data = src->data[i];
uint8_t *dst_data = dst->data[i];
for (unsigned j = 0; j < dst->h[i]; j++) {
memcpy(dst_data, src_data, sizeof(*dst_data) * dst->w[i]);
src_data += src->linesize[i];
dst_data += dst->stride[i];
}
}
return 0;
}
static int do_vmaf(FFFrameSync *fs)
{
AVFilterContext *ctx = fs->parent;
LIBVMAFContext *s = ctx->priv;
VmafPicture pic_ref, pic_dist;
AVFrame *ref, *dist;
int err = 0;
int ret = ff_framesync_dualinput_get(fs, &dist, &ref);
if (ret < 0)
return ret;
if (ctx->is_disabled || !ref)
return ff_filter_frame(ctx->outputs[0], dist);
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 = 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);
}
return ff_filter_frame(ctx->outputs[0], dist);
}
static AVDictionary **delimited_dict_parse(char *str, unsigned *cnt)
{
AVDictionary **dict = NULL;
char *str_copy = NULL;
char *saveptr = NULL;
unsigned cnt2;
int err = 0;
if (!str)
return NULL;
cnt2 = 1;
for (char *p = str; *p; p++) {
if (*p == '|')
cnt2++;
}
dict = av_calloc(cnt2, sizeof(*dict));
if (!dict)
goto fail;
str_copy = av_strdup(str);
if (!str_copy)
goto fail;
*cnt = 0;
for (unsigned i = 0; i < cnt2; i++) {
char *s = av_strtok(i == 0 ? str_copy : NULL, "|", &saveptr);
if (!s)
continue;
err = av_dict_parse_string(&dict[(*cnt)++], s, "=", ":", 0);
if (err)
goto fail;
}
av_free(str_copy);
return dict;
fail:
if (dict) {
for (unsigned i = 0; i < *cnt; i++) {
if (dict[i])
av_dict_free(&dict[i]);
}
av_free(dict);
}
av_free(str_copy);
*cnt = 0;
return NULL;
}
static int parse_features(AVFilterContext *ctx)
{
LIBVMAFContext *s = ctx->priv;
AVDictionary **dict = NULL;
unsigned dict_cnt;
int err = 0;
if (!s->feature_cfg)
return 0;
dict = delimited_dict_parse(s->feature_cfg, &dict_cnt);
if (!dict) {
av_log(ctx, AV_LOG_ERROR,
"could not parse feature config: %s\n", s->feature_cfg);
return AVERROR(EINVAL);
}
for (unsigned i = 0; i < dict_cnt; i++) {
char *feature_name = NULL;
VmafFeatureDictionary *feature_opts_dict = NULL;
AVDictionaryEntry *e = NULL;
while (e = av_dict_get(dict[i], "", e, AV_DICT_IGNORE_SUFFIX)) {
if (av_stristr(e->key, "name")) {
feature_name = e->value;
continue;
}
err = vmaf_feature_dictionary_set(&feature_opts_dict, e->key,
e->value);
if (err) {
av_log(ctx, AV_LOG_ERROR,
"could not set feature option: %s.%s=%s\n",
feature_name, e->key, e->value);
goto exit;
}
}
err = vmaf_use_feature(s->vmaf, feature_name, feature_opts_dict);
if (err) {
av_log(ctx, AV_LOG_ERROR,
"problem during vmaf_use_feature: %s\n", feature_name);
goto exit;
}
}
exit:
for (unsigned i = 0; i < dict_cnt; i++) {
if (dict[i])
av_dict_free(&dict[i]);
}
av_free(dict);
return err;
}
static int parse_models(AVFilterContext *ctx)
{
LIBVMAFContext *s = ctx->priv;
AVDictionary **dict;
unsigned dict_cnt;
int err = 0;
if (!s->model_cfg) return 0;
dict_cnt = 0;
dict = delimited_dict_parse(s->model_cfg, &dict_cnt);
if (!dict) {
av_log(ctx, AV_LOG_ERROR,
"could not parse model config: %s\n", s->model_cfg);
return AVERROR(EINVAL);
}
s->model_cnt = dict_cnt;
s->model = av_calloc(s->model_cnt, sizeof(*s->model));
if (!s->model)
return AVERROR(ENOMEM);
for (unsigned i = 0; i < dict_cnt; i++) {
VmafModelConfig model_cfg = { 0 };
AVDictionaryEntry *e = NULL;
char *version = NULL;
char *path = NULL;
while (e = av_dict_get(dict[i], "", e, AV_DICT_IGNORE_SUFFIX)) {
if (av_stristr(e->key, "disable_clip")) {
model_cfg.flags |= av_stristr(e->value, "true") ?
VMAF_MODEL_FLAG_DISABLE_CLIP : 0;
continue;
}
if (av_stristr(e->key, "enable_transform")) {
model_cfg.flags |= av_stristr(e->value, "true") ?
VMAF_MODEL_FLAG_ENABLE_TRANSFORM : 0;
continue;
}
if (av_stristr(e->key, "name")) {
model_cfg.name = e->value;
continue;
}
if (av_stristr(e->key, "version")) {
version = e->value;
continue;
}
if (av_stristr(e->key, "path")) {
path = e->value;
continue;
}
}
if (version) {
err = vmaf_model_load(&s->model[i], &model_cfg, version);
if (err) {
av_log(ctx, AV_LOG_ERROR,
"could not load libvmaf model with version: %s\n",
version);
goto exit;
}
}
if (path && !s->model[i]) {
err = vmaf_model_load_from_path(&s->model[i], &model_cfg, path);
if (err) {
av_log(ctx, AV_LOG_ERROR,
"could not load libvmaf model with path: %s\n",
path);
goto exit;
}
}
if (!s->model[i]) {
av_log(ctx, AV_LOG_ERROR,
"could not load libvmaf model with config: %s\n",
s->model_cfg);
goto exit;
}
while (e = av_dict_get(dict[i], "", e, AV_DICT_IGNORE_SUFFIX)) {
VmafFeatureDictionary *feature_opts_dict = NULL;
char *feature_opt = NULL;
char *feature_name = av_strtok(e->key, ".", &feature_opt);
if (!feature_opt)
continue;
err = vmaf_feature_dictionary_set(&feature_opts_dict,
feature_opt, e->value);
if (err) {
av_log(ctx, AV_LOG_ERROR,
"could not set feature option: %s.%s=%s\n",
feature_name, feature_opt, e->value);
err = AVERROR(EINVAL);
goto exit;
}
err = vmaf_model_feature_overload(s->model[i], feature_name,
feature_opts_dict);
if (err) {
av_log(ctx, AV_LOG_ERROR,
"could not overload feature: %s\n", feature_name);
err = AVERROR(EINVAL);
goto exit;
}
}
}
for (unsigned i = 0; i < s->model_cnt; i++) {
err = vmaf_use_features_from_model(s->vmaf, s->model[i]);
if (err) {
av_log(ctx, AV_LOG_ERROR,
"problem during vmaf_use_features_from_model\n");
err = AVERROR(EINVAL);
goto exit;
}
}
exit:
for (unsigned i = 0; i < dict_cnt; i++) {
if (dict[i])
av_dict_free(&dict[i]);
}
av_free(dict);
return err;
}
static enum VmafLogLevel log_level_map(int log_level)
{
switch (log_level) {
case AV_LOG_QUIET:
return VMAF_LOG_LEVEL_NONE;
case AV_LOG_ERROR:
return VMAF_LOG_LEVEL_ERROR;
case AV_LOG_WARNING:
return VMAF_LOG_LEVEL_WARNING;
case AV_LOG_INFO:
return VMAF_LOG_LEVEL_INFO;
case AV_LOG_DEBUG:
return VMAF_LOG_LEVEL_DEBUG;
default:
return VMAF_LOG_LEVEL_INFO;
}
}
static int parse_deprecated_options(AVFilterContext *ctx)
{
LIBVMAFContext *s = ctx->priv;
VmafModel *model = NULL;
VmafModelCollection *model_collection = NULL;
enum VmafModelFlags flags = VMAF_MODEL_FLAGS_DEFAULT;
int err = 0;
VmafModelConfig model_cfg = {
.name = "vmaf",
.flags = flags,
};
if (s->enable_transform || s->phone_model)
flags |= VMAF_MODEL_FLAG_ENABLE_TRANSFORM;
if (!s->model_path)
goto extra_metrics_only;
if (s->enable_conf_interval) {
err = vmaf_model_collection_load_from_path(&model, &model_collection,
&model_cfg, s->model_path);
if (err) {
av_log(ctx, AV_LOG_ERROR,
"problem loading model file: %s\n", s->model_path);
goto exit;
}
err = vmaf_use_features_from_model_collection(s->vmaf, model_collection);
if (err) {
av_log(ctx, AV_LOG_ERROR,
"problem loading feature extractors from model file: %s\n",
s->model_path);
goto exit;
}
} else {
err = vmaf_model_load_from_path(&model, &model_cfg, s->model_path);
if (err) {
av_log(ctx, AV_LOG_ERROR,
"problem loading model file: %s\n", s->model_path);
goto exit;
}
err = vmaf_use_features_from_model(s->vmaf, model);
if (err) {
av_log(ctx, AV_LOG_ERROR,
"problem loading feature extractors from model file: %s\n",
s->model_path);
goto exit;
}
}
extra_metrics_only:
if (s->psnr) {
VmafFeatureDictionary *d = NULL;
vmaf_feature_dictionary_set(&d, "enable_chroma", "false");
err = vmaf_use_feature(s->vmaf, "psnr", d);
if (err) {
av_log(ctx, AV_LOG_ERROR,
"problem loading feature extractor: psnr\n");
goto exit;
}
}
if (s->ssim) {
err = vmaf_use_feature(s->vmaf, "float_ssim", NULL);
if (err) {
av_log(ctx, AV_LOG_ERROR,
"problem loading feature extractor: ssim\n");
goto exit;
}
}
if (s->ms_ssim) {
err = vmaf_use_feature(s->vmaf, "float_ms_ssim", NULL);
if (err) {
av_log(ctx, AV_LOG_ERROR,
"problem loading feature extractor: ms_ssim\n");
goto exit;
}
}
exit:
return err;
}
static av_cold int init(AVFilterContext *ctx)
{
LIBVMAFContext *s = ctx->priv;
int err = 0;
VmafConfiguration cfg = {
.log_level = log_level_map(av_log_get_level()),
.n_subsample = s->n_subsample,
.n_threads = s->n_threads,
};
err = vmaf_init(&s->vmaf, cfg);
if (err)
return AVERROR(EINVAL);
err = parse_deprecated_options(ctx);
if (err)
return err;
err = parse_models(ctx);
if (err)
return err;
err = parse_features(ctx);
if (err)
return err;
s->fs.on_event = do_vmaf;
return 0;
}
static const enum AVPixelFormat pix_fmts[] = {
AV_PIX_FMT_YUV444P, AV_PIX_FMT_YUV422P, AV_PIX_FMT_YUV420P,
AV_PIX_FMT_YUV444P10LE, AV_PIX_FMT_YUV422P10LE, AV_PIX_FMT_YUV420P10LE,
AV_PIX_FMT_NONE
};
static int config_input_ref(AVFilterLink *inlink)
{
AVFilterContext *ctx = inlink->dst;
LIBVMAFContext *s = ctx->priv;
const AVPixFmtDescriptor *desc;
int err = 0;
if (ctx->inputs[0]->w != ctx->inputs[1]->w) {
av_log(ctx, AV_LOG_ERROR, "input width must match.\n");
err |= AVERROR(EINVAL);
}
if (ctx->inputs[0]->h != ctx->inputs[1]->h) {
av_log(ctx, AV_LOG_ERROR, "input height must match.\n");
err |= AVERROR(EINVAL);
}
if (ctx->inputs[0]->format != ctx->inputs[1]->format) {
av_log(ctx, AV_LOG_ERROR, "input pix_fmt must match.\n");
err |= AVERROR(EINVAL);
}
if (err)
return err;
desc = av_pix_fmt_desc_get(inlink->format);
s->bpc = desc->comp[0].depth;
return 0;
}
static int config_output(AVFilterLink *outlink)
{
AVFilterContext *ctx = outlink->src;
LIBVMAFContext *s = ctx->priv;
AVFilterLink *mainlink = ctx->inputs[0];
int ret;
ret = ff_framesync_init_dualinput(&s->fs, ctx);
if (ret < 0)
return ret;
outlink->w = mainlink->w;
outlink->h = mainlink->h;
outlink->time_base = mainlink->time_base;
outlink->sample_aspect_ratio = mainlink->sample_aspect_ratio;
outlink->frame_rate = mainlink->frame_rate;
if ((ret = ff_framesync_configure(&s->fs)) < 0)
return ret;
return 0;
}
static int activate(AVFilterContext *ctx)
{
LIBVMAFContext *s = ctx->priv;
return ff_framesync_activate(&s->fs);
}
static enum VmafOutputFormat log_fmt_map(const char *log_fmt)
{
if (log_fmt) {
if (av_stristr(log_fmt, "xml"))
return VMAF_OUTPUT_FORMAT_XML;
if (av_stristr(log_fmt, "json"))
return VMAF_OUTPUT_FORMAT_JSON;
if (av_stristr(log_fmt, "csv"))
return VMAF_OUTPUT_FORMAT_CSV;
if (av_stristr(log_fmt, "sub"))
return VMAF_OUTPUT_FORMAT_SUB;
}
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;
int err = 0;
ff_framesync_uninit(&s->fs);
if (!s->frame_cnt)
goto clean_up;
err = vmaf_read_pictures(s->vmaf, NULL, NULL, 0);
if (err) {
av_log(ctx, AV_LOG_ERROR,
"problem flushing libvmaf context.\n");
}
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),
&vmaf_score, 0, s->frame_cnt - 1);
if (err) {
av_log(ctx, AV_LOG_ERROR,
"problem getting pooled vmaf score.\n");
}
av_log(ctx, AV_LOG_INFO, "VMAF score: %f\n", vmaf_score);
}
if (s->vmaf) {
if (s->log_path && !err)
vmaf_write_output(s->vmaf, s->log_path, log_fmt_map(s->log_fmt));
}
clean_up:
if (s->model) {
for (unsigned i = 0; i < s->model_cnt; i++) {
if (s->model[i])
vmaf_model_destroy(s->model[i]);
}
av_free(s->model);
}
if (s->vmaf)
vmaf_close(s->vmaf);
}
static const AVFilterPad libvmaf_inputs[] = {
{
.name = "main",
.type = AVMEDIA_TYPE_VIDEO,
},{
.name = "reference",
.type = AVMEDIA_TYPE_VIDEO,
.config_props = config_input_ref,
},
};
static const AVFilterPad libvmaf_outputs[] = {
{
.name = "default",
.type = AVMEDIA_TYPE_VIDEO,
.config_props = config_output,
},
};
const AVFilter ff_vf_libvmaf = {
.name = "libvmaf",
.description = NULL_IF_CONFIG_SMALL("Calculate the VMAF between two video streams."),
.preinit = libvmaf_framesync_preinit,
.init = init,
.uninit = uninit,
.activate = activate,
.priv_size = sizeof(LIBVMAFContext),
.priv_class = &libvmaf_class,
FILTER_INPUTS(libvmaf_inputs),
FILTER_OUTPUTS(libvmaf_outputs),
FILTER_PIXFMTS_ARRAY(pix_fmts),
};
Reference in New Issue View Git Blame Copy Permalink