avfilter/vf_vignette: add a wavelet mode reversing the vignette with noise suppression

Dividing by the attenuation amplifies the noise and the quantization
steps of the dark areas by the same factor. The wavelet mode performs
the division in the pixel domain, then transforms the result and its
per pixel noise variance, propagated through the squared filter taps as
if the samples were uncorrelated, with the undecimated 9/7 wavelet.
Each detail coefficient is Wiener shrunk toward 0 with a signal
variance estimated from a 7x7 box of coefficients at the dilation of
the level, weighted by their precision and counted only when it exceeds
the noise by 3 standard errors of independent coefficients. The lowpass
of the top level averages over large areas and is kept, so the shade of
heavily attenuated areas survives.

Assisted-by: Claude

doc/filters.texi

@@ -26388,10 +26388,22 @@ The larger the distance from the central point, the brighter the image becomes.
 This can be used to reverse a vignette effect, though there is no automatic
 detection to extract the lens @option{angle} and other settings (yet). It can
 also be used to create a burning effect.
+
+@item wavelet
+Like @samp{backward}, but the noise amplified in the dark areas is suppressed
+in an undecimated 9/7 wavelet domain: detail coefficients are shrunk toward 0
+according to the local attenuation and the @option{noise} level, while the
+low frequencies, which average over larger areas, are kept.
+It needs about 150 bytes of memory per pixel.
 @end table
 
 Default value is @samp{forward}.
 
+@item noise
+Set the standard deviation of the input noise in 8-bit levels, used by the
+@samp{wavelet} mode. Quantization alone contributes about 0.3-0.4.
+Default is @code{0.5}.
+
 @item eval
 Set evaluation mode for the expressions (@option{angle}, @option{x0}, @option{y0}).
 

libavfilter/Makefile

@@ -573,7 +573,7 @@ OBJS-$(CONFIG_VIBRANCE_FILTER)               += vf_vibrance.o
 OBJS-$(CONFIG_VIDSTABDETECT_FILTER)          += vidstabutils.o vf_vidstabdetect.o
 OBJS-$(CONFIG_VIDSTABTRANSFORM_FILTER)       += vidstabutils.o vf_vidstabtransform.o
 OBJS-$(CONFIG_VIF_FILTER)                    += vf_vif.o framesync.o
-OBJS-$(CONFIG_VIGNETTE_FILTER)               += vf_vignette.o
+OBJS-$(CONFIG_VIGNETTE_FILTER)               += vf_vignette.o wavelet97.o
 OBJS-$(CONFIG_VMAFMOTION_FILTER)             += vf_vmafmotion.o framesync.o
 OBJS-$(CONFIG_VPP_AMF_FILTER)                += vf_vpp_amf.o scale_eval.o vf_amf_common.o
 OBJS-$(CONFIG_FRC_AMF_FILTER)                += vf_frc_amf.o vf_amf_common.o

libavfilter/vf_vignette.c

@@ -23,10 +23,12 @@
 #include "libavutil/mem.h"
 #include "libavutil/opt.h"
 #include "libavutil/eval.h"
+#include "libavutil/internal.h"
 #include "libavutil/pixdesc.h"
 #include "avfilter.h"
 #include "filters.h"
 #include "video.h"
