avfilter/vf_colorlevels: add preserve color option

2 files changed, 147 insertions, 11 deletions

diff --git a/libavfilter/preserve_color.h b/libavfilter/preserve_color.h
new file mode 100644
index 0000000000..ac0587ad1e
--- /dev/null
+++ b/libavfilter/preserve_color.h
@@ -0,0 +1,83 @@
+/*
+ * This file is part of FFmpeg.
+ *
+ * FFmpeg is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public
+ * License as published by the Free Software Foundation; either
+ * version 2.1 of the License, or (at your option) any later version.
+ *
+ * FFmpeg is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public
+ * License along with FFmpeg; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
+ */
+
+#ifndef AVFILTER_PRESERVE_COLOR_H
+#define AVFILTER_PRESERVE_COLOR_H
+
+enum {
+    P_NONE,
+    P_LUM,
+    P_MAX,
+    P_AVG,
+    P_SUM,
+    P_NRM,
+    P_PWR,
+    NB_PRESERVE
+};
+
+static inline float normalize(float r, float g, float b, float max)
+{
+    r /= max;
+    g /= max;
+    b /= max;
+    return sqrtf(r * r + g * g + b * b);
+}
+
+static inline float power(float r, float g, float b, float max)
+{
+    r /= max;
+    g /= max;
+    b /= max;
+    return cbrtf(r * r * r + g * g * g + b * b * b);
+}
+
+static inline void preserve_color(int preserve_color,
+                                  float ir, float ig, float ib,
+                                  float  r, float  g, float  b,
+                                  float max,
+                                  float *icolor, float *ocolor)
+{
+    switch (preserve_color) {
+    case P_LUM:
+        *icolor = FFMAX3(ir, ig, ib) + FFMIN3(ir, ig, ib);
+        *ocolor = FFMAX3( r,  g,  b) + FFMIN3( r,  g,  b);
+        break;
+    case P_MAX:
+        *icolor = FFMAX3(ir, ig, ib);
+        *ocolor = FFMAX3( r,  g,  b);
+        break;
+    case P_AVG:
+        *icolor = (ir + ig + ib + 1.f) / 3.f;
+        *ocolor = ( r +  g +  b + 1.f) / 3.f;
+        break;
+    case P_SUM:
+        *icolor = ir + ig + ib;
+        *ocolor =  r +  g  + b;
+        break;
+    case P_NRM:
+        *icolor = normalize(ir, ig, ib, max);
+        *ocolor = normalize( r,  g,  b, max);
+        break;
+    case P_PWR:
+        *icolor = power(ir, ig, ib, max);
+        *ocolor = power( r,  g,  b, max);
+        break;
+    }
+}
+
+#endif /* AVFILTER_PRESERVE_COLOR_H */
diff --git a/libavfilter/vf_colorlevels.c b/libavfilter/vf_colorlevels.c
index 9fd49d4536..89021e95d0 100644
--- a/libavfilter/vf_colorlevels.c
+++ b/libavfilter/vf_colorlevels.c
@@ -26,6 +26,7 @@
 #include "formats.h"
 #include "internal.h"
 #include "video.h"
+#include "preserve_color.h"
 
 #define R 0
 #define G 1
