Raster-service

app.py

import logging
import time

import mercantile
import io
import numpy as np
from shapely.ops import transform
from pyproj.transformer import Transformer

from shapely.geometry import box

from flask import Flask, send_file

from gdal_boots import RasterDataset, PNG


app = Flask(__name__)
PROJ_4326_3857 = Transformer.from_crs("EPSG:4326", "EPSG:3857", always_xy=True)

ds = RasterDataset.open("input.tiff")

band = ds.ds.GetRasterBand(1)
ct = band.GetRasterColorTable()
ct_size = ct.GetCount()
lookup = [
    np.arange(ct_size),
    np.arange(ct_size),
    np.arange(ct_size),
    np.ones(ct_size) * 255,
]
for i in range(ct_size):
    entry = ct.GetColorEntry(i)
    for c in range(4):
        lookup[c][i] = entry[c]

logging.basicConfig(level=logging.INFO)


@app.route("/<int:z>/<int:x>/<int:y>.png")
def raster(z: int, x: int, y: int):
    ts = time.time()
    box_shape_4326 = box(*mercantile.bounds(x, y, z))
    box_shape_3857 = transform(PROJ_4326_3857.transform, box_shape_4326)
    bbox = np.array(box_shape_3857.bounds).reshape(2,2)

    try:
        ds_out = ds.warp(
            bbox.reshape(-1).tolist(),
            width=256,
            height=256,
            bbox_epsg=3857
        )
    except Exception as e:
        logging.exception(f"fail to handle {z}/{x}/{y}")
        ds_out = None

    # colorize
    ts = time.time()
    ds_res = RasterDataset.create((4, 256, 256), np.uint8)
    if ds_out:
        ds_res[3, :] = (ds_out[:] < 14) * 255
        for band_idx in range(3):
            ds_res[band_idx, :, :] = np.take(lookup[band_idx], ds_out[:])

    data = ds_res.to_bytes(PNG())
    logging.info("generate %s png within %s", (z, x, y), round(time.time() - ts, 5))
    return send_file(io.BytesIO(data), mimetype='image/png')

if __name__ == "__main__":
    import bjoern
    host, port = "127.0.0.1", 6000
    logging.info(f"starting at {host}:{port}")
    try:
        bjoern.run(app, host, port)
    except KeyboardInterrupt:
        pass