Add rgb_range percentile support (#322)

Co-authored-by: Dion Häfner <[email protected]>
Co-authored-by: Dion Häfner <[email protected]>

terracotta/handlers/rgb.py

@@ -10,8 +10,10 @@
 from terracotta import get_settings, get_driver, image, xyz, exceptions
 from terracotta.profile import trace
 
-Number = TypeVar("Number", int, float)
-ListOfRanges = Sequence[Optional[Tuple[Optional[Number], Optional[Number]]]]
+NumberOrString = TypeVar("NumberOrString", int, float, str)
+ListOfRanges = Sequence[
+    Optional[Tuple[Optional[NumberOrString], Optional[NumberOrString]]]
+]
 
 
 @trace("rgb_handler")
@@ -89,11 +91,12 @@ def get_band_future(band_key: str) -> Future:
             band_stretch_range = list(metadata["range"])
             scale_min, scale_max = band_stretch_override
 
+            percentiles = metadata.get("percentiles", [])
             if scale_min is not None:
-                band_stretch_range[0] = scale_min
+                band_stretch_range[0] = image.get_stretch_scale(scale_min, percentiles)
 
             if scale_max is not None:
-                band_stretch_range[1] = scale_max
+                band_stretch_range[1] = image.get_stretch_scale(scale_max, percentiles)
 
             if band_stretch_range[1] < band_stretch_range[0]:
                 raise exceptions.InvalidArgumentsError(

terracotta/handlers/singleband.py

@@ -12,6 +12,10 @@
 from terracotta.profile import trace
 
 Number = TypeVar("Number", int, float)
+NumberOrString = TypeVar("NumberOrString", int, float, str)
+ListOfRanges = Sequence[
+    Optional[Tuple[Optional[NumberOrString], Optional[NumberOrString]]]
+]
 RGBA = Tuple[Number, Number, Number, Number]
 
 
@@ -21,7 +25,7 @@ def singleband(
     tile_xyz: Optional[Tuple[int, int, int]] = None,
     *,
     colormap: Union[str, Mapping[Number, RGBA], None] = None,
-    stretch_range: Optional[Tuple[Number, Number]] = None,
+    stretch_range: Optional[Tuple[NumberOrString, NumberOrString]] = None,
     tile_size: Optional[Tuple[int, int]] = None
 ) -> BinaryIO:
     """Return singleband image as PNG"""
@@ -59,11 +63,12 @@ def singleband(
         # determine stretch range from metadata and arguments
         stretch_range_ = list(metadata["range"])
 
+        percentiles = metadata.get("percentiles", [])
         if stretch_min is not None:
-            stretch_range_[0] = stretch_min
+            stretch_range_[0] = image.get_stretch_scale(stretch_min, percentiles)
 
         if stretch_max is not None:
-            stretch_range_[1] = stretch_max
+            stretch_range_[1] = image.get_stretch_scale(stretch_max, percentiles)
 
         cmap_or_palette = cast(Optional[str], colormap)
         out = image.to_uint8(tile_data, *stretch_range_)

terracotta/image.py

@@ -3,7 +3,7 @@
 Utilities to create and manipulate images.
 """
 
-from typing import Sequence, Tuple, TypeVar, Union
+from typing import List, Sequence, Tuple, TypeVar, Union
 from typing.io import BinaryIO
 
 from io import BytesIO
@@ -15,6 +15,7 @@
 from terracotta import exceptions, get_settings
 
 Number = TypeVar("Number", int, float)
+NumberOrString = TypeVar("NumberOrString", int, float, str)
 RGBA = Tuple[Number, Number, Number, Number]
 Palette = Sequence[RGBA]
 Array = Union[np.ndarray, np.ma.MaskedArray]
@@ -181,3 +182,28 @@ def label(data: Array, labels: Sequence[Number]) -> Array:
         out_data[data == label] = i
 
     return out_data
+
+
+def get_stretch_scale(
+    scale: NumberOrString, percentiles: List[int]
+) -> Union[int, float]:
+    if isinstance(scale, (int, float)):
+        return scale
+    if isinstance(scale, str):
+        # can be a percentile
+        if scale.startswith("p"):
+            try:
+                percentile = int(scale[1:])
+            except ValueError:
+                raise exceptions.InvalidArgumentsError(
+                    f"Invalid percentile value: {scale}"
+                )
+
+            if 0 <= percentile < len(percentiles):
+                return percentiles[percentile]
+
+            raise exceptions.InvalidArgumentsError(
+                f"Invalid percentile, out of range: {scale}"
+            )
+
+    raise exceptions.InvalidArgumentsError(f"Invalid scale value: {scale}")

terracotta/server/fields.py

@@ -0,0 +1,47 @@
+"""server/fields.py
+
+Custom marshmallow fields for the server API.
+"""
+
+import re
+from marshmallow import ValidationError, fields
+
+from typing import Any, Union
+
+
+class StringOrNumber(fields.Field):
+    """
+    Marshmallow type that can be either a string or a number.
+    Uses marshmallow's default serialization/deserialization
+    for `String` or a `Float` depending on the value type.
+    """
+
+    def _serialize(
+        self, value: Union[str, bytes, int, float], attr: Any, obj: Any, **kwargs: Any
+    ) -> Union[str, float, None]:
+        if isinstance(value, (str, bytes)):
+            return fields.String()._serialize(value, attr, obj, **kwargs)
+        elif isinstance(value, (int, float)):
+            return fields.Float()._serialize(value, attr, obj, **kwargs)
+        else:
+            raise ValidationError("Must be a string or a number")
+
+    def _deserialize(
+        self, value: Union[str, bytes, int, float], attr: Any, data: Any, **kwargs: Any
+    ) -> Union[str, float, None]:
+        if isinstance(value, (str, bytes)):
+            return fields.String()._deserialize(value, attr, data, **kwargs)
+        elif isinstance(value, (int, float)):
+            return fields.Float()._deserialize(value, attr, data, **kwargs)
+        else:
+            raise ValidationError("Must be a string or a number")
+
+
+def validate_stretch_range(data: Any) -> None:
+    """
+    Validates that the stretch range is in the format `p<digits>`
+    when a string is used.
+    """
+    if isinstance(data, str):
+        if not re.match("^p\\d+$", data):
+            raise ValidationError("Percentile format is `p<digits>`")

terracotta/server/rgb.py

@@ -9,6 +9,7 @@
 from marshmallow import Schema, fields, validate, pre_load, ValidationError, EXCLUDE
 from flask import request, send_file, Response
 
+from terracotta.server.fields import StringOrNumber, validate_stretch_range
 from terracotta.server.flask_api import TILE_API
 
 
@@ -29,25 +30,43 @@ class Meta:
     g = fields.String(required=True, description="Key value for green band")
     b = fields.String(required=True, description="Key value for blue band")
     r_range = fields.List(
-        fields.Number(allow_none=True),
+        StringOrNumber(allow_none=True, validate=validate_stretch_range),
         validate=validate.Length(equal=2),
         example="[0,1]",
         missing=None,
-        description="Stretch range [min, max] to use for red band as JSON array",
+        description=(
+            "Stretch range [min, max] to use for the red band as JSON array. "
+            "Min and max may be numbers to use as absolute range, or strings "
+            "of the format `p<digits>` with an integer between 0 and 100 "
+            "to use percentiles of the image instead. "
+            "Null values indicate global minimum / maximum."
+        ),
     )
     g_range = fields.List(
-        fields.Number(allow_none=True),
+        StringOrNumber(allow_none=True, validate=validate_stretch_range),
         validate=validate.Length(equal=2),
         example="[0,1]",
         missing=None,
-        description="Stretch range [min, max] to use for green band as JSON array",
+        description=(
+            "Stretch range [min, max] to use for the gren band as JSON array. "
+            "Min and max may be numbers to use as absolute range, or strings "
+            "of the format `p<digits>` with an integer between 0 and 100 "
+            "to use percentiles of the image instead. "
+            "Null values indicate global minimum / maximum."
+        ),
     )
     b_range = fields.List(
-        fields.Number(allow_none=True),
+        StringOrNumber(allow_none=True, validate=validate_stretch_range),
         validate=validate.Length(equal=2),
         example="[0,1]",
         missing=None,
-        description="Stretch range [min, max] to use for blue band as JSON array",
+        description=(
+            "Stretch range [min, max] to use for the blue band as JSON array. "
+            "Min and max may be numbers to use as absolute range, or strings "
+            "of the format `p<digits>` with an integer between 0 and 100 "
+            "to use percentiles of the image instead. "
+            "Null values indicate global minimum / maximum."
+        ),
     )
     tile_size = fields.List(
         fields.Integer(),

terracotta/server/singleband.py

@@ -17,6 +17,7 @@
 )
 from flask import request, send_file, Response
 
+from terracotta.server.fields import StringOrNumber, validate_stretch_range
 from terracotta.server.flask_api import TILE_API
 from terracotta.cmaps import AVAILABLE_CMAPS
 
@@ -35,11 +36,16 @@ class Meta:
         unknown = EXCLUDE
 
     stretch_range = fields.List(
-        fields.Number(allow_none=True),
+        StringOrNumber(allow_none=True, validate=validate_stretch_range),
         validate=validate.Length(equal=2),
         example="[0,1]",
-        description="Stretch range to use as JSON array, uses full range by default. "
-        "Null values indicate global minimum / maximum.",
+        description=(
+            "Stretch range [min, max] to use as JSON array. "
+            "Min and max may be numbers to use as absolute range, or strings "
+            "of the format `p<digits>` with an integer between 0 and 100 "
+            "to use percentiles of the image instead. "
+            "Null values indicate global minimum / maximum."
+        ),
         missing=None,
     )
 

tests/handlers/test_rgb.py

@@ -3,6 +3,8 @@
 
 import pytest
 
+from terracotta import exceptions
+
 
 def test_rgb_handler(use_testdb, raster_file, raster_file_xyz):
     import terracotta
@@ -74,7 +76,13 @@ def test_rgb_lowzoom(use_testdb, raster_file, raster_file_xyz_lowzoom):
 
 
 @pytest.mark.parametrize(
-    "stretch_range", [[0, 20000], [10000, 20000], [-50000, 50000], [100, 100]]
+    "stretch_range",
+    [
+        [0, 20000],
+        [10000, 20000],
+        [-50000, 50000],
+        [100, 100],
+    ],
 )
 def test_rgb_stretch(stretch_range, use_testdb, testdb, raster_file_xyz):
     import terracotta
@@ -131,6 +139,82 @@ def test_rgb_invalid_stretch(use_testdb, raster_file_xyz):
         )
 
 
+def test_rgb_percentile_stretch(use_testdb, testdb, raster_file_xyz):
+    import terracotta
+    from terracotta.xyz import get_tile_data
+    from terracotta.handlers import rgb
+
+    ds_keys = ["val21", "x", "val22"]
+    bands = ["val22", "val23", "val24"]
+    pct_stretch_range = ["p2", "p98"]
+
+    raw_img = rgb.rgb(
+        ds_keys[:2],
+        bands,
+        raster_file_xyz,
+        stretch_ranges=[pct_stretch_range] * 3,
+    )
+    img_data = np.asarray(Image.open(raw_img))[..., 0]
+
+    # get unstretched data to compare to
+    driver = terracotta.get_driver(testdb)
+
+    with driver.connect():
+        tile_data = get_tile_data(
+            driver, ds_keys, tile_xyz=raster_file_xyz, tile_size=img_data.shape
+        )
+        band_metadata = driver.get_metadata(ds_keys)
+
+        stretch_range = [
+            band_metadata["percentiles"][1],
+            band_metadata["percentiles"][97],
+        ]
+
+    # filter transparent values
+    valid_mask = ~tile_data.mask
+    assert np.all(img_data[~valid_mask] == 0)
+
+    valid_img = img_data[valid_mask]
+    valid_data = tile_data.compressed()
+
+    assert np.all(valid_img[valid_data < stretch_range[0]] == 1)
+    stretch_range_mask = (valid_data > stretch_range[0]) & (
+        valid_data < stretch_range[1]
+    )
+    assert np.all(valid_img[stretch_range_mask] >= 1)
+    assert np.all(valid_img[stretch_range_mask] <= 255)
+    assert np.all(valid_img[valid_data > stretch_range[1]] == 255)
+
+
+@pytest.mark.parametrize(
+    "stretch_range_params",
+    [
+        ["s2", "p98", "Invalid scale value"],
+        ["pp2", "p98", "Invalid percentile value"],
+        ["p", "p98", "Invalid percentile value"],
+        ["2", "p8", "Invalid scale value"],
+        [{}, "p98", "Invalid scale value"],
+        ["p-2", "p98", "Invalid percentile, out of range"],
+        ["p2", "p298", "Invalid percentile, out of range"],
+    ],
+)
+def test_rgb_invalid_percentiles(use_testdb, raster_file_xyz, stretch_range_params):
+    from terracotta.handlers import rgb
+
+    ds_keys = ["val21", "x", "val22"]
+    bands = ["val22", "val23", "val24"]
+
+    stretch_range = stretch_range_params[:2]
+    with pytest.raises(exceptions.InvalidArgumentsError) as err:
+        rgb.rgb(
+            ds_keys[:2],
+            bands,
+            raster_file_xyz,
+            stretch_ranges=[stretch_range] * 3,
+        )
+    assert stretch_range_params[2] in str(err.value)
+
+
 def test_rgb_preview(use_testdb):
     import terracotta
     from terracotta.handlers import rgb