avfilter/vf_pseudocolor: add presets

1 files changed, 107 insertions, 2 deletions

diff --git a/libavfilter/vf_pseudocolor.c b/libavfilter/vf_pseudocolor.c
index 4c24f49bd9..2cba46e3da 100644
--- a/libavfilter/vf_pseudocolor.c
+++ b/libavfilter/vf_pseudocolor.c
@@ -59,8 +59,53 @@ enum var_name {
     VAR_VARS_NB
 };
 
+enum ColorMaps {
+    MAGMA,
+    INFERNO,
+    PLASMA,
+    VIRIDIS,
+    TURBO,
+    CIVIDIS,
+    NB_CMAPS,
+};
+
+static const double presets[][3][8] =
+{
+    [MAGMA] = {
+        {-7.5631093e-16,  7.4289183e-13, -2.8525484e-10,  5.4446085e-08, -5.5596238e-06,  3.0569325e-04, -2.3137421e-03,  1.2152095e-02 },
+        { 1.3217636e-15, -1.2214648e-12,  4.4319712e-10, -8.0197993e-08,  7.6598370e-06, -3.6523704e-04,  8.4836670e-03, -2.5536888e-02 },
+        {-1.1446568e-15,  1.0013446e-12, -3.5651575e-10,  6.6775016e-08, -6.7120346e-06,  2.7346619e-04,  4.7969657e-03,  1.1971441e-02 },
+    },
+    [INFERNO] = {
+        {-3.9848859e-18,  9.4821649e-14, -6.7371977e-11,  1.8469937e-08, -2.5359307e-06,  1.7959053e-04,  3.9782564e-04,  2.8845935e-04 },
+        { 6.8408539e-16, -6.5499979e-13,  2.4562526e-10, -4.5989298e-08,  4.5723324e-06, -2.2111913e-04,  5.2023164e-03, -1.1226064e-02 },
+        {-2.9921470e-15,  2.5864165e-12, -8.7403799e-10,  1.4713388e-07, -1.2701505e-05,  4.5159935e-04,  3.1087989e-03,  1.9122831e-02 },
+    },
+    [PLASMA] = {
+        { 3.6196089e-16, -3.3623041e-13,  1.2324010e-10, -2.2769060e-08,  2.2297792e-06, -1.2567829e-04,  9.9791629e-03,  5.7247918e-02 },
+        { 5.0262888e-16, -5.3193896e-13,  2.2451715e-10, -4.7529623e-08,  5.1374873e-06, -2.3260136e-04,  3.1502825e-03,  1.5362491e-02 },
+        {-1.7782261e-16,  2.2487839e-13, -1.0610236e-10,  2.4112644e-08, -2.6331623e-06,  8.9499751e-05,  2.1386328e-03,  5.3824268e-01 },
+    },
+    [VIRIDIS] = {
+        { 9.4850045e-16, -8.6629383e-13,  3.0310944e-10, -5.1340396e-08,  4.6024275e-06, -2.2744239e-04,  4.5559993e-03,  2.5662350e-01 },
+        { 9.6461041e-17, -6.9209477e-14,  1.7625397e-11, -2.0229773e-09,  1.4900110e-07, -1.9315187e-05,  5.8967339e-03,  3.9544827e-03 },
+        { 5.1785449e-16, -3.6663004e-13,  1.0249990e-10, -1.5431998e-08,  1.5007941e-06, -1.2001502e-04,  7.6951526e-03,  3.2292815e-01 },
+    },
+    [TURBO] = {
+        {-4.3683890e-15,  3.7020347e-12, -1.1712592e-09,  1.6401790e-07, -8.6842919e-06, -1.8542465e-06,  8.4485325e-03,  1.6267077e-01 },
+        {-4.0011069e-16,  2.7861423e-13, -6.3388921e-11,  5.8872238e-09, -5.4466522e-07,  1.8037114e-05,  1.0599869e-02,  7.6914696e-02 },
+        {-2.8242609e-15,  2.9234108e-12, -1.1726546e-09,  2.2552115e-07, -2.0059387e-05,  5.0595552e-04,  1.7714932e-02,  2.7271836e-01 },
+    },
+    [CIVIDIS] = {
+        {-9.5484131e-16,  9.6988184e-13, -4.0058766e-10,  8.5743924e-08, -9.9644797e-06,  5.9197908e-04, -1.0361579e-02,  3.3164429e-02 },
+        { 1.2731941e-17, -9.4238449e-15,  2.2808841e-12, -1.1548296e-10, -2.3888913e-08,  3.8986680e-06,  2.5879330e-03,  1.2769733e-01 },
+        { 4.6004608e-16, -5.0686849e-13,  2.2753449e-10, -5.3074099e-08,  6.7196096e-06, -4.4120020e-04,  1.3435551e-02,  2.8293355e-01 },
+    },
+};
+
 typedef struct PseudoColorContext {
     const AVClass *class;
+    int preset;
     int max;
     int index;
     int nb_planes;
@@ -90,6 +135,14 @@ static const AVOption pseudocolor_options[] = {
     { "c2", "set component #2 expression", OFFSET(comp_expr_str[2]), AV_OPT_TYPE_STRING, {.str="val"},   .flags = FLAGS },
     { "c3", "set component #3 expression", OFFSET(comp_expr_str[3]), AV_OPT_TYPE_STRING, {.str="val"},   .flags = FLAGS },
     { "i",  "set component as base",       OFFSET(index),            AV_OPT_TYPE_INT,    {.i64=0}, 0, 3, .flags = FLAGS },
+    { "p",  "set preset",                  OFFSET(preset),           AV_OPT_TYPE_INT,    {.i64=-1},-1,NB_CMAPS-1, .flags = FLAGS, "preset" },
+    { "none",       NULL,                  0,                        AV_OPT_TYPE_CONST,  {.i64=-1},.flags = FLAGS, "preset" },
+    { "magma",      NULL,                  0,                        AV_OPT_TYPE_CONST,  {.i64=MAGMA},   .flags = FLAGS, "preset" },
+    { "inferno",    NULL,                  0,                        AV_OPT_TYPE_CONST,  {.i64=INFERNO}, .flags = FLAGS, "preset" },
+    { "plasma",     NULL,                  0,                        AV_OPT_TYPE_CONST,  {.i64=PLASMA},  .flags = FLAGS, "preset" },
+    { "viridis",    NULL,                  0,                        AV_OPT_TYPE_CONST,  {.i64=VIRIDIS}, .flags = FLAGS, "preset" },
+    { "turbo",      NULL,                  0,                        AV_OPT_TYPE_CONST,  {.i64=TURBO},   .flags = FLAGS, "preset" },
+    { "cividis",    NULL,                  0,                        AV_OPT_TYPE_CONST,  {.i64=CIVIDIS}, .flags = FLAGS, "preset" },
     { NULL }
 };
 
@@ -419,14 +472,39 @@ static void pseudocolor_filter_16_11d(int max, int width, int height,
     }
 }
 
