vf_psnr: add channel weighting based on chroma subsampling.

Also add per-channel psnr stream averages to final log message.

Signed-off-by: Michael Niedermayer <michael@niedermayer.cc>

1 files changed, 23 insertions, 8 deletions

diff --git a/libavfilter/vf_psnr.c b/libavfilter/vf_psnr.c
index 406be881a6..b2c6531621 100644
--- a/libavfilter/vf_psnr.c
+++ b/libavfilter/vf_psnr.c
@@ -25,6 +25,7 @@
  * Caculate the PSNR between two input videos.
  */
 
+#include "libavutil/avstring.h"
 #include "libavutil/opt.h"
 #include "libavutil/pixdesc.h"
 #include "avfilter.h"
@@ -37,7 +38,7 @@
 typedef struct PSNRContext {
     const AVClass *class;
     FFDualInputContext dinput;
-    double mse, min_mse, max_mse;
+    double mse, min_mse, max_mse, mse_comp[4];
     uint64_t nb_frames;
     FILE *stats_file;
     char *stats_file_str;
@@ -48,6 +49,7 @@ typedef struct PSNRContext {
     int nb_components;
     int planewidth[4];
     int planeheight[4];
+    double planeweight[4];
 
     void (*compute_mse)(struct PSNRContext *s,
                         const uint8_t *m[4], const int ml[4],
@@ -158,13 +160,14 @@ static AVFrame *do_psnr(AVFilterContext *ctx, AVFrame *main,
                        main->width, main->height, comp_mse);
 
     for (j = 0; j < s->nb_components; j++)
-        mse += comp_mse[j];
-    mse /= s->nb_components;
+        mse += comp_mse[j] * s->planeweight[j];
 
     s->min_mse = FFMIN(s->min_mse, mse);
     s->max_mse = FFMAX(s->max_mse, mse);
 
     s->mse += mse;
+    for (j = 0; j < s->nb_components; j++)
+        s->mse_comp[j] += comp_mse[j];
     s->nb_frames++;
 
     for (j = 0; j < s->nb_components; j++) {
@@ -243,6 +246,7 @@ static int config_input_ref(AVFilterLink *inlink)
     const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(inlink->format);
     AVFilterContext *ctx  = inlink->dst;
     PSNRContext *s = ctx->priv;
+    unsigned sum;
     int j;
 
     s->nb_components = desc->nb_components;
@@ -290,14 +294,17 @@ static int config_input_ref(AVFilterLink *inlink)
     s->comps[2] = s->is_rgb ? 'b' : 'v' ;
     s->comps[3] = 'a';
 
-    for (j = 0; j < s->nb_components; j++)
-        s->average_max += s->max[j];
-    s->average_max /= s->nb_components;
-
     s->planeheight[1] = s->planeheight[2] = FF_CEIL_RSHIFT(inlink->h, desc->log2_chroma_h);
     s->planeheight[0] = s->planeheight[3] = inlink->h;
     s->planewidth[1]  = s->planewidth[2]  = FF_CEIL_RSHIFT(inlink->w, desc->log2_chroma_w);
     s->planewidth[0]  = s->planewidth[3]  = inlink->w;
+    sum = 0;
+    for (j = 0; j < s->nb_components; j++)
+        sum += s->planeheight[j] * s->planewidth[j];
+    for (j = 0; j < s->nb_components; j++) {
+        s->planeweight[j] = (double) s->planeheight[j] * s->planewidth[j] / sum;
+        s->average_max += s->max[j] * s->planeweight[j];
+    }
 
     s->compute_mse = desc->comp[0].depth_minus1 > 7 ? compute_images_mse_16bit : compute_images_mse;
 
@@ -339,7 +346,15 @@ static av_cold void uninit(AVFilterContext *ctx)
     PSNRContext *s = ctx->priv;
 
     if (s->nb_frames > 0) {
-        av_log(ctx, AV_LOG_INFO, "PSNR average:%0.2f min:%0.2f max:%0.2f\n",
+        int j;
+        char buf[256];
+
+        buf[0] = 0;
+        for (j = 0; j < s->nb_components; j++)
+            av_strlcatf(buf, sizeof(buf), " %c:%0.2f", s->comps[j],
+                        get_psnr(s->mse_comp[j], s->nb_frames, s->max[j]));
+        av_log(ctx, AV_LOG_INFO, "PSNR%s average:%0.2f min:%0.2f max:%0.2f\n",
+               buf,
                get_psnr(s->mse, s->nb_frames, s->average_max),
                get_psnr(s->max_mse, 1, s->average_max),
                get_psnr(s->min_mse, 1, s->average_max));