sketch for raw raster support in pcfuncs

docker-compose.dev.yml

@@ -117,3 +117,19 @@ services:
     depends_on:
       - stac
       - tiler
+
+  funcs-dev:
+    image: pc-apis-funcs-dev
+    platform: linux/amd64
+    build:
+      context: .
+      dockerfile: pcfuncs/Dockerfile.dev
+    env_file: ${PC_FUNCS_ENV_FILE:-./pc-funcs.dev.env}
+    ports:
+      - "8083:80"
+    volumes:
+      - .:/opt/src
+    command: >
+      /bin/bash
+    depends_on:
+      - funcs

pcfuncs/Dockerfile.dev

@@ -0,0 +1,4 @@
+FROM pc-apis-funcs
+
+COPY pcfuncs/requirements-dev.txt requirements-dev.txt
+RUN pip install -r requirements-dev.txt

pcfuncs/funclib/raster.py

@@ -8,6 +8,7 @@
 import mercantile
 from PIL.Image import Image as PILImage
 from pyproj import CRS, Transformer
+from rio_tiler.models import ImageData
 
 T = TypeVar("T", bound="Raster")
 
@@ -170,5 +171,56 @@ def mask(self, geom: Dict[str, Any]) -> "PILRaster":
 
 
 class GDALRaster(Raster):
-    # TODO: Implement
-    ...
+    def __init__(self, extent: RasterExtent, image: ImageData) -> None:
+        self.image = image
+        super().__init__(extent)
+
+    def to_bytes(self, format: str = ExportFormats.PNG) -> io.BytesIO:
+        img_bytes = self.image.render(
+            add_mask=True,
+            img_format=format.upper(),
+        )
+        return io.BytesIO(img_bytes)
+
+    def crop(self, bbox: Bbox) -> "GDALRaster":
+        # Web mercator of user bbox
+        if (
+            not bbox.crs == self.extent.bbox.crs
+            and bbox.crs is not None
+            and self.extent.bbox.crs is not None
+        ):
+            bbox = bbox.reproject(self.extent.bbox.crs)
+
+        col_min, row_min = self.extent.map_to_grid(bbox.xmin, bbox.ymax)
+        col_max, row_max = self.extent.map_to_grid(bbox.xmax, bbox.ymin)
+
+        data = self.image.data[:, row_min:row_max, col_min:col_max]
+        mask = self.image.mask[row_min:row_max, col_min:col_max]
+        cropped = ImageData(
+            data,
+            mask,
+            assets=self.image.assets,
+            crs=self.image.crs,
+            bounds=bbox.to_list(),
+            band_names=self.image.band_names,
+            metadata=self.image.metadata,
+            dataset_statistics=self.image.dataset_statistics,
+        )
+
+        return GDALRaster(
+            extent=RasterExtent(
+                bbox,
+                cols=col_max - col_min,
+                rows=row_max - row_min,
+            ),
+            image=cropped,
+        )
+
+    def resample(self, cols: int, rows: int) -> "GDALRaster":
+        return GDALRaster(
+            extent=RasterExtent(bbox=self.extent.bbox, cols=cols, rows=rows),
+            image=self.image.resize(rows, cols),  # type: ignore
+        )
+
+    def mask(self, geom: Dict[str, Any]) -> "GDALRaster":
+        raise NotImplementedError("GDALRaster does not support masking")

pcfuncs/funclib/tiles.py

@@ -3,11 +3,12 @@
 import logging
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Any, Dict, Generic, List, Optional, Type, TypeVar, Union
+from typing import Any, Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union
 from urllib.parse import urlsplit
 
 import aiohttp
 import mercantile
+import numpy
 from funclib.models import RenderOptions
 from funclib.raster import (
     Bbox,
@@ -20,6 +21,7 @@
 from funclib.settings import BaseExporterSettings
 from mercantile import Tile
 from PIL import Image
+from rio_tiler.models import ImageData
 
 from pccommon.backoff import BackoffStrategy, with_backoff_async
 
@@ -57,6 +59,46 @@ def get_tileset_dimensions(tiles: List[Tile], tile_size: int) -> TileSetDimensio
     )
 
 