+#define RGB_TO_Y_BT709(r, g, b) \
+((0.21260*219.0/255.0) * (r) + (0.71520*219.0/255.0) * (g) + \
+ (0.07220*219.0/255.0) * (b))
+
+#define RGB_TO_U_BT709(r1, g1, b1, max) \
+(-(0.11457*224.0/255.0) * r1 - (0.38543*224.0/255.0) * g1 + \
+    (0.50000*224.0/255.0) * b1 + max * 0.5)
+
+#define RGB_TO_V_BT709(r1, g1, b1, max) \
+((0.50000*224.0/255.0) * r1 - (0.45415*224.0/255.0) * g1 - \
+   (0.04585*224.0/255.0) * b1 + max * 0.5)
+
+static double poly_eval(const double *const poly, double x)
+{
+    double res = 0.;
+
+    for (int i = 0; i < 8; i++) {
+        res += pow(x, i) * poly[7-i];
+    }
+
+    return av_clipd(res, 0., 1.);
+}
+
 static int config_input(AVFilterLink *inlink)
 {
     AVFilterContext *ctx = inlink->dst;
     PseudoColorContext *s = ctx->priv;
     const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(inlink->format);
-    int depth, ret, hsub, vsub, color;
+    int depth, ret, hsub, vsub, color, factor, rgb;
 
+    rgb = desc->flags & AV_PIX_FMT_FLAG_RGB;
     depth = desc->comp[0].depth;
+    factor = 1 << (depth - 8);
     s->max = (1 << depth) - 1;
     s->nb_planes = av_pix_fmt_count_planes(inlink->format);
 
@@ -457,7 +535,7 @@ static int config_input(AVFilterLink *inlink)
     s->var_values[VAR_VMAX] = 240 * (1 << (depth - 8));
     s->var_values[VAR_AMAX] = s->max;
 
-    for (color = 0; color < s->nb_planes; color++) {
+    for (color = 0; color < s->nb_planes && s->preset < 0; color++) {
         double res;
         int val;
 
@@ -488,6 +566,33 @@ static int config_input(AVFilterLink *inlink)
         }
     }
 
+    if (s->preset >= 0) {
+        for (int i = 0; i < 256; i++) {
+            for (int j = 0; j < factor; j++) {
+                const double lf = j / (double)factor;
+                double r, g, b;
+
+                g = poly_eval(presets[s->preset][1], i + lf) * s->max;
+                b = poly_eval(presets[s->preset][2], i + lf) * s->max;
+                r = poly_eval(presets[s->preset][0], i + lf) * s->max;
+
+                if (!rgb) {
+                    double y = RGB_TO_Y_BT709(r, g, b);
+                    double u = RGB_TO_U_BT709(r, g, b, s->max);
+                    double v = RGB_TO_V_BT709(r, g, b, s->max);
+
+                    r = v;
+                    g = y;
+                    b = u;
+                }
+
+                s->lut[0][i*factor+j] = g;
+                s->lut[1][i*factor+j] = b;
+                s->lut[2][i*factor+j] = r;
+            }
+        }
+    }
+
     switch (inlink->format) {
     case AV_PIX_FMT_YUV444P:
     case AV_PIX_FMT_YUVA444P: