avfilter/vf_colorchannelmixer: add extended preserve color support

2 files changed, 94 insertions, 84 deletions

diff --git a/doc/filters.texi b/doc/filters.texi
index b309e8c737..398c6b82fc 100644
--- a/doc/filters.texi
+++ b/doc/filters.texi
@@ -8379,8 +8379,27 @@ Default is @code{1} for @var{aa}, and @code{0} for @var{ar}, @var{ag} and @var{a
 
 Allowed ranges for options are @code{[-2.0, 2.0]}.
 
-@item pl
-Preserve lightness when changing colors. Allowed range is from @code{[0.0, 1.0]}.
+@item pc
+Set preserve color mode. The accepted values are:
+@table @samp
+@item none
+Disable color preserving, this is default.
+@item lum
+Preserve luminance.
+@item max
+Preserve max value of RGB triplet.
+@item avg
+Preserve average value of RGB triplet.
+@item sum
+Preserve sum value of RGB triplet.
+@item nrm
+Preserve normalized value of RGB triplet.
+@item pwr
+Preserve power value of RGB triplet.
+@end table
+
+@item pa
+Set the preserve color amount when changing colors. Allowed range is from @code{[0.0, 1.0]}.
 Default is @code{0.0}, thus disabled.
 @end table
 
diff --git a/libavfilter/vf_colorchannelmixer.c b/libavfilter/vf_colorchannelmixer.c
index e51f9d2d68..f1a193116e 100644
--- a/libavfilter/vf_colorchannelmixer.c
+++ b/libavfilter/vf_colorchannelmixer.c
@@ -27,6 +27,7 @@
 #include "formats.h"
 #include "internal.h"
 #include "video.h"
+#include "preserve_color.h"
 
 #define R 0
 #define G 1
@@ -43,8 +44,8 @@ typedef struct ColorChannelMixerContext {
     double gr, gg, gb, ga;
     double br, bg, bb, ba;
     double ar, ag, ab, aa;
-    double sr, sg, sb;
-    double preserve_lightness;
+    double preserve_amount;
+    int preserve_color;
 
     int *lut[4][4];
 
@@ -75,7 +76,15 @@ static const AVOption colorchannelmixer_options[] = {
     { "ag", "set the green gain for the alpha channel", OFFSET(ag), AV_OPT_TYPE_DOUBLE, {.dbl=0}, -2, 2, FLAGS },
     { "ab", "set the blue gain for the alpha channel",  OFFSET(ab), AV_OPT_TYPE_DOUBLE, {.dbl=0}, -2, 2, FLAGS },
     { "aa", "set the alpha gain for the alpha channel", OFFSET(aa), AV_OPT_TYPE_DOUBLE, {.dbl=1}, -2, 2, FLAGS },
-    { "pl", "preserve lightness",       OFFSET(preserve_lightness), AV_OPT_TYPE_DOUBLE, {.dbl=0},  0, 1, FLAGS },
+    { "pc", "set the preserve color mode",  OFFSET(preserve_color), AV_OPT_TYPE_INT,    {.i64=0},  0, NB_PRESERVE-1, FLAGS, "preserve" },
+    { "none",  "disabled",                     0,                        AV_OPT_TYPE_CONST,  {.i64=P_NONE}, 0, 0, FLAGS, "preserve" },
+    { "lum",   "luminance",                    0,                        AV_OPT_TYPE_CONST,  {.i64=P_LUM},  0, 0, FLAGS, "preserve" },
+    { "max",   "max",                          0,                        AV_OPT_TYPE_CONST,  {.i64=P_MAX},  0, 0, FLAGS, "preserve" },
+    { "avg",   "average",                      0,                        AV_OPT_TYPE_CONST,  {.i64=P_AVG},  0, 0, FLAGS, "preserve" },
+    { "sum",   "sum",                          0,                        AV_OPT_TYPE_CONST,  {.i64=P_SUM},  0, 0, FLAGS, "preserve" },
+    { "nrm",   "norm",                         0,                        AV_OPT_TYPE_CONST,  {.i64=P_NRM},  0, 0, FLAGS, "preserve" },
+    { "pwr",   "power",                        0,                        AV_OPT_TYPE_CONST,  {.i64=P_PWR},  0, 0, FLAGS, "preserve" },
+    { "pa", "set the preserve color amount",    OFFSET(preserve_amount), AV_OPT_TYPE_DOUBLE, {.dbl=0},  0, 1, FLAGS },
     { NULL }
 };
 
@@ -108,24 +117,23 @@ static float lerpf(float v0, float v1, float f)
     return v0 + (v1 - v0) * f;
 }
 
-static void preservel(float *r, float *g, float *b, float lin, float lout)
+static void preservel(float *r, float *g, float *b, float lin, float lout, float max)
 {
-    *r *= lout / lin;
-    *g *= lout / lin;
-    *b *= lout / lin;
+    if (lout <= 0.f)
+        lout = 1.f / (max * 2.f);
+    *r *= lin / lout;
+    *g *= lin / lout;
+    *b *= lin / lout;
 }
 
 static av_always_inline int filter_slice_rgba_planar(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs,
-                                                     int have_alpha, int pl)
+                                                     int have_alpha, int pc)
 {
     ColorChannelMixerContext *s = ctx->priv;
     ThreadData *td = arg;
     AVFrame *in = td->in;
     AVFrame *out = td->out;
-    const float l = s->preserve_lightness;
-    const float sr = s->sr;
-    const float sg = s->sg;
-    const float sb = s->sb;
+    const float pa = s->preserve_amount;
     const int slice_start = (out->height * jobnr) / nb_jobs;
     const int slice_end = (out->height * (jobnr+1)) / nb_jobs;
     const uint8_t *srcg = in->data[0] + slice_start * in->linesize[0];
@@ -159,18 +167,19 @@ static av_always_inline int filter_slice_rgba_planar(AVFilterContext *ctx, void
                    s->lut[B][B][bin] +
                    (have_alpha == 1 ? s->lut[B][A][ain] : 0);
 
-            if (pl) {
-                float lin = FFMAX3(rin, gin, bin) + FFMIN3(rin, gin, bin);
-                float frout = rout / sr;
-                float fgout = gout / sg;
-                float fbout = bout / sb;
-                float lout = FFMAX3(frout, fgout, fbout) + FFMIN3(frout, fgout, fbout);
+            if (pc) {
+                float frout = av_clipf(rout, 0.f, 255.f);
+                float fgout = av_clipf(gout, 0.f, 255.f);
+                float fbout = av_clipf(bout, 0.f, 255.f);
+                float lin, lout;
 
-                preservel(&frout, &fgout, &fbout, lin, lout);
+                preserve_color(s->preserve_color, rin, gin, bin,
+                               rout, gout, bout, 255.f, &lin, &lout);
+                preservel(&frout, &fgout, &fbout, lin, lout, 255.f);
 
-                rout = lrintf(lerpf(rout, frout, l));
-                gout = lrintf(lerpf(gout, fgout, l));
-                bout = lrintf(lerpf(bout, fbout, l));
+                rout = lrintf(lerpf(rout, frout, pa));
+                gout = lrintf(lerpf(gout, fgout, pa));
+                bout = lrintf(lerpf(bout, fbout, pa));
             }
 
             dstr[j] = av_clip_uint8(rout);
@@ -199,16 +208,14 @@ static av_always_inline int filter_slice_rgba_planar(AVFilterContext *ctx, void
 }
 
 static av_always_inline int filter_slice_rgba16_planar(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs,
-                                                       int have_alpha, int depth, int pl)
+                                                       int have_alpha, int depth, int pc)
 {
     ColorChannelMixerContext *s = ctx->priv;
     ThreadData *td = arg;
     AVFrame *in = td->in;
     AVFrame *out = td->out;
-    const float l = s->preserve_lightness;
-    const float sr = s->sr;
-    const float sg = s->sg;
-    const float sb = s->sb;
+    const float pa = s->preserve_amount;
+    const float max = (1 << depth) - 1;
     const int slice_start = (out->height * jobnr) / nb_jobs;
     const int slice_end = (out->height * (jobnr+1)) / nb_jobs;
     const uint16_t *srcg = (const uint16_t *)(in->data[0] + slice_start * in->linesize[0]);
@@ -242,18 +249,19 @@ static av_always_inline int filter_slice_rgba16_planar(AVFilterContext *ctx, voi
                    s->lut[B][B][bin] +
                    (have_alpha == 1 ? s->lut[B][A][ain] : 0);
 
-            if (pl) {
-                float lin = FFMAX3(rin, gin, bin) + FFMIN3(rin, gin, bin);
-                float frout = rout / sr;
-                float fgout = gout / sg;
-                float fbout = bout / sb;
-                float lout = FFMAX3(frout, fgout, fbout) + FFMIN3(frout, fgout, fbout);
+            if (pc) {
+                float frout = av_clipf(rout, 0.f, max);
+                float fgout = av_clipf(gout, 0.f, max);
+                float fbout = av_clipf(bout, 0.f, max);
+                float lin, lout;
 
-                preservel(&frout, &fgout, &fbout, lin, lout);
+                preserve_color(s->preserve_color, rin, gin, bin,
+                               rout, gout, bout, max, &lin, &lout);
+                preservel(&frout, &fgout, &fbout, lin, lout, max);
 
-                rout = lrintf(lerpf(rout, frout, l));
-                gout = lrintf(lerpf(gout, fgout, l));
-                bout = lrintf(lerpf(bout, fbout, l));
+                rout = lrintf(lerpf(rout, frout, pa));
+                gout = lrintf(lerpf(gout, fgout, pa));
+                bout = lrintf(lerpf(bout, fbout, pa));
             }
 
             dstr[j] = av_clip_uintp2(rout, depth);
@@ -382,16 +390,13 @@ static int filter_slice_gbrap16_pl(AVFilterContext *ctx, void *arg, int jobnr, i
 }
 
 static av_always_inline int filter_slice_rgba_packed(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs,
-                                                     int have_alpha, int step, int pl)
+                                                     int have_alpha, int step, int pc)
 {
     ColorChannelMixerContext *s = ctx->priv;
     ThreadData *td = arg;
     AVFrame *in = td->in;
     AVFrame *out = td->out;
-    const float l = s->preserve_lightness;
-    const float sr = s->sr;
-    const float sg = s->sg;
-    const float sb = s->sb;
+    const float pa = s->preserve_amount;
     const int slice_start = (out->height * jobnr) / nb_jobs;
     const int slice_end = (out->height * (jobnr+1)) / nb_jobs;
     const uint8_t roffset = s->rgba_map[R];
@@ -426,18 +431,19 @@ static av_always_inline int filter_slice_rgba_packed(AVFilterContext *ctx, void
                    s->lut[B][B][bin] +
                    (have_alpha == 1 ? s->lut[B][A][ain] : 0);
 
-            if (pl) {
-                float lin = FFMAX3(rin, gin, bin) + FFMIN3(rin, gin, bin);
-                float frout = rout / sr;
-                float fgout = gout / sg;
-                float fbout = bout / sb;
-                float lout = FFMAX3(frout, fgout, fbout) + FFMIN3(frout, fgout, fbout);
+            if (pc) {
+                float frout = av_clipf(rout, 0.f, 255.f);
+                float fgout = av_clipf(gout, 0.f, 255.f);
+                float fbout = av_clipf(bout, 0.f, 255.f);
+                float lin, lout;
 
-                preservel(&frout, &fgout, &fbout, lin, lout);
+                preserve_color(s->preserve_color, rin, gin, bin,
+                               rout, gout, bout, 255.f, &lin, &lout);
+                preservel(&frout, &fgout, &fbout, lin, lout, 255.f);
 
-                rout = lrintf(lerpf(rout, frout, l));
-                gout = lrintf(lerpf(gout, fgout, l));
-                bout = lrintf(lerpf(bout, fbout, l));
+                rout = lrintf(lerpf(rout, frout, pa));
+                gout = lrintf(lerpf(gout, fgout, pa));
+                bout = lrintf(lerpf(bout, fbout, pa));
             }
 
             dst[j + roffset] = av_clip_uint8(rout);
@@ -461,16 +467,13 @@ static av_always_inline int filter_slice_rgba_packed(AVFilterContext *ctx, void
 }
 
 static av_always_inline int filter_slice_rgba16_packed(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs,
-                                                       int have_alpha, int step, int pl)
+                                                       int have_alpha, int step, int pc)
 {
     ColorChannelMixerContext *s = ctx->priv;
     ThreadData *td = arg;
     AVFrame *in = td->in;
     AVFrame *out = td->out;
-    const float l = s->preserve_lightness;
-    const float sr = s->sr;
-    const float sg = s->sg;
-    const float sb = s->sb;
+    const float pa = s->preserve_amount;
     const int slice_start = (out->height * jobnr) / nb_jobs;
     const int slice_end = (out->height * (jobnr+1)) / nb_jobs;
     const uint8_t roffset = s->rgba_map[R];
@@ -505,18 +508,19 @@ static av_always_inline int filter_slice_rgba16_packed(AVFilterContext *ctx, voi
                    s->lut[B][B][bin] +
                    (have_alpha == 1 ? s->lut[B][A][ain] : 0);
 
-            if (pl) {
-                float lin = FFMAX3(rin, gin, bin) + FFMIN3(rin, gin, bin);
-                float frout = rout / sr;
-                float fgout = gout / sg;
-                float fbout = bout / sb;
-                float lout = FFMAX3(frout, fgout, fbout) + FFMIN3(frout, fgout, fbout);
+            if (pc) {
+                float frout = av_clipf(rout, 0.f, 65535.f);
+                float fgout = av_clipf(gout, 0.f, 65535.f);
+                float fbout = av_clipf(bout, 0.f, 65535.f);
+                float lin, lout;
 
-                preservel(&frout, &fgout, &fbout, lin, lout);
+                preserve_color(s->preserve_color, rin, gin, bin,
+                               rout, gout, bout, 65535.f, &lin, &lout);
+                preservel(&frout, &fgout, &fbout, lin, lout, 65535.f);
 
-                rout = lrintf(lerpf(rout, frout, l));
-                gout = lrintf(lerpf(gout, fgout, l));
-                bout = lrintf(lerpf(bout, fbout, l));
+                rout = lrintf(lerpf(rout, frout, pa));
+                gout = lrintf(lerpf(gout, fgout, pa));
+                bout = lrintf(lerpf(bout, fbout, pa));
             }
 
             dst[j + roffset] = av_clip_uint16(rout);
@@ -609,19 +613,6 @@ static int config_output(AVFilterLink *outlink)
                 s->lut[i][j] = buffer;
     }
 
-    s->sr = s->rr + s->rg + s->rb + s->ra;
-    s->sg = s->gr + s->gg + s->gb + s->ga;
-    s->sb = s->br + s->bg + s->bb + s->ba;
-
-    if (fabs(s->sr) <= DBL_EPSILON)
-        s->sr = 1.;
-
-    if (fabs(s->sg) <= DBL_EPSILON)
-        s->sg = 1.;
-
-    if (fabs(s->sb) <= DBL_EPSILON)
-        s->sb = 1.;
-
     for (i = 0; i < size; i++) {
         s->lut[R][R][i] = lrint(i * s->rr);
         s->lut[R][G][i] = lrint(i * s->rg);
@@ -724,7 +715,7 @@ static int filter_frame(AVFilterLink *inlink, AVFrame *in)
     AVFilterContext *ctx = inlink->dst;
     ColorChannelMixerContext *s = ctx->priv;
     AVFilterLink *outlink = ctx->outputs[0];
-    const int pl = s->preserve_lightness > 0.;
+    const int pc = s->preserve_color > 0;
     ThreadData td;
     AVFrame *out;
 
@@ -741,7 +732,7 @@ static int filter_frame(AVFilterLink *inlink, AVFrame *in)
 
     td.in = in;
     td.out = out;
-    ff_filter_execute(ctx, s->filter_slice[pl], &td, NULL,
+    ff_filter_execute(ctx, s->filter_slice[pc], &td, NULL,
                       FFMIN(outlink->h, ff_filter_get_nb_threads(ctx)));
 
     if (in != out)