+#include "wavelet97.h"
 
 static const char *const var_names[] = {
     "w",    // stream width
@@ -56,11 +58,21 @@ enum EvalMode {
     EVAL_MODE_NB
 };
 
+enum Mode {
+    MODE_FORWARD,
+    MODE_BACKWARD,
+    MODE_WAVELET,
+};
+
 typedef struct VignetteContext {
     const AVClass *class;
     const AVPixFmtDescriptor *desc;
-    int backward;
+    int mode;                           ///< Mode
     int eval_mode;                      ///< EvalMode
+    double noise;
+    double sqcoeff[2][5];
+    float *wplanes;
+    size_t wplane_size;
 #define DEF_EXPR_FIELDS(name) AVExpr *name##_pexpr; char *name##_expr; double name
     DEF_EXPR_FIELDS(angle);
     DEF_EXPR_FIELDS(x0);
@@ -83,14 +95,16 @@ static const AVOption vignette_options[] = {
     { "a",     "set lens angle", OFFSET(angle_expr), AV_OPT_TYPE_STRING, {.str="PI/5"}, .flags = FLAGS },
     { "x0", "set circle center position on x-axis", OFFSET(x0_expr), AV_OPT_TYPE_STRING, {.str="w/2"}, .flags = FLAGS },
     { "y0", "set circle center position on y-axis", OFFSET(y0_expr), AV_OPT_TYPE_STRING, {.str="h/2"}, .flags = FLAGS },
-    { "mode", "set forward/backward mode", OFFSET(backward), AV_OPT_TYPE_INT, {.i64 = 0}, 0, 1, FLAGS, .unit = "mode" },
-        { "forward",  NULL, 0, AV_OPT_TYPE_CONST, {.i64 = 0}, INT_MIN, INT_MAX, FLAGS, .unit = "mode"},
-        { "backward", NULL, 0, AV_OPT_TYPE_CONST, {.i64 = 1}, INT_MIN, INT_MAX, FLAGS, .unit = "mode"},
+    { "mode", "set forward/backward mode", OFFSET(mode), AV_OPT_TYPE_INT, {.i64 = MODE_FORWARD}, 0, MODE_WAVELET, FLAGS, .unit = "mode" },
+        { "forward",  NULL, 0, AV_OPT_TYPE_CONST, {.i64 = MODE_FORWARD},  INT_MIN, INT_MAX, FLAGS, .unit = "mode"},
+        { "backward", NULL, 0, AV_OPT_TYPE_CONST, {.i64 = MODE_BACKWARD}, INT_MIN, INT_MAX, FLAGS, .unit = "mode"},
+        { "wavelet",  "backward with noise suppression in the wavelet domain", 0, AV_OPT_TYPE_CONST, {.i64 = MODE_WAVELET}, INT_MIN, INT_MAX, FLAGS, .unit = "mode"},
     { "eval", "specify when to evaluate expressions", OFFSET(eval_mode), AV_OPT_TYPE_INT, {.i64 = EVAL_MODE_INIT}, 0, EVAL_MODE_NB-1, FLAGS, .unit = "eval" },
         { "init",  "eval expressions once during initialization", 0, AV_OPT_TYPE_CONST, {.i64=EVAL_MODE_INIT},  .flags = FLAGS, .unit = "eval" },
         { "frame", "eval expressions for each frame",             0, AV_OPT_TYPE_CONST, {.i64=EVAL_MODE_FRAME}, .flags = FLAGS, .unit = "eval" },
     { "dither", "set dithering", OFFSET(do_dither), AV_OPT_TYPE_BOOL, {.i64 = 1}, 0, 1, FLAGS },
     { "aspect", "set aspect ratio", OFFSET(aspect), AV_OPT_TYPE_RATIONAL, {.dbl = 1}, 0, DBL_MAX, .flags = FLAGS },
+    { "noise", "set input noise standard deviation for wavelet mode", OFFSET(noise), AV_OPT_TYPE_DOUBLE, {.dbl = 0.5}, 0.01, 255, FLAGS },
     { NULL }
 };
 
@@ -113,6 +127,11 @@ static av_cold int init(AVFilterContext *ctx)
     PARSE_EXPR(angle);
     PARSE_EXPR(x0);
     PARSE_EXPR(y0);
+
+    // squared taps propagate the noise variance as if the samples were uncorrelated
+    for (int i = 0; i < 2; i++)
+        for (int j = 0; j < 5; j++)
+            s->sqcoeff[i][j] = ff_wavelet97_coeff[i][j] * ff_wavelet97_coeff[i][j];
     return 0;
 }
 
@@ -120,6 +139,7 @@ static av_cold void uninit(AVFilterContext *ctx)
 {
     VignetteContext *s = ctx->priv;
     av_freep(&s->fmap);
+    av_freep(&s->wplanes);
     av_expr_free(s->angle_pexpr);
     av_expr_free(s->x0_pexpr);
     av_expr_free(s->y0_pexpr);
@@ -173,7 +193,7 @@ static void update_context(VignetteContext *s, AVFilterLink *inlink, AVFrame *fr
 
     s->angle = av_clipf(s->angle, 0, M_PI_2);
 
-    if (s->backward) {
+    if (s->mode == MODE_BACKWARD) {
         for (y = 0; y < inlink->h; y++) {
             for (x = 0; x < inlink->w; x++) {
                 const double f = get_natural_factor(s, x, y);
@@ -204,6 +224,139 @@ static inline double get_dither_value(VignetteContext *s)
     return dv;
 }
 
+static int wavelet_levels(int w, int h)
+{
+    return FFMAX(av_log2(FFMIN(w, h)) - 2, 0);
+}
+
+#define BOX_RADIUS 3
+
+static void box_filter(float *dst, const float *src, float *temp,
+                       int linesize, int step, int w, int h)
+{
+    for (int y = 0; y < h; y++) {
+        for (int x = 0; x < w; x++) {
+            double sum = 0;
+
+            for (int k = -BOX_RADIUS; k <= BOX_RADIUS; k++)
+                sum += src[y * linesize + avpriv_mirror(x + k * step, w - 1)];
+            temp[y * linesize + x] = sum / (2 * BOX_RADIUS + 1);
+        }
+    }
+    for (int y = 0; y < h; y++) {
+        for (int x = 0; x < w; x++) {
+            double sum = 0;
+
+            for (int k = -BOX_RADIUS; k <= BOX_RADIUS; k++)
+                sum += temp[avpriv_mirror(y + k * step, h - 1) * linesize + x];
+            dst[y * linesize + x] = sum / (2 * BOX_RADIUS + 1);
+        }
+    }
+}
+
+static void shrink_band(float *band, const float *var, float *temp[3],
+                        int linesize, int step, int w, int h)
+{
+    // energy a box must have to count as signal: noise + 3 standard errors, assuming independent coefficients
+    const double threshold = 1 + 3 * sqrt(2.0 / ((2 * BOX_RADIUS + 1) * (2 * BOX_RADIUS + 1)));
+
+    for (int y = 0; y < h; y++) {
+        for (int x = 0; x < w; x++) {
+            const int i = y * linesize + x;
+
+            temp[0][i] = band[i] * band[i] / var[i];
+            temp[1][i] = 1 / var[i];
+        }
+    }
+    box_filter(temp[0], temp[0], temp[2], linesize, step, w, h);
+    box_filter(temp[1], temp[1], temp[2], linesize, step, w, h);
+
+    for (int y = 0; y < h; y++) {
+        for (int x = 0; x < w; x++) {
+            const int i = y * linesize + x;
+            const float signal = FFMAX(temp[0][i] - threshold, 0) / temp[1][i];
+
+            band[i] *= signal / (signal + var[i]);
+        }
+    }
+}
+
+static void filter_plane_wavelet(VignetteContext *s,
+                                 uint8_t *dst, int dst_linesize,
+                                 const uint8_t *src, int src_linesize, int pixel_step,
+                                 int w, int h, int hsub, int vsub, int bias)
+{
+    const int linesize   = s->fmap_linesize;
+    const size_t plane_size = s->wplane_size;
+    const int levels     = wavelet_levels(w, h);
+    float *cur       = s->wplanes;
+    float *next      = cur + plane_size;
+    float *var       = next + plane_size;
+    float *temp[2]   = { var + plane_size, var + 2 * plane_size };
+    float *vbands[4] = { var + 3 * plane_size, var + 4 * plane_size,
+                         var + 5 * plane_size, var + 6 * plane_size };
+    float *hbands    = var + 7 * plane_size;
+
+    for (int y = 0; y < h; y++) {
+        for (int x = 0; x < w; x++) {
+            const double a = FFMAX(s->fmap[(y << vsub) * linesize + (x << hsub)], 1e-6);
+
+            cur[y * linesize + x] = av_clipd((src[y * src_linesize + x * pixel_step] - bias) / a,
+                                             -bias, 255 - bias);
+            var[y * linesize + x] = s->noise * s->noise / (a * a);
+        }
+    }
+
+    for (int i = 0; i < levels; i++) {
+        float *bands[4] = { next, hbands + 3 * i * plane_size,
+                            hbands + (3 * i + 1) * plane_size, hbands + (3 * i + 2) * plane_size };
+
+        float *scratch[3] = { temp[0], temp[1], cur };
+
+        ff_wavelet97_decompose(bands,  cur, temp, ff_wavelet97_coeff, linesize, 1 << i, w, h);
+        ff_wavelet97_decompose(vbands, var, temp, s->sqcoeff,         linesize, 1 << i, w, h);
+        for (int j = 1; j < 4; j++)
+            shrink_band(bands[j], vbands[j], scratch, linesize, 1 << i, w, h);
+        FFSWAP(float *, cur, next);
+        FFSWAP(float *, var, vbands[0]);
+    }
+
+    for (int i = levels - 1; i >= 0; i--) {
+        float *bands[4] = { cur, hbands + 3 * i * plane_size,
+                            hbands + (3 * i + 1) * plane_size, hbands + (3 * i + 2) * plane_size };
+
+        ff_wavelet97_compose(next, bands, temp, linesize, 1 << i, w, h);
+        FFSWAP(float *, cur, next);
+    }
+
+    for (int y = 0; y < h; y++)
+        for (int x = 0; x < w; x++)
+            dst[y * dst_linesize + x * pixel_step] =
+                av_clip_uint8(cur[y * linesize + x] + bias + get_dither_value(s));
+}
+
+static void filter_frame_wavelet(VignetteContext *s, AVFrame *out, const AVFrame *in)
+{
+    if (s->desc->flags & AV_PIX_FMT_FLAG_RGB) {
+        for (int c = 0; c < 3; c++)
+            filter_plane_wavelet(s, out->data[0] + c, out->linesize[0],
+                                 in->data[0] + c, in->linesize[0], 3,
+                                 in->width, in->height, 0, 0, 0);
+        return;
+    }
+
+    for (int plane = 0; plane < 4 && in->data[plane] && in->linesize[plane]; plane++) {
+        const int chroma = plane == 1 || plane == 2;
+        const int hsub = chroma ? s->desc->log2_chroma_w : 0;
+        const int vsub = chroma ? s->desc->log2_chroma_h : 0;
+
+        filter_plane_wavelet(s, out->data[plane], out->linesize[plane],
+                             in->data[plane], in->linesize[plane], 1,
+                             AV_CEIL_RSHIFT(in->width, hsub), AV_CEIL_RSHIFT(in->height, vsub),
+                             hsub, vsub, chroma ? 127 : 0);
+    }
+}
+
 static int filter_frame(AVFilterLink *inlink, AVFrame *in)
 {
     unsigned x, y, direct = 0;
@@ -227,7 +380,9 @@ static int filter_frame(AVFilterLink *inlink, AVFrame *in)
     if (s->eval_mode == EVAL_MODE_FRAME)
         update_context(s, inlink, in);
 
-    if (s->desc->flags & AV_PIX_FMT_FLAG_RGB) {
+    if (s->mode == MODE_WAVELET) {
+        filter_frame_wavelet(s, out, in);
+    } else if (s->desc->flags & AV_PIX_FMT_FLAG_RGB) {
         uint8_t       *dst = out->data[0];
         const uint8_t *src = in ->data[0];
         const float *fmap = s->fmap;
@@ -319,6 +474,13 @@ static int config_props(AVFilterLink *inlink)
     if (!s->fmap)
         return AVERROR(ENOMEM);
 
+    if (s->mode == MODE_WAVELET) {
+        s->wplane_size = (size_t)s->fmap_linesize * inlink->h;
+        s->wplanes = av_malloc_array(s->wplane_size, (9 + 3 * wavelet_levels(inlink->w, inlink->h)) * sizeof(*s->wplanes));
+        if (!s->wplanes)
+            return AVERROR(ENOMEM);
+    }
+
     if (s->eval_mode == EVAL_MODE_INIT)
         update_context(s, inlink, NULL);