+def paste_into_image(
+    from_image: ImageData,
+    to_image: ImageData,
+    box: Optional[
+        Union[
+            Tuple[int, int],
+            Tuple[int, int, int, int],
+        ]
+    ],
+) -> None:
+    """Paste image data from one to the other."""
+    if from_image.count != to_image.count:
+        raise Exception("Cannot merge 2 images with different band number")
+
+    if from_image.data.dtype != to_image.data.dtype:
+        raise Exception("Cannot merge 2 images with different datatype")
+
+    # Pastes another image into this image.
+    # The box argument is either a 2-tuple giving the upper left corner,
+    # a 4-tuple defining the left, upper, right, and lower pixel coordinate,
+    # or None (same as (0, 0)). See Coordinate System. If a 4-tuple is given,
+    # the size of the pasted image must match the size of the region.
+    if box is None:
+        box = (0, 0)
+
+    if len(box) == 2:
+        size = (from_image.width, from_image.height)
+        box += (box[0] + size[0], box[1] + size[1])  # type: ignore
+        minx, maxy, maxx, miny = box  # type: ignore
+    elif len(box) == 4:
+        # TODO add more size tests
+        minx, maxy, maxx, miny = box  # type: ignore
+
+    else:
+        raise Exception("Invalid box format")
+
+    to_image.data[:, maxy:miny, minx:maxx] = from_image.data
+    to_image.mask[maxy:miny, minx:maxx] = from_image.mask
+
+
 class TileSet(ABC, Generic[T]):
     def __init__(
         self,
@@ -152,8 +194,89 @@ async def create(
 
 
 class GDALTileSet(TileSet[GDALRaster]):
+    async def _get_tile(
+        self,
+        url: str,
+    ) -> Union[ImageData, None]:
+        async def _f() -> ImageData:
+            async with aiohttp.ClientSession() as session:
+                async with self._async_limit:
+                    # We set Accept-Encoding to make sure the response is compressed
+                    async with session.get(
+                        url, headers={"Accept-Encoding": "gzip"}
+                    ) as resp:
+                        if resp.status == 200:
+                            return ImageData.from_bytes(await resp.read())
+
+                        else:
+                            raise TilerError(
+                                f"Error downloading tile: {url}", resp=resp
+                            )
+
+        try:
+            return await with_backoff_async(
+                _f,
+                is_throttle=lambda e: isinstance(e, TilerError),
+                strategy=BackoffStrategy(waits=[0.2, 0.5, 0.75, 1, 2]),
+            )
+        except Exception:
+            logger.warning(f"Tile request failed with backoff: {url}")
+            return None
+
     async def get_mosaic(self, tiles: List[Tile]) -> GDALRaster:
-        raise NotImplementedError()
+        tasks: List[asyncio.Future[Union[ImageData, None]]] = []
+        for tile in tiles:
+            url = self.get_tile_url(tile.z, tile.x, tile.y)
+            print(f"Downloading {url}")
+            tasks.append(asyncio.ensure_future(self._get_tile(url)))
+
+        tile_images: List[Union[ImageData, None]] = list(await asyncio.gather(*tasks))
+
+        tileset_dimensions = get_tileset_dimensions(tiles, self.tile_size)
+
+        # By default if no tiles where return we create an
+        # empty mosaic with 3 bands and uint8
+        count: int = 3
+        dtype: str = "uint8"
+        for im in tile_images:
+            if im:
+                count = im.count
+                dtype = im.data.dtype
+                break  # Get Count / datatype from the first valid tile_images
+
+        mosaic = ImageData(  # type: ignore
+            numpy.zeros(
+                (count, tileset_dimensions.total_rows, tileset_dimensions.total_cols),
+                dtype=dtype,
+            )
+        )
+
+        x = 0
+        y = 0
+        for i, img in enumerate(tile_images):
+            if not img:
+                continue
+
+            paste_into_image(
+                img,
+                mosaic,
+                (x * self.tile_size, y * self.tile_size),
+            )
+
+            # Increment the row/col position for subsequent tiles
+            if (i + 1) % tileset_dimensions.tile_rows == 0:
+                y = 0
+                x += 1
+            else:
+                y += 1
+
+        raster_extent = RasterExtent(
+            bbox=Bbox.from_tiles(tiles),
+            cols=tileset_dimensions.total_cols,
+            rows=tileset_dimensions.total_rows,
+        )
+
+        return GDALRaster(raster_extent, mosaic)
 
 
 class PILTileSet(TileSet[PILRaster]):

pcfuncs/requirements-dev.txt

@@ -0,0 +1 @@
+uvicorn

pcfuncs/requirements.txt

@@ -14,6 +14,7 @@ pillow==9.3.0
 pyproj==3.3.1
 pydantic>=1.9,<2.0.0
 rasterio==1.3.*
+rio-tiler>=4.1.7  # to make sure we have ImageData.from_bytes
 
 # Deployment needs to copy the local code into
 # the app code directory, so requires a separate

pcfuncs/statistics/__init__.py

@@ -0,0 +1,70 @@
+import json
+import logging
+from typing import Dict
+
+import azure.functions as func
+from funclib.errors import BBoxTooLargeError
+from pydantic import ValidationError
+
+from .models import StatisticsRequest
+from .settings import StatisticsSettings
+from .statistics import PcMosaicImage
+
+
+async def main(req: func.HttpRequest) -> func.HttpResponse:
+    try:
+        body = req.get_json()
+    except ValueError:
+        return func.HttpResponse(
+            status_code=400,
+            mimetype="application/text",
+            body="Error: Invalid JSON",
+        )
+
+    try:
+        parsed_request = StatisticsRequest(**body)
+    except ValidationError as e:
+        return func.HttpResponse(
+            status_code=400,
+            mimetype="application/json",
+            body=e.json(),
+        )
+
+    try:
+        response = await handle_request(parsed_request)
+
+        return func.HttpResponse(
+            status_code=200,
+            mimetype="application/json",
+            body=json.dumps(response),
+        )
+    except BBoxTooLargeError as e:
+        logging.exception(e)
+        return func.HttpResponse(
+            status_code=400,
+            mimetype="application/json",
+            body=json.dumps({"error": str(e)}),
+        )
+    except Exception as e:
+        logging.exception(e)
+        return func.HttpResponse(
+            status_code=500,
+            mimetype="application/json",
+        )
+
+
+async def handle_request(req: StatisticsRequest) -> Dict:
+    settings = StatisticsSettings.get()
+
+    mosaic_image = PcMosaicImage(
+        bbox=req.bbox,
+        zoom=req.zoom,
+        cql=req.cql,
+        render_options=req.get_render_options(),
+        settings=settings,
+        data_api_url_override=req.data_api_url,
+    )
+
+    img = await mosaic_image.get()
+
+    return {k: v.dict() for k, v in img.statistics().items()}

pcfuncs/statistics/constants.py

@@ -0,0 +1,6 @@
+STATISTICS_SETTINGS_PREFIX = "STATISTICS_"
+
+DEFAULT_STATISTICS_CONTAINER_URL = "https://pcexplorer.blob.core.windows.net/statistics"
+MAX_TILE_COUNT = 16
+
+DEFAULT_CONCURRENCY = 10

pcfuncs/statistics/function.json

@@ -0,0 +1,17 @@
+{
+  "scriptFile": "__init__.py",
+  "bindings": [
+    {
+      "authLevel": "anonymous",
+      "type": "httpTrigger",
+      "direction": "in",
+      "name": "req",
+      "methods": ["post"]
+    },
+    {
+      "type": "http",
+      "direction": "out",
+      "name": "$return"
+    }
+  ]
+}

pcfuncs/statistics/models.py

@@ -0,0 +1,36 @@
+from typing import Any, Dict, List, Optional
+
+from funclib.models import RenderOptions
+from pydantic import BaseModel, validator
+
+
+def _get_render_options(render_params: str) -> Dict[str, List[str]]:
+    result: Dict[str, List[str]] = {}
+    for p in render_params.split("&"):
+        k, v = p.split("=")
+        if k not in result:
+            result[k] = []
+        result[k].append(v)
+    return result
+
+
+class StatisticsRequest(BaseModel):
+    bbox: List[float]
+    zoom: int
+    cql: Dict[str, Any]
+    render_params: str
+
+    data_api_url: Optional[str] = None
+    """Override for the data API URL. Useful for testing."""
+
+    @validator("render_params")
+    def _validate_render_params(cls, v: str) -> str:
+        RenderOptions.from_query_params(v)
+        return v
+
+    def get_render_options(self) -> RenderOptions:
+        return RenderOptions.from_query_params(self.render_params)
+
+    def get_collection(self) -> str:
+        render_options = _get_render_options(self.render_params)
+        return render_options["collection"][0]