cog tiler with rio-tiler/fastapi

cog_tiler.py

import abc
from dataclasses import dataclass
from io import BytesIO
from typing import Tuple

import numpy as np

from cogeo_mosaic.backends import MosaicBackend
from rio_tiler.errors import TileOutsideBounds
from rio_tiler.mosaic import mosaic_reader
from rio_tiler_crs import COGReader

from titiler.api.utils import postprocess, reformat
from titiler.ressources.enums import ImageType

from fastapi import APIRouter, FastAPI, Path, Query
from starlette.responses import StreamingResponse


@dataclass
class Tiler(abc.ABC):
    @abc.abstractmethod
    def tile(self, asset: str, *args, **kwargs) -> Tuple[np.ndarray, np.ndarray]:
        ...

    @abc.abstractmethod
    def read_tile(
        self,
        x: int = Path(...),
        y: int = Path(...),
        z: int = Path(...),
        img_format: ImageType = ImageType.jpg,
        url: str = Query(...),
    ) -> StreamingResponse:
        ...


@dataclass
class CogTiler(Tiler):
    def tile(self, asset: str, *args, **kwargs) -> Tuple[np.ndarray, np.ndarray]:
        with COGReader(asset) as cog:
            return cog.tile(*args, **kwargs)

    def read_tile(
        self,
        x: int = Path(...),
        y: int = Path(...),
        z: int = Path(...),
        img_format: ImageType = Query(...),
        url: str = Query(...),
    ) -> StreamingResponse:
        tile, mask = self.tile(url, x, y, z)
        tile = postprocess(tile, mask)
        contents = reformat(tile, mask, img_format=img_format,)
        return StreamingResponse(BytesIO(contents))


@dataclass
class MosaicTiler(CogTiler):
    def read_tile(
        self,
        x: int = Path(...),
        y: int = Path(...),
        z: int = Path(...),
        img_format: ImageType = ImageType.jpg,
        url: str = Query(...),
    ) -> StreamingResponse:
        with MosaicBackend(url) as mosaic:
            assets = mosaic.tile(x=x, y=y, z=z)
            if not assets:
                raise TileOutsideBounds(f"No assets found for tile {z}/{x}/{y}")
        tile, mask = mosaic_reader(assets, self.tile, x, y, z)
        tile = postprocess(tile, mask,)
        contents = reformat(tile, mask, img_format=img_format,)
        return StreamingResponse(BytesIO(contents))


def create_tiler_router(tiler: Tiler) -> APIRouter:
    router = APIRouter()
    router.add_api_route(
        endpoint=tiler.read_tile,
        path="/tiles/{z}/{x}/{y}.{img_format}",
        methods=["GET"],
    )
    return router


def create_app(mosaics: bool = False) -> FastAPI:
    app = FastAPI()
    tile_client = CogTiler()
    app.include_router(create_tiler_router(tile_client), prefix="/cog")

    if mosaics:
        mosaic_client = MosaicTiler()
        app.include_router(create_tiler_router(mosaic_client), prefix="/mosaic")

    return app


app = create_app()