add cache to bbox, add mask to bbox

docker-compose.yml

@@ -42,6 +42,7 @@ services:
     environment:
       - TITILER_API_ROOT_PATH=/tiler
       - TITILER_API_CACHECONTROL=no-cache
+      - REDIS_HOST=redis
     command: ["pdm", "run", "uvicorn", "src.app.app:app", "--host", "0.0.0.0", "--port", "8080", "--reload"]
 
   jupyter:
@@ -51,3 +52,7 @@ services:
       - 8888:8888
     volumes:
       - ./titiler/src:/app/src
+
+  redis:
+    image: redis
+

titiler/pyproject.toml

@@ -17,6 +17,7 @@ dependencies = [
     "jupyter>=1.0.0",
     "notebook>=7.0.7",
     "jupyterlab>=4.1.0",
+    "walrus>=0.9.3",
 ]
 requires-python = ">=3.10"
 readme = "README.md"

titiler/src/app/dependencies.py

@@ -8,6 +8,7 @@
 
 import cv2
 import numpy as np
+import os
 
 from titiler.core.algorithm import BaseAlgorithm
 from rio_tiler.models import ImageData
@@ -16,6 +17,14 @@
 
 from titiler.core.algorithm import algorithms as default_algorithms
 from titiler.core.algorithm import Algorithms
+import walrus
+
+db = walrus.Database(
+    host=os.getenv("REDIS_HOST", "localhost"), port=int(os.getenv("REDIS_PORT", 6379))
+)
+cache = db.cache()
+CACHE_TIMEOUT = int(os.getenv("BBOX_CACHE_TIMEOUT", 60))
+BBOX_SCALE = int(os.getenv("BBOX_SCALE", 6))
 
 
 class StravaHeatmap(BaseAlgorithm):
@@ -84,8 +93,15 @@ def __call__(self, img: ImageData) -> ImageData:
             bounds=img.bounds,
         )
 
-# TODO: use a caching system
-ASSETS = {}
+
+@cache.cached(timeout=CACHE_TIMEOUT)
+def get_stats_by_bbox(url, bbox):
+    with Reader(url) as dst:
+        cov = dst.part(bbox)
+        cov_stats = cov.statistics()
+        bs = cov_stats.get('b1')
+        return ((bs.min, bs.max),)
+
 
 class BBoxStats(BaseAlgorithm):
     input_nbands: int = 1
@@ -96,29 +112,29 @@ class BBoxStats(BaseAlgorithm):
     scale: int = 1
 
     def __call__(self, img: ImageData) -> ImageData:
-        if self.scale > 8:
-            index = img.assets[0] + ','.join([str(v) for v in self.bbox])
-            if index in ASSETS:
-                bstats = ASSETS[index]
-            else:
-                with Reader(img.assets[0]) as dst:
-                    cov = dst.part(self.bbox)
-                    stats = cov.statistics()
-                    bs = stats.get('b1')
-                    bstats = ((bs.min, bs.max),)
-                    ASSETS[index] = bstats
-
-            img.rescale(
-                in_range=bstats,
-                out_range=((0, 255),)
-            )
+        # compute the mask, find all the nan values
+        mask = np.isnan(img.data)[0]
+        # generate an array mask, with 0 and 255
+        modified_mask = np.where(mask, 255, 0)
+
+        if self.scale > BBOX_SCALE:
+            stats = get_stats_by_bbox(img.assets[0], self.bbox)
         else:
             stats = img.dataset_statistics
-            img.rescale(
-                in_range=stats,
-                out_range=((0, 255),)
-            )
-        return img
+
+        img.rescale(
+            in_range=stats,
+            out_range=((0, 255),)
+        )
+        data = np.where(~mask, img.data, 0)
+        masked_array = np.ma.MaskedArray(data, mask=modified_mask)
+
+        return ImageData(
+            masked_array,
+            assets=img.assets,
+            crs=img.crs,
+            bounds=img.bounds,
+        )
 
 
 algorithms: Algorithms = default_algorithms.register({