avfilter/vf_deband: add planes coupling mode

Signed-off-by: Paul B Mahol <onemda@gmail.com>

1 files changed, 160 insertions, 2 deletions

diff --git a/libavfilter/vf_deband.c b/libavfilter/vf_deband.c
index fbf27ea04f..ffec037fd4 100644
--- a/libavfilter/vf_deband.c
+++ b/libavfilter/vf_deband.c
@@ -30,6 +30,7 @@
 typedef struct DebandContext {
     const AVClass *class;
 
+    int coupling;
     float threshold[4];
     int range;
     int blur;
@@ -38,6 +39,7 @@ typedef struct DebandContext {
     int nb_components;
     int planewidth[4];
     int planeheight[4];
+    int shift[2];
     int thr[4];
 
     int *x_pos;
@@ -59,6 +61,9 @@ static const AVOption deband_options[] = {
     { "direction", "set direction",           OFFSET(direction),    AV_OPT_TYPE_FLOAT, {.dbl=2*M_PI},-2*M_PI,  2*M_PI, FLAGS },
     { "d",         "set direction",           OFFSET(direction),    AV_OPT_TYPE_FLOAT, {.dbl=2*M_PI},-2*M_PI,  2*M_PI, FLAGS },
     { "blur",      "set blur",                OFFSET(blur),         AV_OPT_TYPE_BOOL,  {.i64=1},           0,       1, FLAGS },
+    { "b",         "set blur",                OFFSET(blur),         AV_OPT_TYPE_BOOL,  {.i64=1},           0,       1, FLAGS },
+    { "coupling",  "set plane coupling",      OFFSET(coupling),     AV_OPT_TYPE_BOOL,  {.i64=0},           0,       1, FLAGS },
+    { "c",         "set plane coupling",      OFFSET(coupling),     AV_OPT_TYPE_BOOL,  {.i64=0},           0,       1, FLAGS },
     { NULL }
 };
 
@@ -66,6 +71,8 @@ AVFILTER_DEFINE_CLASS(deband);
 
 static int query_formats(AVFilterContext *ctx)
 {
+    DebandContext *s = ctx->priv;
+
     static const enum AVPixelFormat pix_fmts[] = {
         AV_PIX_FMT_GRAY8, AV_PIX_FMT_GRAY16,
         AV_PIX_FMT_YUV444P,  AV_PIX_FMT_YUV422P,  AV_PIX_FMT_YUV420P,
@@ -86,7 +93,21 @@ static int query_formats(AVFilterContext *ctx)
         AV_PIX_FMT_YUVA420P16, AV_PIX_FMT_YUVA422P16, AV_PIX_FMT_YUVA444P16,
         AV_PIX_FMT_NONE
     };
-    AVFilterFormats *fmts_list = ff_make_format_list(pix_fmts);
+
+    static const enum AVPixelFormat cpix_fmts[] = {
+        AV_PIX_FMT_YUV444P, AV_PIX_FMT_YUVJ444P,
+        AV_PIX_FMT_YUVA444P, AV_PIX_FMT_YUV444P9,
+        AV_PIX_FMT_YUVA444P9, AV_PIX_FMT_YUVA444P10,
+        AV_PIX_FMT_YUV444P12, AV_PIX_FMT_YUV444P14,
+        AV_PIX_FMT_YUV444P16, AV_PIX_FMT_YUVA444P16,
+        AV_PIX_FMT_GBRP, AV_PIX_FMT_GBRAP,
+        AV_PIX_FMT_GBRP9, AV_PIX_FMT_GBRP10,
+        AV_PIX_FMT_GBRP12, AV_PIX_FMT_GBRP14,
+        AV_PIX_FMT_GBRP16, AV_PIX_FMT_GBRAP16,
+        AV_PIX_FMT_NONE
+    };
+
+    AVFilterFormats *fmts_list = ff_make_format_list(s->coupling ? cpix_fmts : pix_fmts);
     if (!fmts_list)
         return AVERROR(ENOMEM);
 
@@ -158,6 +179,138 @@ static int deband_8_c(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs)
     return 0;
 }
 
+static int deband_8_coupling_c(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs)
+{
+    DebandContext *s = ctx->priv;
+    ThreadData *td = arg;
+    AVFrame *in = td->in;
+    AVFrame *out = td->out;
+    const int start = (s->planeheight[0] *  jobnr   ) / nb_jobs;
+    const int end   = (s->planeheight[0] * (jobnr+1)) / nb_jobs;
+    int x, y, p;
+
+    for (y = start; y < end; y++) {
+        const int pos = y * s->planewidth[0];
+
+        for (x = 0; x < s->planewidth[p]; x++) {
+            const int x_pos = s->x_pos[pos + x];
+            const int y_pos = s->y_pos[pos + x];
+            int avg[4], cmp[4] = { 0 }, src[4];
+
+            for (p = 0; p < s->nb_components; p++) {
+                const uint8_t *src_ptr = (const uint8_t *)in->data[p];
+                const int src_linesize = in->linesize[p];
+                const int thr = s->thr[p];
+                const int w = s->planewidth[p] - 1;
+                const int h = s->planeheight[p] - 1;
+                const int ref0 = src_ptr[av_clip(y +  y_pos, 0, h) * src_linesize + av_clip(x +  x_pos, 0, w)];
+                const int ref1 = src_ptr[av_clip(y + -y_pos, 0, h) * src_linesize + av_clip(x +  x_pos, 0, w)];
+                const int ref2 = src_ptr[av_clip(y + -y_pos, 0, h) * src_linesize + av_clip(x + -x_pos, 0, w)];
+                const int ref3 = src_ptr[av_clip(y +  y_pos, 0, h) * src_linesize + av_clip(x + -x_pos, 0, w)];
+                const int src0 = src_ptr[y * src_linesize + x];
+
+                src[p] = src0;
+                avg[p] = get_avg(ref0, ref1, ref2, ref3);
+
+                if (s->blur) {
+                    cmp[p] = FFABS(src0 - avg[p]) < thr;
+                } else {
+                    cmp[p] = (FFABS(src0 - ref0) < thr) &&
+                             (FFABS(src0 - ref1) < thr) &&
+                             (FFABS(src0 - ref2) < thr) &&
+                             (FFABS(src0 - ref3) < thr);
+                }
+            }
+
+            for (p = 0; p < s->nb_components; p++)
+                if (!cmp[p])
+                    break;
+            if (p == s->nb_components) {
+                for (p = 0; p < s->nb_components; p++) {
+                    const int dst_linesize = out->linesize[p];
+
+                    out->data[p][y * dst_linesize + x] = avg[p];
+                }
+            } else {
+                for (p = 0; p < s->nb_components; p++) {
+                    const int dst_linesize = out->linesize[p];
+
+                    out->data[p][y * dst_linesize + x] = src[p];
+                }
+            }
+        }
+    }
+
+    return 0;
+}
+
+static int deband_16_coupling_c(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs)
+{
+    DebandContext *s = ctx->priv;
+    ThreadData *td = arg;
+    AVFrame *in = td->in;
+    AVFrame *out = td->out;
+    const int start = (s->planeheight[0] *  jobnr   ) / nb_jobs;
+    const int end   = (s->planeheight[0] * (jobnr+1)) / nb_jobs;
+    int x, y, p, z;
+
+    for (y = start; y < end; y++) {
+        const int pos = y * s->planewidth[0];
+
+        for (x = 0; x < s->planewidth[p]; x++) {
+            const int x_pos = s->x_pos[pos + x];
+            const int y_pos = s->y_pos[pos + x];
+            int avg[4], cmp[4] = { 0 }, src[4];
+
+            for (p = 0; p < s->nb_components; p++) {
+                const uint16_t *src_ptr = (const uint16_t *)in->data[p];
+                const int src_linesize = in->linesize[p] / 2;
+                const int thr = s->thr[p];
+                const int w = s->planewidth[p] - 1;
+                const int h = s->planeheight[p] - 1;
+                const int ref0 = src_ptr[av_clip(y +  y_pos, 0, h) * src_linesize + av_clip(x +  x_pos, 0, w)];
+                const int ref1 = src_ptr[av_clip(y + -y_pos, 0, h) * src_linesize + av_clip(x +  x_pos, 0, w)];
+                const int ref2 = src_ptr[av_clip(y + -y_pos, 0, h) * src_linesize + av_clip(x + -x_pos, 0, w)];
+                const int ref3 = src_ptr[av_clip(y +  y_pos, 0, h) * src_linesize + av_clip(x + -x_pos, 0, w)];
+                const int src0 = src_ptr[y * src_linesize + x];
+
+                src[p] = src0;
+                avg[p] = get_avg(ref0, ref1, ref2, ref3);
+
+                if (s->blur) {
+                    cmp[p] = FFABS(src0 - avg[p]) < thr;
+                } else {
+                    cmp[p] = (FFABS(src0 - ref0) < thr) &&
+                             (FFABS(src0 - ref1) < thr) &&
+                             (FFABS(src0 - ref2) < thr) &&
+                             (FFABS(src0 - ref3) < thr);
+                }
+            }
+
+            for (z = 0; z < s->nb_components; z++)
+                if (!cmp[z])
+                    break;
+            if (z == s->nb_components) {
+                for (p = 0; p < s->nb_components; p++) {
+                    const int dst_linesize = out->linesize[p] / 2;
+                    uint16_t *dst = (uint16_t *)out->data[p] + y * dst_linesize + x;
+
+                    dst[0] = avg[p];
+                }
+            } else {
+                for (p = 0; p < s->nb_components; p++) {
+                    const int dst_linesize = out->linesize[p] / 2;
+                    uint16_t *dst = (uint16_t *)out->data[p] + y * dst_linesize + x;
+
+                    dst[0] = src[p];
+                }
+            }
+        }
+    }
+
+    return 0;
+}
+
 static int deband_16_c(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs)
 {
     DebandContext *s = ctx->priv;
@@ -222,8 +375,13 @@ static int config_input(AVFilterLink *inlink)
     s->planeheight[0] = s->planeheight[3] = inlink->h;
     s->planewidth[1]  = s->planewidth[2]  = AV_CEIL_RSHIFT(inlink->w, desc->log2_chroma_w);
     s->planewidth[0]  = s->planewidth[3]  = inlink->w;
+    s->shift[0] = desc->log2_chroma_w;
+    s->shift[1] = desc->log2_chroma_h;
 
-    s->deband = desc->comp[0].depth > 8 ? deband_16_c : deband_8_c;
+    if (s->coupling)
+        s->deband = desc->comp[0].depth > 8 ? deband_16_coupling_c : deband_8_coupling_c;
+    else
+        s->deband = desc->comp[0].depth > 8 ? deband_16_c : deband_8_c;
 
     s->thr[0] = ((1 << desc->comp[0].depth) - 1) * s->threshold[0];
     s->thr[1] = ((1 << desc->comp[1].depth) - 1) * s->threshold[1];