add terrainrgb and terrarium examples

Author: vincentsarago

rio_viz/algorithm/__init__.py

@@ -3,11 +3,13 @@
 from typing import Dict, Type
 
 from rio_viz.algorithm.base import AlgorithmMetadata, BaseAlgorithm  # noqa
-from rio_viz.algorithm.dem import Contours, HillShade
+from rio_viz.algorithm.dem import Contours, HillShade, TerrainRGB, Terrarium
 from rio_viz.algorithm.index import NormalizedIndex
 
 AVAILABLE_ALGORITHM: Dict[str, Type[BaseAlgorithm]] = {
     "hillshade": HillShade,
     "contours": Contours,
     "normalizedIndex": NormalizedIndex,
+    "terrarium": Terrarium,
+    "terrainrgb": TerrainRGB,
 }

rio_viz/algorithm/dem.py

@@ -81,3 +81,83 @@ def apply(self, img: ImageData) -> ImageData:
             crs=img.crs,
             bounds=img.bounds,
         )
+
+
+class Terrarium(BaseAlgorithm):
+    """Encode DEM into RGB (Mapzen Terrarium)."""
+
+    input_nbands: int = 1
+
+    output_nbands: int = 3
+    output_dtype: str = "uint8"
+
+    def apply(self, img: ImageData) -> ImageData:
+        """Encode DEM into RGB."""
+        data = numpy.clip(img.data[0] + 32768.0, 0.0, 65535.0)
+        r = data / 256
+        g = data % 256
+        b = (data * 256) % 256
+        arr = numpy.stack([r, g, b]).astype(numpy.uint8)
+        print(arr.shape)
+        return ImageData(
+            arr,
+            img.mask,
+            assets=img.assets,
+            crs=img.crs,
+            bounds=img.bounds,
+        )
+
+
+class TerrainRGB(BaseAlgorithm):
+    """Encode DEM into RGB (Mapbox Terrain RGB)."""
+
+    interval: int = 1
+    baseval: int = -10000
+
+    input_nbands: int = 1
+
+    output_nbands: int = 3
+    output_dtype: str = "uint8"
+
+    def apply(self, img: ImageData) -> ImageData:
+        """Encode DEM into RGB (Mapbox Terrain RGB).
+
+        Code from https://github.com/mapbox/rio-rgbify/blob/master/rio_rgbify/encoders.py (MIT)
+
+        """
+
+        def _range_check(datarange):
+            """
+            Utility to check if data range is outside of precision for 3 digit base 256
+            """
+            maxrange = 256**3
+
+            return datarange > maxrange
+
+        round_digits = 0
+
+        data = img.data[0].astype(numpy.float64)
+        data -= self.baseval
+        data /= self.interval
+
+        data = numpy.around(data / 2**round_digits) * 2**round_digits
+
+        rows, cols = data.shape
+        datarange = data.max() - data.min()
+        if _range_check(datarange):
+            raise ValueError("Data of {} larger than 256 ** 3".format(datarange))
+
+        rgb = numpy.zeros((3, rows, cols), dtype=numpy.uint8)
+        rgb[2] = ((data / 256) - (data // 256)) * 256
+        rgb[1] = (((data // 256) / 256) - ((data // 256) // 256)) * 256
+        rgb[0] = (
+            (((data // 256) // 256) / 256) - (((data // 256) // 256) // 256)
+        ) * 256
+
+        return ImageData(
+            rgb,
+            img.mask,
+            assets=img.assets,
+            crs=img.crs,
+            bounds=img.bounds,
+        )