From 3eb487cbcaefa703aa362390ae0b90b58024a023 Mon Sep 17 00:00:00 2001
From: Elias Carotti <eliascrt _at_ amazon _dot_ it>
Date: Fri, 15 Sep 2023 20:05:43 +0200
Subject: [PATCH] Add the SSE computation for libx264.

Since libx264 only provides a per-frame per-channel PSNR, this is inverted to get back the SSE.
---
 libavcodec/libx264.c | 37 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 36 insertions(+), 1 deletion(-)

diff --git a/libavcodec/libx264.c b/libavcodec/libx264.c
index 77a9f173b4..4c643c9066 100644
--- a/libavcodec/libx264.c
+++ b/libavcodec/libx264.c
@@ -129,6 +129,8 @@ typedef struct X264Context {
     int roi_warned;

     int mb_info;
+
+    int64_t sse[3];
 } X264Context;

 static void X264_log(void *p, int level, const char *fmt, va_list args)
@@ -726,7 +728,40 @@ FF_ENABLE_DEPRECATION_WARNINGS

     pkt->flags |= AV_PKT_FLAG_KEY*pic_out.b_keyframe;
     if (ret) {
-        ff_side_data_set_encoder_stats(pkt, (pic_out.i_qpplus1 - 1) * FF_QP2LAMBDA, NULL, 0, pict_type);
+        const AVPixFmtDescriptor *pix_desc = av_pix_fmt_desc_get(csp_to_pixfmt(pic_out.img.i_csp));
+        int error_count = 0;
+        int64_t *errors = NULL;
+
+        if (ctx->flags & AV_CODEC_FLAG_PSNR) {
+            double scale[3] = { 1,
+                (1 << pix_desc->log2_chroma_h) * (double)(1 << pix_desc->log2_chroma_w),
+                (1 << pix_desc->log2_chroma_h) * (double)(1 << pix_desc->log2_chroma_w),
+            };
+            double mse;
+            int i;
+
+            error_count = pix_desc->nb_components;
+
+            av_log(ctx, AV_LOG_DEBUG, "PSNR values from libx264: %.3f %.3f %.3f.\n",
+                   pic_out.prop.f_psnr[0], pic_out.prop.f_psnr[1], pic_out.prop.f_psnr[2]);
+
+            for (i = 0; i < pix_desc->nb_components; ++i) {
+                double max_value = (double)(1 << pix_desc->comp[i].depth) - 1.0;
+                double plane_size = ctx->width * (double)ctx->height / scale[i];
+
+                /* psnr = -10 * log10(max_value * max_value / mse) */
+                mse = (max_value * max_value) / pow(10, -pic_out.prop.f_psnr[i] / 10.0);
+
+                /* SSE = MSE * width * height / scale -> because of possible chroma downsampling */
+                x4->sse[i] = (int64_t)floor(mse * plane_size + .5);
+            };
+
+            errors = x4->sse;
+        }
+
+        ff_side_data_set_encoder_stats(pkt, (pic_out.i_qpplus1 - 1) * FF_QP2LAMBDA,
+                                       errors, error_count, pict_type);
+
         if (wallclock)
             ff_side_data_set_prft(pkt, wallclock);
     }
--
2.34.1