add floating point values support in colormap Keys (#765)

CHANGES.md

@@ -1,3 +1,7 @@
+# 7.2.1 (2024-11-14)
+
+* add official support for floating point values in ColorMap
+
 # 7.2.0 (2024-11-05)
 
 * Ensure compatibility between XarrayReader and other Readers by adding `**kwargs` on class methods (https://github.com/cogeotiff/rio-tiler/pull/762)

rio_tiler/colormap.py

@@ -18,6 +18,7 @@
 from rio_tiler.types import (
     ColorMapType,
     DataMaskType,
+    DiscreteColorMapType,
     GDALColorMapType,
     IntervalColorMapType,
 )
@@ -116,10 +117,16 @@ def apply_cmap(data: numpy.ndarray, colormap: ColorMapType) -> DataMaskType:
     # rio_tiler.colormap.make_lut, because we don't want to create a `lookup table`
     # with more than 256 entries (256 x 4) array. In this case we use `apply_discrete_cmap`
     # which can work with arbitrary colormap dict.
-    if len(colormap) != 256 or max(colormap) >= 256 or min(colormap) < 0:
+    if (
+        len(colormap) != 256
+        or max(colormap) >= 256
+        or min(colormap) < 0
+        or any(isinstance(k, float) for k in colormap)
+    ):
         return apply_discrete_cmap(data, colormap)
 
-    lookup_table = make_lut(colormap)
+    cm = {int(k): v for k, v in colormap.items()}
+    lookup_table = make_lut(cm)  # type: ignore
     data = lookup_table[data[0], :]
 
     data = numpy.transpose(data, [2, 0, 1])
@@ -132,12 +139,14 @@ def apply_cmap(data: numpy.ndarray, colormap: ColorMapType) -> DataMaskType:
     return data[:-1], data[-1]
 
 
-def apply_discrete_cmap(data: numpy.ndarray, colormap: GDALColorMapType) -> DataMaskType:
+def apply_discrete_cmap(
+    data: numpy.ndarray, colormap: Union[GDALColorMapType, DiscreteColorMapType]
+) -> DataMaskType:
     """Apply discrete colormap.
 
     Args:
         data (numpy.ndarray): 1D image array to translate to RGB.
-        colormap (GDALColorMapType): Discrete ColorMap dictionary.
+        colormap (GDALColorMapType or DiscreteColorMapType): Discrete ColorMap dictionary.
 
     Returns:
         tuple: Data (numpy.ndarray) and Alpha band (numpy.ndarray).

rio_tiler/types.py

@@ -18,11 +18,15 @@
 # ColorMap Dict: {1: (0, 0, 0, 255), ...}
 GDALColorMapType = Dict[int, ColorTuple]
 
+# Discrete Colormap, like GDALColorMapType but accept Float: {0.1: (0, 0, 0, 255), ...}
+DiscreteColorMapType = Dict[NumType, ColorTuple]
+
 # Intervals ColorMap: [((0, 1), (0, 0, 0, 0)), ...]
 IntervalColorMapType = Sequence[Tuple[IntervalTuple, ColorTuple]]
 
 ColorMapType = Union[
     GDALColorMapType,
+    DiscreteColorMapType,
     IntervalColorMapType,
 ]
 

tests/test_cmap.py

@@ -301,6 +301,8 @@ def test_parse_color_bad():
 
 def test_discrete_float():
     """test for titiler issue 738."""
+
+    # make sure we apply discrete colormap when we have less than 256 cmap entries
     cm = {
         0: (0, 255, 255, 255),
         1: (83, 151, 145, 255),
@@ -325,3 +327,20 @@ def test_discrete_float():
     dd, mm = colormap.apply_discrete_cmap(data.copy(), cm)
     assert d.dtype == numpy.uint8
     assert m.dtype == numpy.uint8
+    numpy.testing.assert_array_equal(d, dd)
+    numpy.testing.assert_array_equal(m, mm)
+
+    # make we allow float keys in discrete colormap
+    cm = {
+        0.5: (0, 255, 255, 255),
+        1.5: (83, 151, 145, 255),
+        2.5: (87, 194, 23, 255),
+    }
+
+    data = numpy.random.choice([0.5, 2.5], 256 * 256).reshape(1, 256, 256)
+    d, m = colormap.apply_cmap(data.copy(), cm)
+    dd, mm = colormap.apply_discrete_cmap(data.copy(), cm)
+    assert d.dtype == numpy.uint8
+    assert m.dtype == numpy.uint8
+    numpy.testing.assert_array_equal(d, dd)
+    numpy.testing.assert_array_equal(m, mm)