avfilter: add apsnr filter

1 files changed, 64 insertions, 3 deletions

diff --git a/libavfilter/af_asdr.c b/libavfilter/af_asdr.c
index 7d778b7f6b..b0401804f6 100644
--- a/libavfilter/af_asdr.c
+++ b/libavfilter/af_asdr.c
@@ -18,6 +18,8 @@
  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  */
 
+#include <float.h>
+
 #include "libavutil/channel_layout.h"
 #include "libavutil/common.h"
 
@@ -27,6 +29,8 @@
 
 typedef struct AudioSDRContext {
     int channels;
+    uint64_t nb_samples;
+    double max;
     double *sum_u;
     double *sum_uv;
 
@@ -67,6 +71,34 @@ static int sdr_##name(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs)\
 SDR_FILTER(fltp, float)
 SDR_FILTER(dblp, double)
 
+#define PSNR_FILTER(name, type)                                               \
+static int psnr_##name(AVFilterContext *ctx, void *arg, int jobnr,int nb_jobs)\
+{                                                                             \
+    AudioSDRContext *s = ctx->priv;                                           \
+    AVFrame *u = s->cache[0];                                                 \
+    AVFrame *v = s->cache[1];                                                 \
+    const int channels = u->ch_layout.nb_channels;                            \
+    const int start = (channels * jobnr) / nb_jobs;                           \
+    const int end = (channels * (jobnr+1)) / nb_jobs;                         \
+    const int nb_samples = u->nb_samples;                                     \
+                                                                              \
+    for (int ch = start; ch < end; ch++) {                                    \
+        const type *const us = (type *)u->extended_data[ch];                  \
+        const type *const vs = (type *)v->extended_data[ch];                  \
+        double sum_uv = 0.;                                                   \
+                                                                              \
+        for (int n = 0; n < nb_samples; n++)                                  \
+            sum_uv += (us[n] - vs[n]) * (us[n] - vs[n]);                      \
+                                                                              \
+        s->sum_uv[ch] += sum_uv;                                              \
+    }                                                                         \
+                                                                              \
+    return 0;                                                                 \
+}
+
+PSNR_FILTER(fltp, float)
+PSNR_FILTER(dblp, double)
+
 static int activate(AVFilterContext *ctx)
 {
     AudioSDRContext *s = ctx->priv;
@@ -97,6 +129,7 @@ static int activate(AVFilterContext *ctx)
         out = s->cache[0];
         s->cache[0] = NULL;
 
+        s->nb_samples += available;
         return ff_filter_frame(outlink, out);
     }
 
@@ -126,7 +159,12 @@ static int config_output(AVFilterLink *outlink)
     AudioSDRContext *s = ctx->priv;
 
     s->channels = inlink->ch_layout.nb_channels;
-    s->filter = inlink->format == AV_SAMPLE_FMT_FLTP ? sdr_fltp : sdr_dblp;
+
+    if (!strcmp(ctx->filter->name, "asdr"))
+        s->filter = inlink->format == AV_SAMPLE_FMT_FLTP ? sdr_fltp : sdr_dblp;
+    else
+        s->filter = inlink->format == AV_SAMPLE_FMT_FLTP ? psnr_fltp : psnr_dblp;
+    s->max = inlink->format == AV_SAMPLE_FMT_FLTP ? FLT_MAX : DBL_MAX;
 
     s->sum_u  = av_calloc(outlink->ch_layout.nb_channels, sizeof(*s->sum_u));
     s->sum_uv = av_calloc(outlink->ch_layout.nb_channels, sizeof(*s->sum_uv));
@@ -140,8 +178,16 @@ static av_cold void uninit(AVFilterContext *ctx)
 {
     AudioSDRContext *s = ctx->priv;
 
-    for (int ch = 0; ch < s->channels; ch++)
-        av_log(ctx, AV_LOG_INFO, "SDR ch%d: %g dB\n", ch, 20. * log10(s->sum_u[ch] / s->sum_uv[ch]));
+    if (!strcmp(ctx->filter->name, "asdr")) {
+        for (int ch = 0; ch < s->channels; ch++)
+            av_log(ctx, AV_LOG_INFO, "SDR ch%d: %g dB\n", ch, 20. * log10(s->sum_u[ch] / s->sum_uv[ch]));
+    } else {
+        for (int ch = 0; ch < s->channels; ch++) {
+            double psnr = s->sum_uv[ch] > 0.0 ? 2.0 * log(s->max) - log(s->nb_samples / s->sum_uv[ch]) : INFINITY;
+
+            av_log(ctx, AV_LOG_INFO, "PSNR ch%d: %g dB\n", ch, psnr);
+        }
+    }
 
     av_frame_free(&s->cache[0]);
     av_frame_free(&s->cache[1]);
@@ -183,3 +229,18 @@ const AVFilter ff_af_asdr = {
     FILTER_SAMPLEFMTS(AV_SAMPLE_FMT_FLTP,
                       AV_SAMPLE_FMT_DBLP),
 };
+
+const AVFilter ff_af_apsnr = {
+    .name           = "apsnr",
+    .description    = NULL_IF_CONFIG_SMALL("Measure Audio Peak Signal-to-Noise Ratio."),
+    .priv_size      = sizeof(AudioSDRContext),
+    .activate       = activate,
+    .uninit         = uninit,
+    .flags          = AVFILTER_FLAG_METADATA_ONLY |
+                      AVFILTER_FLAG_SLICE_THREADS |
+                      AVFILTER_FLAG_SUPPORT_TIMELINE_INTERNAL,
+    FILTER_INPUTS(inputs),
+    FILTER_OUTPUTS(outputs),
+    FILTER_SAMPLEFMTS(AV_SAMPLE_FMT_FLTP,
+                      AV_SAMPLE_FMT_DBLP),
+};