libavfilter/vf_psnr.c

   1 /*
   2  * Copyright (c) 2011 Roger Pau Monné <roger.pau@entel.upc.edu>
   3  * Copyright (c) 2011 Stefano Sabatini
   4  * Copyright (c) 2013 Paul B Mahol
   5  *
   6  * This file is part of FFmpeg.
   7  *
   8  * FFmpeg is free software; you can redistribute it and/or
   9  * modify it under the terms of the GNU Lesser General Public
  10  * License as published by the Free Software Foundation; either
  11  * version 2.1 of the License, or (at your option) any later version.
  12  *
  13  * FFmpeg is distributed in the hope that it will be useful,
  14  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  15  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  16  * Lesser General Public License for more details.
  17  *
  18  * You should have received a copy of the GNU Lesser General Public
  19  * License along with FFmpeg; if not, write to the Free Software
  20  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  21  */
  22
  23 /**
  24  * @file
  25  * Calculate the PSNR between two input videos.
  26  */
  27
  28 #include "libavutil/avstring.h"
  29 #include "libavutil/file_open.h"
  30 #include "libavutil/mem.h"
  31 #include "libavutil/opt.h"
  32 #include "libavutil/pixdesc.h"
  33 #include "avfilter.h"
  34 #include "drawutils.h"
  35 #include "filters.h"
  36 #include "framesync.h"
  37 #include "psnr.h"
  38
  39 typedef struct PSNRContext {
  40     const AVClass *class;
  41     FFFrameSync fs;
  42     double mse, min_mse, max_mse, mse_comp[4];
  43     uint64_t nb_frames;
  44     FILE *stats_file;
  45     char *stats_file_str;
  46     int stats_version;
  47     int stats_header_written;
  48     int stats_add_max;
  49     int max[4], average_max;
  50     int is_rgb;
  51     uint8_t rgba_map[4];
  52     char comps[4];
  53     int nb_components;
  54     int nb_threads;
  55     int planewidth[4];
  56     int planeheight[4];
  57     double planeweight[4];
  58     uint64_t **score;
  59     PSNRDSPContext dsp;
  60 } PSNRContext;
  61
  62 #define OFFSET(x) offsetof(PSNRContext, x)
  63 #define FLAGS AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_VIDEO_PARAM
  64
  65 static const AVOption psnr_options[] = {
  66     {"stats_file", "Set file where to store per-frame difference information", OFFSET(stats_file_str), AV_OPT_TYPE_STRING, {.str=NULL}, 0, 0, FLAGS },
  67     {"f",          "Set file where to store per-frame difference information", OFFSET(stats_file_str), AV_OPT_TYPE_STRING, {.str=NULL}, 0, 0, FLAGS },
  68     {"stats_version", "Set the format version for the stats file.",               OFFSET(stats_version),  AV_OPT_TYPE_INT,    {.i64=1},    1, 2, FLAGS },
  69     {"output_max",  "Add raw stats (max values) to the output log.",            OFFSET(stats_add_max), AV_OPT_TYPE_BOOL, {.i64=0}, 0, 1, FLAGS},
  70     { NULL }
  71 };
  72
  73 FRAMESYNC_DEFINE_CLASS(psnr, PSNRContext, fs);
  74
  75 static inline unsigned pow_2(unsigned base)
  76 {
  77     return base*base;
  78 }
  79
  80 static inline double get_psnr(double mse, uint64_t nb_frames, int max)
  81 {
  82     return 10.0 * log10(pow_2(max) / (mse / nb_frames));
  83 }
  84
  85 typedef struct ThreadData {
  86     const uint8_t *main_data[4];
  87     const uint8_t *ref_data[4];
  88     int main_linesize[4];
  89     int ref_linesize[4];
  90     int planewidth[4];
  91     int planeheight[4];
  92     uint64_t **score;
  93     int nb_components;
  94     PSNRDSPContext *dsp;
  95 } ThreadData;
  96
  97 static
  98 int compute_images_mse(AVFilterContext *ctx, void *arg,
  99                        int jobnr, int nb_jobs)
 100 {
 101     ThreadData *td = arg;
 102     uint64_t *score = td->score[jobnr];
 103
 104     for (int c = 0; c < td->nb_components; c++) {
 105         const int outw = td->planewidth[c];
 106         const int outh = td->planeheight[c];
 107         const int slice_start = (outh * jobnr) / nb_jobs;
 108         const int slice_end = (outh * (jobnr+1)) / nb_jobs;
 109         const int ref_linesize = td->ref_linesize[c];
 110         const int main_linesize = td->main_linesize[c];
 111         const uint8_t *main_line = td->main_data[c] + main_linesize * slice_start;
 112         const uint8_t *ref_line = td->ref_data[c] + ref_linesize * slice_start;
 113         uint64_t m = 0;
 114         for (int i = slice_start; i < slice_end; i++) {
 115             m += td->dsp->sse_line(main_line, ref_line, outw);
 116             ref_line += ref_linesize;
 117             main_line += main_linesize;
 118         }
 119         score[c] = m;
 120     }
 121
 122     return 0;
 123 }
 124
 125 static void set_meta(AVDictionary **metadata, const char *key, char comp, float d)
 126 {
 127     char value[128];
 128     snprintf(value, sizeof(value), "%f", d);
 129     if (comp) {
 130         char key2[128];
 131         snprintf(key2, sizeof(key2), "%s%c", key, comp);
 132         av_dict_set(metadata, key2, value, 0);
 133     } else {
 134         av_dict_set(metadata, key, value, 0);
 135     }
 136 }
 137
 138 static int do_psnr(FFFrameSync *fs)
 139 {
 140     AVFilterContext *ctx = fs->parent;
 141     PSNRContext *s = ctx->priv;
 142     AVFrame *master, *ref;
 143     double comp_mse[4], mse = 0.;
 144     uint64_t comp_sum[4] = { 0 };
 145     AVDictionary **metadata;
 146     ThreadData td;
 147     int ret;
 148
 149     ret = ff_framesync_dualinput_get(fs, &master, &ref);
 150     if (ret < 0)
 151         return ret;
 152     if (ctx->is_disabled || !ref)
 153         return ff_filter_frame(ctx->outputs[0], master);
 154     metadata = &master->metadata;
 155
 156     td.nb_components = s->nb_components;
 157     td.dsp = &s->dsp;
 158     td.score = s->score;
 159     for (int c = 0; c < s->nb_components; c++) {
 160         td.main_data[c] = master->data[c];
 161         td.ref_data[c] = ref->data[c];
 162         td.main_linesize[c] = master->linesize[c];
 163         td.ref_linesize[c] = ref->linesize[c];
 164         td.planewidth[c] = s->planewidth[c];
 165         td.planeheight[c] = s->planeheight[c];
 166     }
 167
 168     if (master->color_range != ref->color_range) {
 169         av_log(ctx, AV_LOG_WARNING, "master and reference "
 170                "frames use different color ranges (%s != %s)\n",
 171                av_color_range_name(master->color_range),
 172                av_color_range_name(ref->color_range));
 173     }
 174
 175     ff_filter_execute(ctx, compute_images_mse, &td, NULL,
 176                       FFMIN(s->planeheight[1], s->nb_threads));
 177
 178     for (int j = 0; j < s->nb_threads; j++) {
 179         for (int c = 0; c < s->nb_components; c++)
 180             comp_sum[c] += s->score[j][c];
 181     }
 182
 183     for (int c = 0; c < s->nb_components; c++)
 184         comp_mse[c] = comp_sum[c] / ((double)s->planewidth[c] * s->planeheight[c]);
 185
 186     for (int c = 0; c < s->nb_components; c++)
 187         mse += comp_mse[c] * s->planeweight[c];
 188
 189     s->min_mse = FFMIN(s->min_mse, mse);
 190     s->max_mse = FFMAX(s->max_mse, mse);
 191
 192     s->mse += mse;
 193
 194     for (int j = 0; j < s->nb_components; j++)
 195         s->mse_comp[j] += comp_mse[j];
 196     s->nb_frames++;
 197
 198     for (int j = 0; j < s->nb_components; j++) {
 199         int c = s->is_rgb ? s->rgba_map[j] : j;
 200         set_meta(metadata, "lavfi.psnr.mse.", s->comps[j], comp_mse[c]);
 201         set_meta(metadata, "lavfi.psnr.psnr.", s->comps[j], get_psnr(comp_mse[c], 1, s->max[c]));
 202     }
 203     set_meta(metadata, "lavfi.psnr.mse_avg", 0, mse);
 204     set_meta(metadata, "lavfi.psnr.psnr_avg", 0, get_psnr(mse, 1, s->average_max));
 205
 206     if (s->stats_file) {
 207         if (s->stats_version == 2 && !s->stats_header_written) {
 208             fprintf(s->stats_file, "psnr_log_version:2 fields:n");
 209             fprintf(s->stats_file, ",mse_avg");
 210             for (int j = 0; j < s->nb_components; j++) {
 211                 fprintf(s->stats_file, ",mse_%c", s->comps[j]);
 212             }
 213             fprintf(s->stats_file, ",psnr_avg");
 214             for (int j = 0; j < s->nb_components; j++) {
 215                 fprintf(s->stats_file, ",psnr_%c", s->comps[j]);
 216             }
 217             if (s->stats_add_max) {
 218                 fprintf(s->stats_file, ",max_avg");
 219                 for (int j = 0; j < s->nb_components; j++) {
 220                     fprintf(s->stats_file, ",max_%c", s->comps[j]);
 221                 }
 222             }
 223             fprintf(s->stats_file, "\n");
 224             s->stats_header_written = 1;
 225         }
 226         fprintf(s->stats_file, "n:%"PRId64" mse_avg:%0.2f ", s->nb_frames, mse);
 227         for (int j = 0; j < s->nb_components; j++) {
 228             int c = s->is_rgb ? s->rgba_map[j] : j;
 229             fprintf(s->stats_file, "mse_%c:%0.2f ", s->comps[j], comp_mse[c]);
 230         }
 231         fprintf(s->stats_file, "psnr_avg:%0.2f ", get_psnr(mse, 1, s->average_max));
 232         for (int j = 0; j < s->nb_components; j++) {
 233             int c = s->is_rgb ? s->rgba_map[j] : j;
 234             fprintf(s->stats_file, "psnr_%c:%0.2f ", s->comps[j],
 235                     get_psnr(comp_mse[c], 1, s->max[c]));
 236         }
 237         if (s->stats_version == 2 && s->stats_add_max) {
 238             fprintf(s->stats_file, "max_avg:%d ", s->average_max);
 239             for (int j = 0; j < s->nb_components; j++) {
 240                 int c = s->is_rgb ? s->rgba_map[j] : j;
 241                 fprintf(s->stats_file, "max_%c:%d ", s->comps[j], s->max[c]);
 242             }
 243         }
 244         fprintf(s->stats_file, "\n");
 245     }
 246
 247     return ff_filter_frame(ctx->outputs[0], master);
 248 }
 249
 250 static av_cold int init(AVFilterContext *ctx)
 251 {
 252     PSNRContext *s = ctx->priv;
 253
 254     s->min_mse = +INFINITY;
 255     s->max_mse = -INFINITY;
 256
 257     if (s->stats_file_str) {
 258         if (s->stats_version < 2 && s->stats_add_max) {
 259             av_log(ctx, AV_LOG_ERROR,
 260                 "stats_add_max was specified but stats_version < 2.\n" );
 261             return AVERROR(EINVAL);
 262         }
 263         if (!strcmp(s->stats_file_str, "-")) {
 264             s->stats_file = stdout;
 265         } else {
 266             s->stats_file = avpriv_fopen_utf8(s->stats_file_str, "w");
 267             if (!s->stats_file) {
 268                 int err = AVERROR(errno);
 269                 av_log(ctx, AV_LOG_ERROR, "Could not open stats file %s: %s\n",
 270                        s->stats_file_str, av_err2str(err));
 271                 return err;
 272             }
 273         }
 274     }
 275
 276     s->fs.on_event = do_psnr;
 277     return 0;
 278 }
 279
 280 static const enum AVPixelFormat pix_fmts[] = {
 281     AV_PIX_FMT_GRAY8, AV_PIX_FMT_GRAY9, AV_PIX_FMT_GRAY10, AV_PIX_FMT_GRAY12, AV_PIX_FMT_GRAY14, AV_PIX_FMT_GRAY16,
 282 #define PF_NOALPHA(suf) AV_PIX_FMT_YUV420##suf,  AV_PIX_FMT_YUV422##suf,  AV_PIX_FMT_YUV444##suf
 283 #define PF_ALPHA(suf)   AV_PIX_FMT_YUVA420##suf, AV_PIX_FMT_YUVA422##suf, AV_PIX_FMT_YUVA444##suf
 284 #define PF(suf)         PF_NOALPHA(suf), PF_ALPHA(suf)
 285     PF(P), PF(P9), PF(P10), PF_NOALPHA(P12), PF_NOALPHA(P14), PF(P16),
 286     AV_PIX_FMT_YUV440P, AV_PIX_FMT_YUV411P, AV_PIX_FMT_YUV410P,
 287     AV_PIX_FMT_YUVJ411P, AV_PIX_FMT_YUVJ420P, AV_PIX_FMT_YUVJ422P,
 288     AV_PIX_FMT_YUVJ440P, AV_PIX_FMT_YUVJ444P,
 289     AV_PIX_FMT_GBRP, AV_PIX_FMT_GBRP9, AV_PIX_FMT_GBRP10,
 290     AV_PIX_FMT_GBRP12, AV_PIX_FMT_GBRP14, AV_PIX_FMT_GBRP16,
 291     AV_PIX_FMT_GBRAP, AV_PIX_FMT_GBRAP10, AV_PIX_FMT_GBRAP12, AV_PIX_FMT_GBRAP16,
 292     AV_PIX_FMT_NONE
 293 };
 294
 295 static int config_input_ref(AVFilterLink *inlink)
 296 {
 297     const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(inlink->format);
 298     AVFilterContext *ctx  = inlink->dst;
 299     PSNRContext *s = ctx->priv;
 300     double average_max;
 301     unsigned sum;
 302     int j;
 303
 304     s->nb_threads = ff_filter_get_nb_threads(ctx);
 305     s->nb_components = desc->nb_components;
 306     if (ctx->inputs[0]->w != ctx->inputs[1]->w ||
 307         ctx->inputs[0]->h != ctx->inputs[1]->h) {
 308         av_log(ctx, AV_LOG_ERROR, "Width and height of input videos must be same.\n");
 309         return AVERROR(EINVAL);
 310     }
 311
 312     s->max[0] = (1 << desc->comp[0].depth) - 1;
 313     s->max[1] = (1 << desc->comp[1].depth) - 1;
 314     s->max[2] = (1 << desc->comp[2].depth) - 1;
 315     s->max[3] = (1 << desc->comp[3].depth) - 1;
 316
 317     s->is_rgb = ff_fill_rgba_map(s->rgba_map, inlink->format) >= 0;
 318     s->comps[0] = s->is_rgb ? 'r' : 'y' ;
 319     s->comps[1] = s->is_rgb ? 'g' : 'u' ;
 320     s->comps[2] = s->is_rgb ? 'b' : 'v' ;
 321     s->comps[3] = 'a';
 322
 323     s->planeheight[1] = s->planeheight[2] = AV_CEIL_RSHIFT(inlink->h, desc->log2_chroma_h);
 324     s->planeheight[0] = s->planeheight[3] = inlink->h;
 325     s->planewidth[1]  = s->planewidth[2]  = AV_CEIL_RSHIFT(inlink->w, desc->log2_chroma_w);
 326     s->planewidth[0]  = s->planewidth[3]  = inlink->w;
 327     sum = 0;
 328     for (j = 0; j < s->nb_components; j++)
 329         sum += s->planeheight[j] * s->planewidth[j];
 330     average_max = 0;
 331     for (j = 0; j < s->nb_components; j++) {
 332         s->planeweight[j] = (double) s->planeheight[j] * s->planewidth[j] / sum;
 333         average_max += s->max[j] * s->planeweight[j];
 334     }
 335     s->average_max = lrint(average_max);
 336
 337     ff_psnr_init(&s->dsp, desc->comp[0].depth);
 338
 339     s->score = av_calloc(s->nb_threads, sizeof(*s->score));
 340     if (!s->score)
 341         return AVERROR(ENOMEM);
 342
 343     for (int t = 0; t < s->nb_threads; t++) {
 344         s->score[t] = av_calloc(s->nb_components, sizeof(*s->score[0]));
 345         if (!s->score[t])
 346             return AVERROR(ENOMEM);
 347     }
 348
 349     return 0;
 350 }
 351
 352 static int config_output(AVFilterLink *outlink)
 353 {
 354     AVFilterContext *ctx = outlink->src;
 355     PSNRContext *s = ctx->priv;
 356     AVFilterLink *mainlink = ctx->inputs[0];
 357     FilterLink *il = ff_filter_link(mainlink);
 358     FilterLink *ol = ff_filter_link(outlink);
 359     int ret;
 360
 361     ret = ff_framesync_init_dualinput(&s->fs, ctx);
 362     if (ret < 0)
 363         return ret;
 364     outlink->w = mainlink->w;
 365     outlink->h = mainlink->h;
 366     outlink->time_base = mainlink->time_base;
 367     outlink->sample_aspect_ratio = mainlink->sample_aspect_ratio;
 368     ol->frame_rate = il->frame_rate;
 369     if ((ret = ff_framesync_configure(&s->fs)) < 0)
 370         return ret;
 371
 372     outlink->time_base = s->fs.time_base;
 373
 374     if (av_cmp_q(mainlink->time_base, outlink->time_base) ||
 375         av_cmp_q(ctx->inputs[1]->time_base, outlink->time_base))
 376         av_log(ctx, AV_LOG_WARNING, "not matching timebases found between first input: %d/%d and second input %d/%d, results may be incorrect!\n",
 377                mainlink->time_base.num, mainlink->time_base.den,
 378                ctx->inputs[1]->time_base.num, ctx->inputs[1]->time_base.den);
 379
 380     return 0;
 381 }
 382
 383 static int activate(AVFilterContext *ctx)
 384 {
 385     PSNRContext *s = ctx->priv;
 386     return ff_framesync_activate(&s->fs);
 387 }
 388
 389 static av_cold void uninit(AVFilterContext *ctx)
 390 {
 391     PSNRContext *s = ctx->priv;
 392
 393     if (s->nb_frames > 0) {
 394         int j;
 395         char buf[256];
 396
 397         buf[0] = 0;
 398         for (j = 0; j < s->nb_components; j++) {
 399             int c = s->is_rgb ? s->rgba_map[j] : j;
 400             av_strlcatf(buf, sizeof(buf), " %c:%f", s->comps[j],
 401                         get_psnr(s->mse_comp[c], s->nb_frames, s->max[c]));
 402         }
 403         av_log(ctx, AV_LOG_INFO, "PSNR%s average:%f min:%f max:%f\n",
 404                buf,
 405                get_psnr(s->mse, s->nb_frames, s->average_max),
 406                get_psnr(s->max_mse, 1, s->average_max),
 407                get_psnr(s->min_mse, 1, s->average_max));
 408     }
 409
 410     ff_framesync_uninit(&s->fs);
 411     for (int t = 0; t < s->nb_threads && s->score; t++)
 412         av_freep(&s->score[t]);
 413     av_freep(&s->score);
 414
 415     if (s->stats_file && s->stats_file != stdout)
 416         fclose(s->stats_file);
 417 }
 418
 419 static const AVFilterPad psnr_inputs[] = {
 420     {
 421         .name         = "main",
 422         .type         = AVMEDIA_TYPE_VIDEO,
 423     },{
 424         .name         = "reference",
 425         .type         = AVMEDIA_TYPE_VIDEO,
 426         .config_props = config_input_ref,
 427     },
 428 };
 429
 430 static const AVFilterPad psnr_outputs[] = {
 431     {
 432         .name          = "default",
 433         .type          = AVMEDIA_TYPE_VIDEO,
 434         .config_props  = config_output,
 435     },
 436 };
 437
 438 const FFFilter ff_vf_psnr = {
 439     .p.name        = "psnr",
 440     .p.description = NULL_IF_CONFIG_SMALL("Calculate the PSNR between two video streams."),
 441     .p.priv_class  = &psnr_class,
 442     .p.flags       = AVFILTER_FLAG_SUPPORT_TIMELINE_INTERNAL |
 443                      AVFILTER_FLAG_SLICE_THREADS             |
 444                      AVFILTER_FLAG_METADATA_ONLY,
 445     .preinit       = psnr_framesync_preinit,
 446     .init          = init,
 447     .uninit        = uninit,
 448     .activate      = activate,
 449     .priv_size     = sizeof(PSNRContext),
 450     FILTER_INPUTS(psnr_inputs),
 451     FILTER_OUTPUTS(psnr_outputs),
 452     FILTER_PIXFMTS_ARRAY(pix_fmts),
 453 };