From 993ed667f9189f12cb2671dda2187a4898e6158f Mon Sep 17 00:00:00 2001
From: Biser Hong <biser@skyfi.com>
Date: Thu, 12 Dec 2024 11:58:34 +0200
Subject: [PATCH] use color_operations for custom color transforms

---
 terracotta/handlers/rgb.py | 23 ++++++++++++++++-------
 terracotta/image.py        | 16 ++++++++++++++++
 terracotta/server/rgb.py   |  4 ++++
 3 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/terracotta/handlers/rgb.py b/terracotta/handlers/rgb.py
index 4d3b6d8f..f5f2b46e 100644
--- a/terracotta/handlers/rgb.py
+++ b/terracotta/handlers/rgb.py
@@ -24,6 +24,7 @@ def rgb(
     *,
     stretch_ranges: Optional[ListOfRanges] = None,
     gamma_factor: Optional[float] = None,
+    color_transform: Optional[str] = None,
     tile_size: Optional[Tuple[int, int]] = None
 ) -> BinaryIO:
     """Return RGB image as PNG
@@ -81,6 +82,7 @@ def get_band_future(band_key: str) -> Future:
         futures = [get_band_future(key) for key in rgb_values]
         band_items = zip(rgb_values, stretch_ranges_, futures)
 
+        out_ranges = []
         out_arrays = []
 
         for i, (band_key, band_stretch_override, band_data_future) in enumerate(
@@ -107,14 +109,21 @@ def get_band_future(band_key: str) -> Future:
 
             band_data = band_data_future.result()
 
-            if gamma_factor:
-                # gamma correction is monotonic and preserves percentiles
-                band_stretch_range_arr = np.array(band_stretch_range, dtype=band_data.dtype)
-                band_stretch_range = list(image.gamma_correction(band_stretch_range_arr, gamma_factor, band_range))
-                # gamma correct band data
-                band_data = image.gamma_correction(band_data, gamma_factor, band_range)
+            out_ranges.append(band_stretch_range)
+            out_arrays.append(band_data)
 
-            out_arrays.append(image.to_uint8(band_data, *band_stretch_range))
+    out = np.ma.stack(out_arrays, axis=0)
+
+    if color_transform:
+        band_stretch_range_arr = [np.array(band_rng, dtype=band_data.dtype) for band_rng in out_ranges]
+        band_stretch_range_arr = np.ma.stack(band_stretch_range_arr, axis=0)
+
+        band_stretch_range_arr = image.apply_color_transform(band_stretch_range_arr, color_transform)
+        band_data = image.apply_color_transform(out, color_transform)
+
+    out_arrays = []
+    for k in range(band_data.shape[0]):
+        out_arrays.append(image.to_uint8(band_data[k], *band_stretch_range_arr[k]))
 
     out = np.ma.stack(out_arrays, axis=-1)
     return image.array_to_png(out)
diff --git a/terracotta/image.py b/terracotta/image.py
index 62788b14..76ec45dc 100755
--- a/terracotta/image.py
+++ b/terracotta/image.py
@@ -11,6 +11,7 @@
 
 import numpy as np
 from PIL import Image
+from color_operations import parse_operations
 from color_operations.operations import gamma
 from color_operations.utils import to_math_type, scale_dtype
 
@@ -187,6 +188,21 @@ def gamma_correction(
     return arr
 
 
+def apply_color_transform(
+        masked_data: Array,
+        color_transform: str,
+        out_dtype: type = np.uint16,
+) -> Array:
+    """Apply gamma correction to the input array and scale it to the output dtype."""
+    arr = to_math_type(masked_data)
+
+    for func in parse_operations(color_transform):
+        arr = func(arr)
+
+    arr = scale_dtype(arr, out_dtype)
+    return arr
+
+
 def label(data: Array, labels: Sequence[Number]) -> Array:
     """Create a labelled uint8 version of data, with output values starting at 1.
 
diff --git a/terracotta/server/rgb.py b/terracotta/server/rgb.py
index 74f12a38..2627a290 100644
--- a/terracotta/server/rgb.py
+++ b/terracotta/server/rgb.py
@@ -73,6 +73,10 @@ class Meta:
         missing=None,
         description="Gamma factor to perform gamma correction."
     )
+    color_transform = fields.String(
+        missing=None,
+        description="Gamma factor to perform gamma correction."
+    )
     tile_size = fields.List(
         fields.Integer(),
         validate=validate.Length(equal=2),