sketch algorithm support

Author: vincentsarago

pyproject.toml

@@ -20,6 +20,7 @@ dynamic = ["version"]
 dependencies = [
     "braceexpand",
     "rio-cogeo>=3.1",
+    "rio-tiler>=3.1.5",
     "titiler.core>=0.5,<0.8",
     "starlette-cramjam>=0.3,<0.4",
     "uvicorn",

rio_viz/algorithm/__init__.py

@@ -0,0 +1,13 @@
+"""rio_viz.algorithm."""
+
+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.index import NormalizedIndex
+
+AVAILABLE_ALGORITHM: Dict[str, Type[BaseAlgorithm]] = {
+    "hillshade": HillShade,
+    "contours": Contours,
+    "normalizedIndex": NormalizedIndex,
+}

rio_viz/algorithm/base.py

@@ -0,0 +1,37 @@
+"""Algorithm base class."""
+
+import abc
+from typing import Dict, Optional, Sequence
+
+from pydantic import BaseModel
+from rio_tiler.models import ImageData
+
+
+class BaseAlgorithm(BaseModel, metaclass=abc.ABCMeta):
+    """Algorithm baseclass."""
+
+    input_nbands: int
+
+    output_nbands: int
+    output_dtype: str
+    output_min: Optional[Sequence]
+    output_max: Optional[Sequence]
+
+    @abc.abstractmethod
+    def apply(self, img: ImageData) -> ImageData:
+        """Apply"""
+        ...
+
+    class Config:
+        """Config for model."""
+
+        extra = "allow"
+
+
+class AlgorithmMetadata(BaseModel):
+    """Algorithm metadata."""
+
+    name: str
+    inputs: Dict
+    outputs: Dict
+    params: Dict

rio_viz/algorithm/dem.py

@@ -0,0 +1,83 @@
+"""rio_viz.algorithm DEM."""
+
+import numpy
+from rio_tiler.colormap import apply_cmap, cmap
+from rio_tiler.models import ImageData
+from rio_tiler.utils import linear_rescale
+
+from rio_viz.algorithm.base import BaseAlgorithm
+
+
+class HillShade(BaseAlgorithm):
+    """Hillshade."""
+
+    azimuth: int = 90
+    angle_altitude: float = 90
+
+    input_nbands: int = 1
+
+    output_nbands: int = 1
+    output_dtype: str = "uint8"
+
+    def apply(self, img: ImageData) -> ImageData:
+        """Create hillshade from DEM dataset."""
+        data = img.data[0]
+        mask = img.mask
+
+        x, y = numpy.gradient(data)
+
+        slope = numpy.pi / 2.0 - numpy.arctan(numpy.sqrt(x * x + y * y))
+        aspect = numpy.arctan2(-x, y)
+        azimuthrad = self.azimuth * numpy.pi / 180.0
+        altituderad = self.angle_altitude * numpy.pi / 180.0
+        shaded = numpy.sin(altituderad) * numpy.sin(slope) + numpy.cos(
+            altituderad
+        ) * numpy.cos(slope) * numpy.cos(azimuthrad - aspect)
+        hillshade_array = 255 * (shaded + 1) / 2
+
+        # ImageData only accept image in form of (count, height, width)
+        arr = numpy.expand_dims(hillshade_array, axis=0).astype(dtype=numpy.uint8)
+
+        return ImageData(
+            arr,
+            mask,
+            assets=img.assets,
+            crs=img.crs,
+            bounds=img.bounds,
+        )
+
+
+class Contours(BaseAlgorithm):
+    """Contours.
+
+    Original idea from https://custom-scripts.sentinel-hub.com/dem/contour-lines/
+    """
+
+    increment: int = 35
+    thickness: int = 1
+    minz: int = -12000
+    maxz: int = 8000
+
+    input_nbands: int = 1
+
+    output_nbands: int = 3
+    output_dtype: str = "uint8"
+
+    def apply(self, img: ImageData) -> ImageData:
+        """Add contours."""
+        data = img.data
+
+        # Apply rescaling for minz,maxz to 1->255 and apply Terrain colormap
+        arr = linear_rescale(data, (self.minz, self.maxz), (1, 255)).astype("uint8")
+        arr, _ = apply_cmap(arr, cmap.get("terrain"))
+
+        # set black (0) for contour lines
+        arr = numpy.where(data % self.increment < self.thickness, 0, arr)
+
+        return ImageData(
+            arr,
+            img.mask,
+            assets=img.assets,
+            crs=img.crs,
+            bounds=img.bounds,
+        )

rio_viz/algorithm/index.py

@@ -0,0 +1,37 @@
+"""rio_viz.algorithm Normalized Index."""
+
+from typing import Sequence
+
+import numpy
+from rio_tiler.models import ImageData
+
+from rio_viz.algorithm.base import BaseAlgorithm
+
+
+class NormalizedIndex(BaseAlgorithm):
+    """Normalized Difference Index."""
+
+    input_nbands: int = 2
+
+    output_nbands: int = 1
+    output_dtype: str = "float32"
+    output_min: Sequence[float] = [-1.0]
+    output_max: Sequence[float] = [1.0]
+
+    def apply(self, img: ImageData) -> ImageData:
+        """Normalized difference."""
+        b1 = img.data[0]
+        b2 = img.data[1]
+
+        arr = numpy.where(img.mask, (b2 - b1) / (b2 + b1), 0)
+
+        # ImageData only accept image in form of (count, height, width)
+        arr = numpy.expand_dims(arr, axis=0).astype(self.output_dtype)
+
+        return ImageData(
+            arr,
+            img.mask,
+            assets=img.assets,
+            crs=img.crs,
+            bounds=img.bounds,
+        )

rio_viz/app.py

@@ -22,6 +22,8 @@
 from starlette.types import ASGIApp
 from starlette_cramjam.middleware import CompressionMiddleware
 
+from rio_viz.algorithm import AVAILABLE_ALGORITHM, AlgorithmMetadata
+from rio_viz.dependency import PostProcessParams
 from rio_viz.resources.enums import RasterFormat, VectorTileFormat, VectorTileType
 
 from titiler.core.dependencies import (
@@ -36,7 +38,6 @@
     HistogramParams,
     ImageParams,
     ImageRenderingParams,
-    PostProcessParams,
     StatisticsParams,
 )
 from titiler.core.models.mapbox import TileJSON
@@ -295,19 +296,26 @@ async def preview(
                 # Adapt options for each reader type
                 self._update_params(src_dst, layer_params)
 
-                data = await src_dst.preview(
+                img = await src_dst.preview(
                     **layer_params,
                     **dataset_params,
                     **img_params,
                 )
                 dst_colormap = getattr(src_dst, "colormap", None)
 
-            if not format:
-                format = RasterFormat.jpeg if data.mask.all() else RasterFormat.png
+            if postprocess_params.image_process:
+                img = postprocess_params.image_process.apply(img)
+
+            if postprocess_params.rescale:
+                img.rescale(postprocess_params.rescale)
+
+            if postprocess_params.color_formula:
+                img.apply_color_formula(postprocess_params.color_formula)
 
-            image = data.post_process(**postprocess_params)
+            if not format:
+                format = RasterFormat.jpeg if img.mask.all() else RasterFormat.png
 
-            content = image.render(
+            content = img.render(
                 img_format=format.driver,
                 colormap=colormap or dst_colormap,
                 **format.profile,
@@ -360,17 +368,24 @@ async def part(
                 # Adapt options for each reader type
                 self._update_params(src_dst, layer_params)
 
-                data = await src_dst.part(
+                img = await src_dst.part(
                     [minx, miny, maxx, maxy],
                     **layer_params,
                     **dataset_params,
                     **img_params,
                 )
                 dst_colormap = getattr(src_dst, "colormap", None)
 
-            image = data.post_process(**postprocess_params)
+            if postprocess_params.image_process:
+                img = postprocess_params.image_process.apply(img)
+
+            if postprocess_params.rescale:
+                img.rescale(postprocess_params.rescale)
+
+            if postprocess_params.color_formula:
+                img.apply_color_formula(postprocess_params.color_formula)
 
-            content = image.render(
+            content = img.render(
                 img_format=format.driver,
                 colormap=colormap or dst_colormap,
                 **format.profile,
@@ -415,17 +430,24 @@ async def geojson_part(
                 # Adapt options for each reader type
                 self._update_params(src_dst, layer_params)
 
-                data = await src_dst.feature(
+                img = await src_dst.feature(
                     geom.dict(exclude_none=True), **layer_params, **dataset_params
                 )
                 dst_colormap = getattr(src_dst, "colormap", None)
 
-            if not format:
-                format = RasterFormat.jpeg if data.mask.all() else RasterFormat.png
+            if postprocess_params.image_process:
+                img = postprocess_params.image_process.apply(img)
+
+            if postprocess_params.rescale:
+                img.rescale(postprocess_params.rescale)
 
-            image = data.post_process(**postprocess_params)
+            if postprocess_params.color_formula:
+                img.apply_color_formula(postprocess_params.color_formula)
 
-            content = image.render(
+            if not format:
+                format = RasterFormat.jpeg if img.mask.all() else RasterFormat.png
+
+            content = img.render(
                 img_format=format.driver,
                 colormap=colormap or dst_colormap,
                 **format.profile,
@@ -475,7 +497,7 @@ async def tile(
                 # Adapt options for each reader type
                 self._update_params(src_dst, layer_params)
 
-                tile_data = await src_dst.tile(
+                img = await src_dst.tile(
                     x,
                     y,
                     z,
@@ -502,22 +524,27 @@ async def tile(
                 _mvt_encoder = partial(run_in_threadpool, pixels_encoder)
 
                 content = await _mvt_encoder(
-                    tile_data.data,
-                    tile_data.mask,
-                    tile_data.band_names,
+                    img.data,
+                    img.mask,
+                    img.band_names,
                     feature_type=feature_type.value,
                 )  # type: ignore
 
             # Raster Tile
             else:
-                if not format:
-                    format = (
-                        RasterFormat.jpeg if tile_data.mask.all() else RasterFormat.png
-                    )
+                if postprocess_params.image_process:
+                    img = postprocess_params.image_process.apply(img)
+
+                if postprocess_params.rescale:
+                    img.rescale(postprocess_params.rescale)
+
+                if postprocess_params.color_formula:
+                    img.apply_color_formula(postprocess_params.color_formula)
 
-                image = tile_data.post_process(**postprocess_params)
+                if not format:
+                    format = RasterFormat.jpeg if img.mask.all() else RasterFormat.png
 
-                content = image.render(
+                content = img.render(
                     img_format=format.driver,
                     colormap=colormap or dst_colormap,
                     **format.profile,
@@ -646,6 +673,35 @@ async def wmts(
                 media_type="application/xml",
             )
 
+        @self.router.get(
+            "/algorithm",
+            response_model=List[AlgorithmMetadata],
+        )
+        def algo(request: Request):
+            """Handle /algorithm."""
+            algos = []
+            for k, v in AVAILABLE_ALGORITHM.items():
+                props = v.schema()["properties"]
+                ins = {
+                    k.replace("input_", ""): v
+                    for k, v in props.items()
+                    if k.startswith("input_")
+                }
+                outs = {
+                    k.replace("output_", ""): v
+                    for k, v in props.items()
+                    if k.startswith("output_")
+                }
+                params = {
+                    k: v
+                    for k, v in props.items()
+                    if not k.startswith("input_") and not k.startswith("output_")
+                }
+                algos.append(
+                    AlgorithmMetadata(name=k, inputs=ins, outputs=outs, params=params)
+                )
+            return algos
+
         @self.router.get(
             "/",
             responses={200: {"description": "Simple COG viewer."}},