vo_gpu: improve tone mapping desaturation

Instead of desaturating towards luma, we desaturate towards the
per-channel tone mapped version. This essentially proves a smooth
roll-off towards the "hollywood"-style (non-chromatic) tone mapping
algorithm, which works better for bright content, while continuing to
use the "linear" style (chromatic) tone mapping algorithm for primarily
in-gamut content.

We also split up the desaturation algorithm into strength and exponent,
which allows users to use less aggressive desaturation settings without
affecting the overall curve.

4 files changed, 87 insertions, 76 deletions

diff --git a/video/out/gpu/video.c b/video/out/gpu/video.c
index c12fb8536c..9ffdc62d20 100644
--- a/video/out/gpu/video.c
+++ b/video/out/gpu/video.c
@@ -313,9 +313,12 @@ static const struct gl_video_opts gl_video_opts_def = {
     .alpha_mode = ALPHA_BLEND_TILES,
     .background = {0, 0, 0, 255},
     .gamma = 1.0f,
-    .tone_mapping = TONE_MAPPING_HABLE,
-    .tone_mapping_param = NAN,
-    .tone_mapping_desat = 0.5,
+    .tone_map = {
+        .curve = TONE_MAPPING_HABLE,
+        .curve_param = NAN,
+        .desat = 0.75,
+        .desat_exp = 1.5,
+    },
     .early_flush = -1,
     .hwdec_interop = "auto",
 };