@@ -40,6 +41,7 @@ typedef struct Range {
 typedef struct ColorLevelsContext {
     const AVClass *class;
     Range range[4];
+    int preserve_color;
 
     int nb_comp;
     int bpp;
@@ -47,7 +49,7 @@ typedef struct ColorLevelsContext {
     uint8_t rgba_map[4];
     int linesize;
 
-    int (*colorlevels_slice)(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs);
+    int (*colorlevels_slice[2])(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs);
 } ColorLevelsContext;
 
 #define OFFSET(x) offsetof(ColorLevelsContext, x)
@@ -69,6 +71,14 @@ static const AVOption colorlevels_options[] = {
     { "gomax", "set output green white point", OFFSET(range[G].out_max), AV_OPT_TYPE_DOUBLE, {.dbl=1},  0, 1, FLAGS },
     { "bomax", "set output blue white point",  OFFSET(range[B].out_max), AV_OPT_TYPE_DOUBLE, {.dbl=1},  0, 1, FLAGS },
     { "aomax", "set output alpha white point", OFFSET(range[A].out_max), AV_OPT_TYPE_DOUBLE, {.dbl=1},  0, 1, FLAGS },
+    { "preserve", "set 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" },
     { NULL }
 };
 
@@ -104,7 +114,7 @@ typedef struct ThreadData {
     int omin[4];
 } ThreadData;
 
-#define DO_COMMON(type, clip)                                                   \
+#define DO_COMMON(type, clip, preserve)                                         \
     ColorLevelsContext *s = ctx->priv;                                          \
     const ThreadData *td = arg;                                                 \
     const int linesize = s->linesize;                                           \
@@ -143,9 +153,35 @@ typedef struct ThreadData {
                                                                                 \
     for (int y = slice_start; y < slice_end; y++) {                             \
         for (int x = 0; x < linesize; x += step) {                              \
-            dst_r[x] = clip((src_r[x] - imin_r) * coeff_r + omin_r);            \
-            dst_g[x] = clip((src_g[x] - imin_g) * coeff_g + omin_g);            \
-            dst_b[x] = clip((src_b[x] - imin_b) * coeff_b + omin_b);            \
+            int ir, ig, ib, or, og, ob;                                         \
+            ir = src_r[x];                                                      \
+            ig = src_g[x];                                                      \
+            ib = src_b[x];                                                      \
+            if (preserve) {                                                     \
+                float ratio, icolor, ocolor, max = (1<<(8*sizeof(type)))-1;     \
+                                                                                \
+                or = (ir - imin_r) * coeff_r + omin_r;                          \
+                og = (ig - imin_g) * coeff_g + omin_g;                          \
+                ob = (ib - imin_b) * coeff_b + omin_b;                          \
+                                                                                \
+                preserve_color(s->preserve_color, ir, ig, ib, or, og, ob, max,  \
+                              &icolor, &ocolor);                                \
+                if (ocolor > 0.f) {                                             \
+                    ratio = icolor / ocolor;                                    \
+                                                                                \
+                    or *= ratio;                                                \
+                    og *= ratio;                                                \
+                    ob *= ratio;                                                \
+                }                                                               \
+                                                                                \
+                dst_r[x] = clip(or);                                            \
+                dst_g[x] = clip(og);                                            \
+                dst_b[x] = clip(ob);                                            \
+            } else {                                                            \
+                dst_r[x] = clip((ir - imin_r) * coeff_r + omin_r);              \
+                dst_g[x] = clip((ig - imin_g) * coeff_g + omin_g);              \
+                dst_b[x] = clip((ib - imin_b) * coeff_b + omin_b);              \
+            }                                                                   \
         }                                                                       \
                                                                                 \
         for (int x = 0; x < linesize && s->nb_comp == 4; x += step)             \
@@ -164,14 +200,28 @@ typedef struct ThreadData {
 
 static int colorlevels_slice_8(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs)
 {
-    DO_COMMON(uint8_t, av_clip_uint8)
+    DO_COMMON(uint8_t, av_clip_uint8, 0)
 
     return 0;
 }
 
 static int colorlevels_slice_16(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs)
 {
-    DO_COMMON(uint16_t, av_clip_uint16)
+    DO_COMMON(uint16_t, av_clip_uint16, 0)
+
+    return 0;
+}
+
+static int colorlevels_preserve_slice_8(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs)
+{
+    DO_COMMON(uint8_t, av_clip_uint8, 1)
+
+    return 0;
+}
+
+static int colorlevels_preserve_slice_16(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs)
+{
+    DO_COMMON(uint16_t, av_clip_uint16, 1)
 
     return 0;
 }
@@ -188,9 +238,12 @@ static int config_input(AVFilterLink *inlink)
     s->linesize = inlink->w * s->step;
     ff_fill_rgba_map(s->rgba_map, inlink->format);
 
-    s->colorlevels_slice = colorlevels_slice_8;
-    if (s->bpp == 2)
-        s->colorlevels_slice = colorlevels_slice_16;
+    s->colorlevels_slice[0] = colorlevels_slice_8;
+    s->colorlevels_slice[1] = colorlevels_preserve_slice_8;
+    if (s->bpp == 2) {
+        s->colorlevels_slice[0] = colorlevels_slice_16;
+        s->colorlevels_slice[1] = colorlevels_preserve_slice_16;
+    }
 
     return 0;
 }
@@ -304,7 +357,7 @@ static int filter_frame(AVFilterLink *inlink, AVFrame *in)
         break;
     }
 
-    ff_filter_execute(ctx, s->colorlevels_slice, &td, NULL,
+    ff_filter_execute(ctx, s->colorlevels_slice[s->preserve_color > 0], &td, NULL,
                       FFMIN(inlink->h, ff_filter_get_nb_threads(ctx)));
 
     if (in != out)