cast data to uint8 (#766)

CHANGES.md

@@ -1,6 +1,7 @@
 # 7.2.1 (2024-11-14)
 
 * add official support for floating point values in ColorMap
+* cast data to `uint8` datatype when applying linear colormap
 
 # 7.2.0 (2024-11-05)
 

rio_tiler/colormap.py

@@ -4,6 +4,7 @@
 import os
 import pathlib
 import re
+import warnings
 from typing import Dict, List, Sequence, Tuple, Union
 
 import attr
@@ -125,17 +126,19 @@ def apply_cmap(data: numpy.ndarray, colormap: ColorMapType) -> DataMaskType:
     ):
         return apply_discrete_cmap(data, colormap)
 
-    cm = {int(k): v for k, v in colormap.items()}
-    lookup_table = make_lut(cm)  # type: ignore
+    # For now we assume ColorMap are in uint8
+    if data.dtype != numpy.uint8:
+        warnings.warn(
+            f"Input array is of type {data.dtype} and `will be converted to Int in order to apply the ColorMap.",
+            UserWarning,
+        )
+        data = data.astype(numpy.uint8)
+
+    lookup_table = make_lut(colormap)  # type: ignore
     data = lookup_table[data[0], :]
 
     data = numpy.transpose(data, [2, 0, 1])
 
-    # If the colormap has values between 0-255
-    # we cast the output array to Uint8.
-    if data.min() >= 0 and data.max() <= 255:
-        data = data.astype("uint8")
-
     return data[:-1], data[-1]
 
 

tests/test_io_rasterio.py

@@ -17,6 +17,7 @@
 from rasterio.vrt import WarpedVRT
 from rasterio.warp import transform_bounds
 
+from rio_tiler.colormap import cmap
 from rio_tiler.constants import WEB_MERCATOR_TMS, WGS84_CRS
 from rio_tiler.errors import (
     ExpressionMixingWarning,
@@ -1120,3 +1121,27 @@ def test_inverted_latitude():
     with pytest.warns(UserWarning):
         with Reader(COG_INVERTED) as src:
             _ = src.tile(0, 0, 0)
+
+
+def test_int16_colormap():
+    """Should raise a warning about invalid data type for applying colormap.
+
+    ref: https://github.com/developmentseed/titiler/issues/1023
+    """
+    data = os.path.join(PREFIX, "cog_int16.tif")
+    color_map = cmap.get("viridis")
+
+    with Reader(data) as src:
+        info = src.info()
+        assert info.dtype == "int16"
+        assert info.nodata_type == "Nodata"
+        assert info.nodata_value == -32768
+
+        img = src.preview()
+        assert img.mask.any()
+
+        with pytest.warns(UserWarning):
+            im = img.apply_colormap(color_map)
+
+            # make sure we keep the nodata part masked
+            assert not im.mask.all()