[FFmpeg-devel] [PATCH] avfilter/vf_psnr: >8 bit support

[Paul B Mahol]

Signed-off-by: Paul B Mahol <onemda at gmail.com>
---
 libavfilter/vf_psnr.c | 75 ++++++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 59 insertions(+), 16 deletions(-)

diff --git a/libavfilter/vf_psnr.c b/libavfilter/vf_psnr.c
index ef638ff..8aa1543 100644
--- a/libavfilter/vf_psnr.c
+++ b/libavfilter/vf_psnr.c
@@ -48,6 +48,11 @@ typedef struct PSNRContext {
     int nb_components;
     int planewidth[4];
     int planeheight[4];
+
+    void (*compute_mse)(struct PSNRContext *s,
+                        const uint8_t *m[4], const int ml[4],
+                        const uint8_t *r[4], const int rl[4],
+                        int w, int h, double mse[4]);
 } PSNRContext;
 
 #define OFFSET(x) offsetof(PSNRContext, x)
@@ -61,7 +66,7 @@ static const AVOption psnr_options[] = {
 
 AVFILTER_DEFINE_CLASS(psnr);
 
-static inline int pow2(int base)
+static inline unsigned pow2(unsigned base)
 {
     return base*base;
 }
@@ -98,6 +103,33 @@ void compute_images_mse(PSNRContext *s,
     }
 }
 
+static inline
+void compute_images_mse_16bit(PSNRContext *s,
+                        const uint8_t *main_data[4], const int main_linesizes[4],
+                        const uint8_t *ref_data[4], const int ref_linesizes[4],
+                        int w, int h, double mse[4])
+{
+    int i, c, j;
+
+    for (c = 0; c < s->nb_components; c++) {
+        const int outw = s->planewidth[c];
+        const int outh = s->planeheight[c];
+        const uint16_t *main_line = (uint16_t *)main_data[c];
+        const uint16_t *ref_line = (uint16_t *)ref_data[c];
+        const int ref_linesize = ref_linesizes[c] / 2;
+        const int main_linesize = main_linesizes[c] / 2;
+        int64_t m = 0;
+
+        for (i = 0; i < outh; i++) {
+            for (j = 0; j < outw; j++)
+                m += pow2(main_line[j] - ref_line[j]);
+            ref_line += ref_linesize;
+            main_line += main_linesize;
+        }
+        mse[c] = m / (outw * outh);
+    }
+}
+
 static void set_meta(AVDictionary **metadata, const char *key, char comp, float d)
 {
     char value[128];
@@ -119,8 +151,8 @@ static AVFrame *do_psnr(AVFilterContext *ctx, AVFrame *main,
     int j, c;
     AVDictionary **metadata = avpriv_frame_get_metadatap(main);
 
-    compute_images_mse(s, (const uint8_t **)main->data, main->linesize,
-                       (const uint8_t **)ref->data, ref->linesize,
+    s->compute_mse(s, (const uint8_t **)main->data, main->linesize,
+                      (const uint8_t **)ref->data, ref->linesize,
                        main->width, main->height, comp_mse);
 
     for (j = 0; j < s->nb_components; j++)
@@ -191,6 +223,12 @@ static int query_formats(AVFilterContext *ctx)
         AV_PIX_FMT_YUVJ444P, AV_PIX_FMT_YUVJ440P, AV_PIX_FMT_YUVJ422P,
         AV_PIX_FMT_YUVJ420P, AV_PIX_FMT_YUVJ411P,
         AV_PIX_FMT_YUVA444P, AV_PIX_FMT_YUVA422P, AV_PIX_FMT_YUVA420P,
+        AV_PIX_FMT_YUV444P9, AV_PIX_FMT_YUV422P9, AV_PIX_FMT_YUV420P9,
+        AV_PIX_FMT_YUV444P10, AV_PIX_FMT_YUV422P10, AV_PIX_FMT_YUV420P10,
+        AV_PIX_FMT_YUV444P12, AV_PIX_FMT_YUV422P12, AV_PIX_FMT_YUV420P12,
+        AV_PIX_FMT_YUV444P14, AV_PIX_FMT_YUV422P14, AV_PIX_FMT_YUV420P14,
+        AV_PIX_FMT_YUV444P16, AV_PIX_FMT_YUV422P16, AV_PIX_FMT_YUV420P16,
+        AV_PIX_FMT_YUVA444P16, AV_PIX_FMT_YUVA422P16, AV_PIX_FMT_YUVA420P16,
         AV_PIX_FMT_NONE
     };
 
@@ -217,21 +255,24 @@ static int config_input_ref(AVFilterLink *inlink)
     }
 
     switch (inlink->format) {
-    case AV_PIX_FMT_YUV410P:
-    case AV_PIX_FMT_YUV411P:
-    case AV_PIX_FMT_YUV420P:
-    case AV_PIX_FMT_YUV422P:
-    case AV_PIX_FMT_YUV440P:
-    case AV_PIX_FMT_YUV444P:
-    case AV_PIX_FMT_YUVA420P:
-    case AV_PIX_FMT_YUVA422P:
-    case AV_PIX_FMT_YUVA444P:
-        s->max[0] = 235;
-        s->max[3] = 255;
-        s->max[1] = s->max[2] = 240;
+    case AV_PIX_FMT_GRAY8:
+    case AV_PIX_FMT_GBRP:
+    case AV_PIX_FMT_GBRAP:
+    case AV_PIX_FMT_YUVJ411P:
+    case AV_PIX_FMT_YUVJ420P:
+    case AV_PIX_FMT_YUVJ422P:
+    case AV_PIX_FMT_YUVJ440P:
+    case AV_PIX_FMT_YUVJ444P:
+        s->max[0] = (1 << (desc->comp[0].depth_minus1 + 1)) - 1;
+        s->max[1] = (1 << (desc->comp[1].depth_minus1 + 1)) - 1;
+        s->max[2] = (1 << (desc->comp[2].depth_minus1 + 1)) - 1;
+        s->max[3] = (1 << (desc->comp[3].depth_minus1 + 1)) - 1;
         break;
     default:
-        s->max[0] = s->max[1] = s->max[2] = s->max[3] = 255;
+        s->max[0] = 235 * (1 << (desc->comp[0].depth_minus1 - 7));
+        s->max[1] = 240 * (1 << (desc->comp[1].depth_minus1 - 7));
+        s->max[2] = 240 * (1 << (desc->comp[2].depth_minus1 - 7));
+        s->max[3] = (1 << (desc->comp[3].depth_minus1 + 1)) - 1;
     }
 
     s->is_rgb = ff_fill_rgba_map(s->rgba_map, inlink->format) >= 0;
@@ -249,6 +290,8 @@ static int config_input_ref(AVFilterLink *inlink)
     s->planewidth[1]  = s->planewidth[2]  = FF_CEIL_RSHIFT(inlink->w, desc->log2_chroma_w);
     s->planewidth[0]  = s->planewidth[3]  = inlink->w;
 
+    s->compute_mse = desc->comp[0].depth_minus1 > 7 ? compute_images_mse_16bit : compute_images_mse;
+
     return 0;
 }
 
-- 
1.7.11.2