sketch for raw raster support in pcfuncs

pcfuncs/funclib/models.py

@@ -1,7 +1,13 @@
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 from urllib.parse import quote
 
+import attr
+import numpy
 from pydantic import BaseModel
+from rasterio.enums import ColorInterp, Resampling
+from rasterio.io import DatasetReader, DatasetWriter, MemoryFile
+from rio_tiler.models import ImageData
+from rio_tiler.utils import get_array_statistics, non_alpha_indexes, resize_array
 
 
 class RenderOptions(BaseModel):
@@ -104,3 +110,137 @@ def from_query_params(cls, render_params: str) -> "RenderOptions":
                 result[k] = v
 
         return RenderOptions(**result)
+
+
+@attr.s
+class RIOImage(ImageData):
+    @property
+    def size(self) -> Tuple[int, int]:
+        return (self.width, self.height)
+
+    def paste(
+        self,
+        img: "RIOImage",
+        box: Optional[
+            Union[
+                Tuple[int, int],
+                Tuple[int, int, int, int],
+            ]
+        ],
+    ) -> None:
+        if img.count != self.count:
+            raise Exception("Cannot merge 2 images with different band number")
+
+        if img.data.dtype != self.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 box:
+            if len(box) == 2:
+                minx, maxy = box  # type: ignore
+                self.data[:, maxy:, minx:] = img.data
+                self.mask[maxy:, minx:] = img.mask
+
+            elif len(box) == 4:
+                # TODO add more size tests
+                minx, maxy, maxy, miny = box  # type: ignore
+                self.data[:, maxy:miny, minx:maxy] = img.data
+                self.mask[maxy:miny, minx:maxy] = img.mask
+
+            else:
+                raise Exception("Invalid box format")
+
+    def crop(self, bbox: Tuple[int, int, int, int]) -> "RIOImage":
+        """From rio-tiler 4.0"""
+        col_min, row_min, col_max, row_max = bbox
+
+        data = self.data[:, row_min:row_max, col_min:col_max]
+        mask = self.mask[row_min:row_max, col_min:col_max]
+
+        return RIOImage(  # type: ignore
+            data,
+            mask,
+            assets=self.assets,
+            crs=self.crs,
+            bounds=bbox,
+            band_names=self.band_names,
+            metadata=self.metadata,
+            dataset_statistics=self.dataset_statistics,
+        )
+
+    # This is slightly different from the resize method in rio-tiler 4.0
+    def resize(  # type: ignore
+        self,
+        size: Tuple[int, int],
+        resampling_method: Resampling = "nearest",
+    ) -> "RIOImage":
+        """From rio-tiler 4.0"""
+        width, height = size
+
+        data = resize_array(self.data, height, width, resampling_method)
+        mask = resize_array(self.mask, height, width, resampling_method)
+
+        return RIOImage(  # type: ignore
+            data,
+            mask,
+            assets=self.assets,
+            crs=self.crs,
+            bounds=self.bounds,
+            band_names=self.band_names,
+            metadata=self.metadata,
+            dataset_statistics=self.dataset_statistics,
+        )
+
+    @classmethod
+    def from_rio(
+        cls, dataset: Union[DatasetReader, DatasetWriter, MemoryFile]
+    ) -> "RIOImage":
+        indexes = non_alpha_indexes(dataset)
+
+        if ColorInterp.alpha in dataset.colorinterp:
+            # If dataset has an alpha band we need to get the mask using
+            # the alpha band index and then split the data and mask values
+            alpha_idx = dataset.colorinterp.index(ColorInterp.alpha) + 1
+            idx = tuple(indexes) + (alpha_idx,)
+            data = dataset.read(indexes=idx)
+            data, mask = data[0:-1], data[-1].astype("uint8")
+
+        else:
+            data = dataset.read(indexes=indexes)
+            mask = dataset.dataset_mask()
+
+        return cls(  # type: ignore
+            data,
+            mask,
+            crs=dataset.crs,
+            bounds=dataset.bounds,
+        )
+
+    def statistics(
+        self,
+        categorical: bool = False,
+        categories: Optional[List[float]] = None,
+        percentiles: List[int] = [2, 98],
+        hist_options: Optional[Dict] = None,
+    ) -> Dict:
+        data = numpy.ma.array(self.data)
+        data.mask = self.mask == 0
+
+        hist_options = hist_options or {}
+
+        stats = get_array_statistics(
+            data,
+            categorical=categorical,
+            categories=categories,
+            percentiles=percentiles,
+            **hist_options,
+        )
+
+        return {f"{self.band_names[ix]}": stats[ix] for ix in range(len(stats))}