@@ -353,20 +356,22 @@ const struct m_sub_options gl_video_conf = {
         OPT_CHOICE_C("target-trc", target_trc, 0, mp_csp_trc_names),
         OPT_CHOICE_OR_INT("target-peak", target_peak, 0, 10, 10000,
                           ({"auto", 0})),
-        OPT_CHOICE("tone-mapping", tone_mapping, 0,
+        OPT_CHOICE("tone-mapping", tone_map.curve, 0,
                    ({"clip",     TONE_MAPPING_CLIP},
                     {"mobius",   TONE_MAPPING_MOBIUS},
                     {"reinhard", TONE_MAPPING_REINHARD},
                     {"hable",    TONE_MAPPING_HABLE},
                     {"gamma",    TONE_MAPPING_GAMMA},
                     {"linear",   TONE_MAPPING_LINEAR})),
-        OPT_CHOICE("hdr-compute-peak", compute_hdr_peak, 0,
+        OPT_CHOICE("hdr-compute-peak", tone_map.compute_peak, 0,
                    ({"auto", 0},
                     {"yes", 1},
                     {"no", -1})),
-        OPT_FLOAT("tone-mapping-param", tone_mapping_param, 0),
-        OPT_FLOAT("tone-mapping-desaturate", tone_mapping_desat, 0),
-        OPT_FLAG("gamut-warning", gamut_warning, 0),
+        OPT_FLOAT("tone-mapping-param", tone_map.curve_param, 0),
+        OPT_FLOAT("tone-mapping-desaturate", tone_map.desat, 0),
+        OPT_FLOATRANGE("tone-mapping-desaturate-exponent",
+                       tone_map.desat_exp, 0, 0.0, 20.0),
+        OPT_FLAG("gamut-warning", tone_map.gamut_warning, 0),
         OPT_FLAG("opengl-pbo", pbo, 0),
         SCALER_OPTS("scale",  SCALER_SCALE),
         SCALER_OPTS("dscale", SCALER_DSCALE),
@@ -2472,7 +2477,8 @@ static void pass_colormanage(struct gl_video *p, struct mp_colorspace src, bool
     if (!dst.sig_peak)
         dst.sig_peak = mp_trc_nom_peak(dst.gamma);
 
-    bool detect_peak = p->opts.compute_hdr_peak >= 0 && mp_trc_is_hdr(src.gamma);
+    struct gl_tone_map_opts tone_map = p->opts.tone_map;
+    bool detect_peak = tone_map.compute_peak >= 0 && mp_trc_is_hdr(src.gamma);
     if (detect_peak && !p->hdr_peak_ssbo) {
         struct {
             uint32_t counter;
@@ -2493,8 +2499,8 @@ static void pass_colormanage(struct gl_video *p, struct mp_colorspace src, bool
         p->hdr_peak_ssbo = ra_buf_create(ra, &params);
         if (!p->hdr_peak_ssbo) {
             MP_WARN(p, "Failed to create HDR peak detection SSBO, disabling.\n");
+            tone_map.compute_peak = p->opts.tone_map.compute_peak = -1;
             detect_peak = false;
-            p->opts.compute_hdr_peak = -1;
         }
     }
 
@@ -2515,9 +2521,7 @@ static void pass_colormanage(struct gl_video *p, struct mp_colorspace src, bool
     }
 
     // Adapt from src to dst as necessary
-    pass_color_map(p->sc, src, dst, p->opts.tone_mapping,
-                   p->opts.tone_mapping_param, p->opts.tone_mapping_desat,
-                   detect_peak, p->opts.gamut_warning, p->use_linear && !osd);
+    pass_color_map(p->sc, p->use_linear && !osd, src, dst, &tone_map);
 
     if (p->use_lut_3d) {
         gl_sc_uniform_texture(p->sc, "lut_3d", p->lut_3d_texture);
@@ -3583,12 +3587,12 @@ static void check_gl_features(struct gl_video *p)
     }
 
     bool have_compute_peak = have_compute && have_ssbo;
-    if (!have_compute_peak && p->opts.compute_hdr_peak >= 0) {
-        int msgl = p->opts.compute_hdr_peak == 1 ? MSGL_WARN : MSGL_V;
+    if (!have_compute_peak && p->opts.tone_map.compute_peak >= 0) {
+        int msgl = p->opts.tone_map.compute_peak == 1 ? MSGL_WARN : MSGL_V;
         MP_MSG(p, msgl, "Disabling HDR peak computation (one or more of the "
                         "following is not supported: compute shaders=%d, "
                         "SSBO=%d).\n", have_compute, have_ssbo);
-        p->opts.compute_hdr_peak = -1;
+        p->opts.tone_map.compute_peak = -1;
     }
 
     p->forced_dumb_mode = p->opts.dumb_mode > 0 || !have_fbo || !have_texrg;
@@ -3610,7 +3614,6 @@ static void check_gl_features(struct gl_video *p)
             .alpha_mode = p->opts.alpha_mode,
             .use_rectangle = p->opts.use_rectangle,
             .background = p->opts.background,
-            .compute_hdr_peak = p->opts.compute_hdr_peak,
             .dither_algo = p->opts.dither_algo,
             .dither_depth = p->opts.dither_depth,
             .dither_size = p->opts.dither_size,
@@ -3618,9 +3621,7 @@ static void check_gl_features(struct gl_video *p)
             .temporal_dither_period = p->opts.temporal_dither_period,
             .tex_pad_x = p->opts.tex_pad_x,
             .tex_pad_y = p->opts.tex_pad_y,
-            .tone_mapping = p->opts.tone_mapping,
-            .tone_mapping_param = p->opts.tone_mapping_param,
-            .tone_mapping_desat = p->opts.tone_mapping_desat,
+            .tone_map = p->opts.tone_map,
             .early_flush = p->opts.early_flush,
             .icc_opts = p->opts.icc_opts,
             .hwdec_interop = p->opts.hwdec_interop,
diff --git a/video/out/gpu/video.h b/video/out/gpu/video.h
index ca8b6f65d4..ee5c0a2861 100644
--- a/video/out/gpu/video.h
+++ b/video/out/gpu/video.h
@@ -98,6 +98,15 @@ enum tone_mapping {
 // How many frames to average over for HDR peak detection
 #define PEAK_DETECT_FRAMES 63
 
+struct gl_tone_map_opts {
+    int curve;
+    float curve_param;
+    int compute_peak;
+    float desat;
+    float desat_exp;
+    int gamut_warning; // bool
+};
+
 struct gl_video_opts {
     int dumb_mode;
     struct scaler_config scaler[4];
@@ -107,11 +116,7 @@ struct gl_video_opts {
     int target_prim;
     int target_trc;
     int target_peak;
-    int tone_mapping;
-    int compute_hdr_peak;
-    float tone_mapping_param;
-    float tone_mapping_desat;
-    int gamut_warning;
+    struct gl_tone_map_opts tone_map;
     int correct_downscaling;
     int linear_downscaling;
     int linear_upscaling;
diff --git a/video/out/gpu/video_shaders.c b/video/out/gpu/video_shaders.c
index 342fb39ded..315e15cc89 100644
--- a/video/out/gpu/video_shaders.c
+++ b/video/out/gpu/video_shaders.c
@@ -580,7 +580,7 @@ static void hdr_update_peak(struct gl_shader_cache *sc)
 
     // Have each thread update the work group sum with the local value
     GLSL(barrier();)
-    GLSLF("atomicAdd(wg_sum, uint(sig * %f));\n", MP_REF_WHITE);
+    GLSLF("atomicAdd(wg_sum, uint(sig_max * %f));\n", MP_REF_WHITE);
 
     // Have one thread per work group update the global atomics. We use the
     // work group average even for the global sum, to make the values slightly
@@ -642,48 +642,42 @@ static void hdr_update_peak(struct gl_shader_cache *sc)
 
 // Tone map from a known peak brightness to the range [0,1]. If ref_peak
 // is 0, we will use peak detection instead
-static void pass_tone_map(struct gl_shader_cache *sc, bool detect_peak,
+static void pass_tone_map(struct gl_shader_cache *sc,
                           float src_peak, float dst_peak,
-                          enum tone_mapping algo, float param, float desat)
+                          const struct gl_tone_map_opts *opts)
 {
     GLSLF("// HDR tone mapping\n");
 
     // To prevent discoloration due to out-of-bounds clipping, we need to make
     // sure to reduce the value range as far as necessary to keep the entire
     // signal in range, so tone map based on the brightest component.
-    GLSL(float sig = max(max(color.r, color.g), color.b);)
+    GLSL(int sig_idx = 0;)
+    GLSL(if (color[1] > color[sig_idx]) sig_idx = 1;)
+    GLSL(if (color[2] > color[sig_idx]) sig_idx = 2;)
+    GLSL(float sig_max = color[sig_idx];)
     GLSLF("float sig_peak = %f;\n", src_peak);
     GLSLF("float sig_avg = %f;\n", sdr_avg);
 
-    if (detect_peak)
+    if (opts->compute_peak >= 0)
         hdr_update_peak(sc);
 
+    GLSLF("vec3 sig = color.rgb;\n");
+
     // Rescale the variables in order to bring it into a representation where
     // 1.0 represents the dst_peak. This is because all of the tone mapping
     // algorithms are defined in such a way that they map to the range [0.0, 1.0].
     if (dst_peak > 1.0) {
-        GLSLF("sig *= %f;\n", 1.0 / dst_peak);
-        GLSLF("sig_peak *= %f;\n", 1.0 / dst_peak);
+        GLSLF("sig *= 1.0/%f;\n", dst_peak);
+        GLSLF("sig_peak *= 1.0/%f;\n", dst_peak);
     }
 
-    GLSL(float sig_orig = sig;)
+    GLSL(float sig_orig = sig[sig_idx];)
     GLSLF("float slope = min(1.0, %f / sig_avg);\n", sdr_avg);
     GLSL(sig *= slope;)
     GLSL(sig_peak *= slope;)
 
-    // Desaturate the color using a coefficient dependent on the signal.
-    // Do this after peak detection in order to prevent over-desaturating
-    // overly bright souces
-    if (desat > 0) {
-        float base = 0.18 * dst_peak;
-        GLSL(float luma = dot(dst_luma, color.rgb);)
-        GLSLF("float coeff = max(sig - %f, 1e-6) / max(sig, 1e-6);\n", base);
-        GLSLF("coeff = pow(coeff, %f);\n", 10.0 / desat);
-        GLSL(color.rgb = mix(color.rgb, vec3(luma), coeff);)
-        GLSL(sig = mix(sig, luma * slope, coeff);) // also make sure to update `sig`
-    }
-
-    switch (algo) {
+    float param = opts->curve_param;
+    switch (opts->curve) {
     case TONE_MAPPING_CLIP:
         GLSLF("sig = %f * sig;\n", isnan(param) ? 1.0 : param);
         break;
@@ -697,14 +691,15 @@ static void pass_tone_map(struct gl_shader_cache *sc, bool detect_peak,
         GLSLF("float b = (j*j - 2.0*j*sig_peak + sig_peak) / "
               "max(1e-6, sig_peak - 1.0);\n");
         GLSLF("float scale = (b*b + 2.0*b*j + j*j) / (b-a);\n");
-        GLSL(sig = sig > j ? scale * (sig + a) / (sig + b) : sig;)
+        GLSLF("sig = mix(sig, scale * (sig + vec3(a)) / (sig + vec3(b)),"
+              "          greaterThan(sig, vec3(j)));\n");
         GLSLF("}\n");
         break;
 
     case TONE_MAPPING_REINHARD: {
         float contrast = isnan(param) ? 0.5 : param,
               offset = (1.0 - contrast) / contrast;
-        GLSLF("sig = sig / (sig + %f);\n", offset);
+        GLSLF("sig = sig / (sig + vec3(%f));\n", offset);
         GLSLF("float scale = (sig_peak + %f) / sig_peak;\n", offset);
         GLSL(sig *= scale;)
         break;
@@ -712,19 +707,25 @@ static void pass_tone_map(struct gl_shader_cache *sc, bool detect_peak,
 
     case TONE_MAPPING_HABLE: {
         float A = 0.15, B = 0.50, C = 0.10, D = 0.20, E = 0.02, F = 0.30;
-        GLSLHF("float hable(float x) {\n");
-        GLSLHF("return ((x * (%f*x + %f)+%f)/(x * (%f*x + %f) + %f)) - %f;\n",
-               A, C*B, D*E, A, B, D*F, E/F);
+        GLSLHF("vec3 hable(vec3 x) {\n");
+        GLSLHF("return (x * (%f*x + vec3(%f)) + vec3(%f)) / "
+               "       (x * (%f*x + vec3(%f)) + vec3(%f)) "
+               "       - vec3(%f);\n",
+               A, C*B, D*E,
+               A, B, D*F,
+               E/F);
         GLSLHF("}\n");
-        GLSL(sig = hable(sig) / hable(sig_peak);)
+        GLSLF("sig = hable(max(vec3(0.0), sig)) / hable(vec3(sig_peak)).x;\n");
         break;
     }
 
     case TONE_MAPPING_GAMMA: {
         float gamma = isnan(param) ? 1.8 : param;
-        GLSLF("const float cutoff = 0.05, gamma = %f;\n", 1.0/gamma);
-        GLSL(float scale = pow(cutoff / sig_peak, gamma) / cutoff;)
-        GLSL(sig = sig > cutoff ? pow(sig / sig_peak, gamma) : scale * sig;)
+        GLSLF("const float cutoff = 0.05, gamma = 1.0/%f;\n", gamma);
+        GLSL(float scale = pow(cutoff / sig_peak, gamma.x) / cutoff;)
+        GLSLF("sig = mix(scale * sig,"
+              "          pow(sig / sig_peak, vec3(gamma)),"
+              "          greaterThan(sig, vec3(cutoff)));\n");
         break;
     }
 
@@ -738,24 +739,32 @@ static void pass_tone_map(struct gl_shader_cache *sc, bool detect_peak,
         abort();
     }
 
-    // Apply the computed scale factor to the color, linearly to prevent
-    // discoloration
-    GLSL(sig = min(sig, 1.0);)
-    GLSL(color.rgb *= vec3(sig / sig_orig);)
+    GLSL(sig = min(sig, vec3(1.0));)
+    GLSL(vec3 sig_lin = color.rgb * (sig[sig_idx] / sig_orig);)
+
+    // Mix between the per-channel tone mapped and the linear tone mapped
+    // signal based on the desaturation strength
+    if (opts->desat > 0) {
+        float base = 0.18 * dst_peak;
+        GLSLF("float coeff = max(sig[sig_idx] - %f, 1e-6) / "
+              "              max(sig[sig_idx], 1.0);\n", base);
+        GLSLF("coeff = %f * pow(coeff, %f);\n", opts->desat, opts->desat_exp);
+        GLSLF("color.rgb = mix(sig_lin, %f * sig, coeff);\n", dst_peak);
+    } else {
+        GLSL(color.rgb = sig_lin;)
+    }
 }
 
 // Map colors from one source space to another. These source spaces must be
 // known (i.e. not MP_CSP_*_AUTO), as this function won't perform any
 // auto-guessing. If is_linear is true, we assume the input has already been
-// linearized (e.g. for linear-scaling). If `detect_peak` is true, we will
-// detect the peak instead of relying on metadata. Note that this requires
-// the caller to have already bound the appropriate SSBO and set up the
-// compute shader metadata
-void pass_color_map(struct gl_shader_cache *sc,
+// linearized (e.g. for linear-scaling). If `opts->compute_peak` is true, we
+// will detect the peak instead of relying on metadata. Note that this requires
+// the caller to have already bound the appropriate SSBO and set up the compute
+// shader metadata
+void pass_color_map(struct gl_shader_cache *sc, bool is_linear,
                     struct mp_colorspace src, struct mp_colorspace dst,
-                    enum tone_mapping algo, float tone_mapping_param,
-                    float tone_mapping_desat, bool detect_peak,
-                    bool gamut_warning, bool is_linear)
+                    const struct gl_tone_map_opts *opts)
 {
     GLSLF("// color mapping\n");
 
@@ -803,10 +812,8 @@ void pass_color_map(struct gl_shader_cache *sc,
 
     // Tone map to prevent clipping when the source signal peak exceeds the
     // encodable range or we've reduced the gamut
-    if (src.sig_peak > dst.sig_peak) {
-        pass_tone_map(sc, detect_peak, src.sig_peak, dst.sig_peak, algo,
-                      tone_mapping_param, tone_mapping_desat);
-    }
+    if (src.sig_peak > dst.sig_peak)
+        pass_tone_map(sc, src.sig_peak, dst.sig_peak, opts);
 
     if (need_ootf)
         pass_inverse_ootf(sc, dst.light, dst.sig_peak);
@@ -821,7 +828,7 @@ void pass_color_map(struct gl_shader_cache *sc,
     GLSLF("color.rgb *= vec3(%f);\n", 1.0 / dst_range);
 
     // Warn for remaining out-of-gamut colors is enabled
-    if (gamut_warning) {
+    if (opts->gamut_warning) {
         GLSL(if (any(greaterThan(color.rgb, vec3(1.01)))))
             GLSL(color.rgb = vec3(1.0) - color.rgb;) // invert
     }
diff --git a/video/out/gpu/video_shaders.h b/video/out/gpu/video_shaders.h
index cd395d6377..f20d643e99 100644
--- a/video/out/gpu/video_shaders.h
+++ b/video/out/gpu/video_shaders.h
@@ -40,11 +40,9 @@ void pass_sample_oversample(struct gl_shader_cache *sc, struct scaler *scaler,
 void pass_linearize(struct gl_shader_cache *sc, enum mp_csp_trc trc);
 void pass_delinearize(struct gl_shader_cache *sc, enum mp_csp_trc trc);
 
-void pass_color_map(struct gl_shader_cache *sc,
+void pass_color_map(struct gl_shader_cache *sc, bool is_linear,
                     struct mp_colorspace src, struct mp_colorspace dst,
-                    enum tone_mapping algo, float tone_mapping_param,
-                    float tone_mapping_desat, bool use_detected_peak,
-                    bool gamut_warning, bool is_linear);
+                    const struct gl_tone_map_opts *opts);
 
 void pass_sample_deband(struct gl_shader_cache *sc, struct deband_opts *opts,
                         AVLFG *lfg, enum mp_csp_trc trc);