avfilter/vf_pseudocolor: allow more advanced presets

1 files changed, 137 insertions, 44 deletions

diff --git a/libavfilter/vf_pseudocolor.c b/libavfilter/vf_pseudocolor.c
index 21c054adff..192839342b 100644
--- a/libavfilter/vf_pseudocolor.c
+++ b/libavfilter/vf_pseudocolor.c
@@ -59,48 +59,98 @@ enum var_name {
     VAR_VARS_NB
 };
 
-enum ColorMaps {
+enum Curves {
     MAGMA,
     INFERNO,
     PLASMA,
     VIRIDIS,
     TURBO,
     CIVIDIS,
-    NB_CMAPS,
+    NB_CURVES,
 };
 
-static const double presets[][3][8] =
+enum Presets {
+    PRESET_MAGMA,
+    PRESET_INFERNO,
+    PRESET_PLASMA,
+    PRESET_VIRIDIS,
+    PRESET_TURBO,
+    PRESET_CIVIDIS,
+    PRESET_RANGE1,
+    PRESET_RANGE2,
+    NB_PRESETS,
+};
+
+typedef struct Curve {
+    double coef[3][8];
+} Curve;
+
+typedef struct Fill {
+    float fill[3];
+} Fill;
+
+typedef struct Range {
+    int start, end;
+} Range;
+
+typedef struct Preset {
+    int nb_segments;
+    const Range *ranges;
+    const Curve *curves;
+    const Fill  *fills;
+} Preset;
+
+static const Range full_range   = {0, 256};
+static const Range spec1_range[] = {{0, 16}, {16, 236}, {236, 256}};
+static const Range spec2_range[] = {{0, 16}, {16, 22}, {22, 226}, {226, 236}, {236, 256}};
+
+static const Fill spec1_fills[] = {{0.5f, 0.f, .5f}, {-1.f, -1.f, -1.f}, {1.f, 0.f, 0.f}};
+static const Fill spec2_fills[] = {{0.5f, 0.f, .5f}, {0.f, 1.f, 1.f}, {-1.f, -1.f, -1.f}, {1.f, 1.f, 0.f}, {1.f, 0.f, 0.f}};
+
+static const Curve curves[] =
 {
-    [MAGMA] = {
+    [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] = {
+    }},
+    [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] = {
+    }},
+    [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] = {
+    }},
+    [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] = {
+    }},
+    [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] = {
+    }},
+    [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 },
-    },
+    }},
+};
+
+static const Preset presets[] =
+{
+    [PRESET_MAGMA]   = { 1, &full_range, &curves[MAGMA],   NULL },
+    [PRESET_INFERNO] = { 1, &full_range, &curves[INFERNO], NULL },
+    [PRESET_PLASMA]  = { 1, &full_range, &curves[PLASMA],  NULL },
+    [PRESET_VIRIDIS] = { 1, &full_range, &curves[VIRIDIS], NULL },
+    [PRESET_TURBO]   = { 1, &full_range, &curves[TURBO],   NULL },
+    [PRESET_CIVIDIS] = { 1, &full_range, &curves[CIVIDIS], NULL },
+    [PRESET_RANGE1]  = { 3, spec1_range, NULL,             spec1_fills },
+    [PRESET_RANGE2]  = { 5, spec2_range, NULL,             spec2_fills },
 };
 
 typedef struct PseudoColorContext {
@@ -135,14 +185,16 @@ 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" },
+    { "p",  "set preset",                  OFFSET(preset),           AV_OPT_TYPE_INT,    {.i64=-1},-1, NB_PRESETS-1, .flags = FLAGS, "preset" },
+    { "none",       NULL,                  0,                        AV_OPT_TYPE_CONST,  {.i64=-1},             .flags = FLAGS, "preset" },
+    { "magma",      NULL,                  0,                        AV_OPT_TYPE_CONST,  {.i64=PRESET_MAGMA},   .flags = FLAGS, "preset" },
+    { "inferno",    NULL,                  0,                        AV_OPT_TYPE_CONST,  {.i64=PRESET_INFERNO}, .flags = FLAGS, "preset" },
+    { "plasma",     NULL,                  0,                        AV_OPT_TYPE_CONST,  {.i64=PRESET_PLASMA},  .flags = FLAGS, "preset" },
+    { "viridis",    NULL,                  0,                        AV_OPT_TYPE_CONST,  {.i64=PRESET_VIRIDIS}, .flags = FLAGS, "preset" },
+    { "turbo",      NULL,                  0,                        AV_OPT_TYPE_CONST,  {.i64=PRESET_TURBO},   .flags = FLAGS, "preset" },
+    { "cividis",    NULL,                  0,                        AV_OPT_TYPE_CONST,  {.i64=PRESET_CIVIDIS}, .flags = FLAGS, "preset" },
+    { "range1",     NULL,                  0,                        AV_OPT_TYPE_CONST,  {.i64=PRESET_RANGE1},  .flags = FLAGS, "preset" },
+    { "range2",     NULL,                  0,                        AV_OPT_TYPE_CONST,  {.i64=PRESET_RANGE2},  .flags = FLAGS, "preset" },
     { NULL }
 };
 
@@ -534,28 +586,69 @@ 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;
+        int nb_segments = presets[s->preset].nb_segments;
+
+        for (int seg = 0; seg < nb_segments; seg++) {
+            int start = presets[s->preset].ranges[seg].start;
+            int end   = presets[s->preset].ranges[seg].end;
+
+            for (int i = start; i < end; i++) {
+                if (!presets[s->preset].curves) {
+                    const Fill fill = presets[s->preset].fills[seg];
+
+                    for (int j = 0; j < factor; j++) {
+                        double r, g, b;
+
+                        g = fill.fill[1];
+                        b = fill.fill[2];
+                        r = fill.fill[0];
+
+                        if (g >= 0.f && b >= 0.f && r >= 0.f) {
+                            g *= s->max;
+                            b *= s->max;
+                            r *= 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;
+                    }
+                } else {
+                    const Curve curve = presets[s->preset].curves[seg];
+
+                    for (int j = 0; j < factor; j++) {
+                        const double lf = j / (double)factor;
+                        double r, g, b;
+
+                        g = poly_eval(curve.coef[1], i + lf) * s->max;
+                        b = poly_eval(curve.coef[2], i + lf) * s->max;
+                        r = poly_eval(curve.coef[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;
+                    }
                 }
-
-                s->lut[0][i*factor+j] = g;
-                s->lut[1][i*factor+j] = b;
-                s->lut[2][i*factor+j] = r;
             }
         }
     }