pcfuncs/funclib/raster.py

@@ -6,6 +6,7 @@
 from typing import Any, Dict, List, Optional, Tuple, TypeVar
 
 import mercantile
+from funclib.models import RIOImage
 from PIL.Image import Image as PILImage
 from pyproj import CRS, Transformer
 
@@ -170,5 +171,45 @@ def mask(self, geom: Dict[str, Any]) -> "PILRaster":
 
 
 class GDALRaster(Raster):
-    # TODO: Implement
-    ...
+    def __init__(self, extent: RasterExtent, image: RIOImage) -> 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)
+
+        box: Any = (col_min, row_min, col_max, row_max)
+        cropped = self.image.crop(box)
+        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((cols, rows)),
+        )
+
+    def mask(self, geom: Dict[str, Any]) -> "GDALRaster":
+        raise NotImplementedError("GDALRaster does not support masking")

pcfuncs/funclib/tiles.py

@@ -8,7 +8,8 @@
 
 import aiohttp
 import mercantile
-from funclib.models import RenderOptions
+import numpy
+from funclib.models import RenderOptions, RIOImage
 from funclib.raster import (
     Bbox,
     ExportFormats,
@@ -20,6 +21,7 @@
 from funclib.settings import BaseExporterSettings
 from mercantile import Tile
 from PIL import Image
+from rasterio.io import MemoryFile
 
 from pccommon.backoff import BackoffStrategy, with_backoff_async
 
@@ -152,8 +154,82 @@ async def create(
 
 
 class GDALTileSet(TileSet[GDALRaster]):
+    async def _get_tile(self, url: str) -> Union[RIOImage, None]:
+        async def _f() -> RIOImage:
+            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:
+                            with MemoryFile(io.BytesIO(await resp.read())) as mem:
+                                with mem.open() as src:
+                                    return RIOImage.from_rio(src)
+                        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[RIOImage, 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[RIOImage, None]] = list(await asyncio.gather(*tasks))
+
+        tileset_dimensions = get_tileset_dimensions(tiles, self.tile_size)
+
+        # Get Count / datatype from the first tile_images
+        count: int
+        dtype: str
+        for im in tile_images:
+            if im:
+                count = im.count
+                dtype = im.data.dtype
+                break
+
+        mosaic = RIOImage(  # 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
+
+            mosaic.paste(tile, (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.txt

@@ -14,6 +14,7 @@ pillow==9.3.0
 pyproj==3.3.1
 pydantic>=1.9,<2.0.0
 rasterio==1.3.*
+rio-tiler==3.1.*  # same as titiler 0.7.0
 
 # Deployment needs to copy the local code into
 # the app code directory, so requires a separate

pcfuncs/tests/funclib/test_raster.py

@@ -1,6 +1,8 @@
 from pathlib import Path
 
-from funclib.raster import Bbox, PILRaster, RasterExtent
+import rasterio
+from funclib.models import RIOImage
+from funclib.raster import Bbox, GDALRaster, PILRaster, RasterExtent
 from PIL import Image
 
 HERE = Path(__file__).parent
@@ -35,3 +37,22 @@ def test_raster_extent_map_to_grid() -> None:
 
     assert x == 5
     assert y == 5
+
+
+def test_rio_crop() -> None:
+    with rasterio.open(DATA_FILES / "s2.png") as src:
+        img = RIOImage.from_rio(src)
+
+    raster = GDALRaster(
+        extent=RasterExtent(
+            bbox=Bbox(0, 0, 10, 10),
+            cols=img.size[0],
+            rows=img.size[1],
+        ),
+        image=img,
+    )
+
+    cropped = raster.crop(Bbox(0, 0, 5, 5))
+
+    assert abs(cropped.extent.cols - (img.size[0] / 2)) < 1
+    assert abs(cropped.extent.rows - (img.size[1] / 2)) < 1