1 files changed, 225 insertions, 0 deletions

diff --git a/media/libjxl/src/lib/extras/tone_mapping.cc b/media/libjxl/src/lib/extras/tone_mapping.cc
new file mode 100644
index 0000000000..ac4306f8ed
--- /dev/null
+++ b/media/libjxl/src/lib/extras/tone_mapping.cc
@@ -0,0 +1,225 @@
+// Copyright (c) the JPEG XL Project Authors. All rights reserved.
+//
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+#include "lib/extras/tone_mapping.h"
+
+#undef HWY_TARGET_INCLUDE
+#define HWY_TARGET_INCLUDE "lib/extras/tone_mapping.cc"
+#include <hwy/foreach_target.h>
+#include <hwy/highway.h>
+
+#include "lib/jxl/enc_color_management.h"
+#include "lib/jxl/transfer_functions-inl.h"
+
+HWY_BEFORE_NAMESPACE();
+namespace jxl {
+namespace HWY_NAMESPACE {
+
+Status ToneMapFrame(const std::pair<float, float> display_nits,
+                    ImageBundle* const ib, ThreadPool* const pool) {
+  // Perform tone mapping as described in Report ITU-R BT.2390-8, section 5.4
+  // (pp. 23-25).
+  // https://www.itu.int/pub/R-REP-BT.2390-8-2020
+
+  HWY_FULL(float) df;
+  using V = decltype(Zero(df));
+
+  ColorEncoding linear_rec2020;
+  linear_rec2020.SetColorSpace(ColorSpace::kRGB);
+  linear_rec2020.primaries = Primaries::k2100;
+  linear_rec2020.white_point = WhitePoint::kD65;
+  linear_rec2020.tf.SetTransferFunction(TransferFunction::kLinear);
+  JXL_RETURN_IF_ERROR(linear_rec2020.CreateICC());
+  JXL_RETURN_IF_ERROR(ib->TransformTo(linear_rec2020, GetJxlCms(), pool));
+
+  const auto eotf_inv = [&df](const V luminance) -> V {
+    return TF_PQ().EncodedFromDisplay(df, luminance * Set(df, 1. / 10000));
+  };
+
+  const V pq_mastering_min =
+      eotf_inv(Set(df, ib->metadata()->tone_mapping.min_nits));
+  const V pq_mastering_max =
+      eotf_inv(Set(df, ib->metadata()->tone_mapping.intensity_target));
+  const V pq_mastering_range = pq_mastering_max - pq_mastering_min;
+  const V inv_pq_mastering_range =
+      Set(df, 1) / (pq_mastering_max - pq_mastering_min);
+  const V min_lum = (eotf_inv(Set(df, display_nits.first)) - pq_mastering_min) *
+                    inv_pq_mastering_range;
+  const V max_lum =
+      (eotf_inv(Set(df, display_nits.second)) - pq_mastering_min) *
+      inv_pq_mastering_range;
+  const V ks = MulAdd(Set(df, 1.5f), max_lum, Set(df, -0.5f));
+  const V b = min_lum;
+
+  const V inv_one_minus_ks = Set(df, 1) / Max(Set(df, 1e-6f), Set(df, 1) - ks);
+  const auto T = [ks, inv_one_minus_ks](const V a) {
+    return (a - ks) * inv_one_minus_ks;
+  };
+  const auto P = [&T, &df, ks, max_lum](const V b) {
+    const V t_b = T(b);
+    const V t_b_2 = t_b * t_b;
+    const V t_b_3 = t_b_2 * t_b;
+    return MulAdd(
+        MulAdd(Set(df, 2), t_b_3, MulAdd(Set(df, -3), t_b_2, Set(df, 1))), ks,
+        MulAdd(t_b_3 + MulAdd(Set(df, -2), t_b_2, t_b), Set(df, 1) - ks,
+               MulAdd(Set(df, -2), t_b_3, Set(df, 3) * t_b_2) * max_lum));
+  };
+
+  const V inv_max_display_nits = Set(df, 1 / display_nits.second);
+
+  JXL_RETURN_IF_ERROR(RunOnPool(
+      pool, 0, ib->ysize(), ThreadPool::NoInit,
+      [&](const uint32_t y, size_t /* thread */) {
+        float* const JXL_RESTRICT row_r = ib->color()->PlaneRow(0, y);
+        float* const JXL_RESTRICT row_g = ib->color()->PlaneRow(1, y);
+        float* const JXL_RESTRICT row_b = ib->color()->PlaneRow(2, y);
+        for (size_t x = 0; x < ib->xsize(); x += Lanes(df)) {
+          V red = Load(df, row_r + x);
+          V green = Load(df, row_g + x);
+          V blue = Load(df, row_b + x);
+          const V luminance = Set(df, ib->metadata()->IntensityTarget()) *
+                              (MulAdd(Set(df, 0.2627f), red,
+                                      MulAdd(Set(df, 0.6780f), green,
+                                             Set(df, 0.0593f) * blue)));
+          const V normalized_pq =
+              Min(Set(df, 1.f), (eotf_inv(luminance) - pq_mastering_min) *
+                                    inv_pq_mastering_range);
+          const V e2 =
+              IfThenElse(normalized_pq < ks, normalized_pq, P(normalized_pq));
+          const V one_minus_e2 = Set(df, 1) - e2;
+          const V one_minus_e2_2 = one_minus_e2 * one_minus_e2;
+          const V one_minus_e2_4 = one_minus_e2_2 * one_minus_e2_2;
+          const V e3 = MulAdd(b, one_minus_e2_4, e2);
+          const V e4 = MulAdd(e3, pq_mastering_range, pq_mastering_min);
+          const V new_luminance =
+              Min(Set(df, display_nits.second),
+                  ZeroIfNegative(Set(df, 10000) *
+                                 TF_PQ().DisplayFromEncoded(df, e4)));
+
+          const V ratio = new_luminance / luminance;
+          const V normalizer =
+              Set(df, ib->metadata()->IntensityTarget()) * inv_max_display_nits;
+
+          for (V* const val : {&red, &green, &blue}) {
+            *val = IfThenElse(luminance <= Set(df, 1e-6f), new_luminance,
+                              *val * ratio) *
+                   normalizer;
+          }
+
+          Store(red, df, row_r + x);
+          Store(green, df, row_g + x);
+          Store(blue, df, row_b + x);
+        }
+      },
+      "ToneMap"));
+
+  return true;
+}
+
+Status GamutMapFrame(ImageBundle* const ib, float preserve_saturation,
+                     ThreadPool* const pool) {
+  HWY_FULL(float) df;
+  using V = decltype(Zero(df));
+
+  ColorEncoding linear_rec2020;
+  linear_rec2020.SetColorSpace(ColorSpace::kRGB);
+  linear_rec2020.primaries = Primaries::k2100;
+  linear_rec2020.white_point = WhitePoint::kD65;
+  linear_rec2020.tf.SetTransferFunction(TransferFunction::kLinear);
+  JXL_RETURN_IF_ERROR(linear_rec2020.CreateICC());
+  JXL_RETURN_IF_ERROR(ib->TransformTo(linear_rec2020, GetJxlCms(), pool));
+
+  JXL_RETURN_IF_ERROR(RunOnPool(
+      pool, 0, ib->ysize(), ThreadPool::NoInit,
+      [&](const uint32_t y, size_t /* thread*/) {
+        float* const JXL_RESTRICT row_r = ib->color()->PlaneRow(0, y);
+        float* const JXL_RESTRICT row_g = ib->color()->PlaneRow(1, y);
+        float* const JXL_RESTRICT row_b = ib->color()->PlaneRow(2, y);
+        for (size_t x = 0; x < ib->xsize(); x += Lanes(df)) {
+          V red = Load(df, row_r + x);
+          V green = Load(df, row_g + x);
+          V blue = Load(df, row_b + x);
+          const V luminance =
+              MulAdd(Set(df, 0.2627f), red,
+                     MulAdd(Set(df, 0.6780f), green, Set(df, 0.0593f) * blue));
+
+          // Desaturate out-of-gamut pixels. This is done by mixing each pixel
+          // with just enough gray of the target luminance to make all
+          // components non-negative.
+          // - For saturation preservation, if a component is still larger than
+          // 1 then the pixel is normalized to have a maximum component of 1.
+          // That will reduce its luminance.
+          // - For luminance preservation, getting all components below 1 is
+          // done by mixing in yet more gray. That will desaturate it further.
+          V gray_mix_saturation = Zero(df);
+          V gray_mix_luminance = Zero(df);
+          for (const V val : {red, green, blue}) {
+            const V inv_val_minus_gray = Set(df, 1) / (val - luminance);
+            gray_mix_saturation =
+                IfThenElse(val >= luminance, gray_mix_saturation,
+                           Max(gray_mix_saturation, val * inv_val_minus_gray));
+            gray_mix_luminance =
+                Max(gray_mix_luminance,
+                    IfThenElse(val <= luminance, gray_mix_saturation,
+                               (val - Set(df, 1)) * inv_val_minus_gray));
+          }
+          const V gray_mix =
+              Clamp(Set(df, preserve_saturation) *
+                            (gray_mix_saturation - gray_mix_luminance) +
+                        gray_mix_luminance,
+                    Zero(df), Set(df, 1));
+          for (V* const val : {&red, &green, &blue}) {
+            *val = MulAdd(gray_mix, luminance - *val, *val);
+          }
+          const V normalizer =
+              Set(df, 1) / Max(Set(df, 1), Max(red, Max(green, blue)));
+          for (V* const val : {&red, &green, &blue}) {
+            *val = *val * normalizer;
+          }
+
+          Store(red, df, row_r + x);
+          Store(green, df, row_g + x);
+          Store(blue, df, row_b + x);
+        }
+      },
+      "GamutMap"));
+
+  return true;
+}
+
+// NOLINTNEXTLINE(google-readability-namespace-comments)
+}  // namespace HWY_NAMESPACE
+}  // namespace jxl
+HWY_AFTER_NAMESPACE();
+
+#if HWY_ONCE
+namespace jxl {
+
+namespace {
+HWY_EXPORT(ToneMapFrame);
+HWY_EXPORT(GamutMapFrame);
+}  // namespace
+
+Status ToneMapTo(const std::pair<float, float> display_nits,
+                 CodecInOut* const io, ThreadPool* const pool) {
+  const auto tone_map_frame = HWY_DYNAMIC_DISPATCH(ToneMapFrame);
+  for (ImageBundle& ib : io->frames) {
+    JXL_RETURN_IF_ERROR(tone_map_frame(display_nits, &ib, pool));
+  }
+  io->metadata.m.SetIntensityTarget(display_nits.second);
+  return true;
+}
+
+Status GamutMap(CodecInOut* const io, float preserve_saturation,
+                ThreadPool* const pool) {
+  const auto gamut_map_frame = HWY_DYNAMIC_DISPATCH(GamutMapFrame);
+  for (ImageBundle& ib : io->frames) {
+    JXL_RETURN_IF_ERROR(gamut_map_frame(&ib, preserve_saturation, pool));
+  }
+  return true;
+}
+
+}  // namespace jxl
+#endif