tilefetch 0.1.0__py3-none-any.whl

tilefetch/mosaic.py

@@ -0,0 +1,191 @@
+"""Stitch tiles into one image and resample it to an exact bbox / CRS."""
+
+from __future__ import annotations
+
+import io
+import math
+from typing import Callable, Iterable
+
+from PIL import Image
+
+from .tiles import EARTH_RADIUS, BBox, Tile, lonlat_to_mercator, tile_bounds_mercator
+
+__all__ = ["stitch", "crop_mercator", "warp_to_lonlat", "decode_tile", "RESAMPLING"]
+
+RESAMPLING = {
+    "nearest": Image.Resampling.NEAREST,
+    "bilinear": Image.Resampling.BILINEAR,
+    "bicubic": Image.Resampling.BICUBIC,
+}
+
+
+def _resampling(name: str) -> Image.Resampling:
+    try:
+        return RESAMPLING[name]
+    except KeyError:
+        raise ValueError(f"resampling must be one of {sorted(RESAMPLING)}") from None
+
+
+def decode_tile(data: bytes) -> Image.Image | None:
+    """Decode tile bytes to RGBA, or ``None`` if they are not an image."""
+    try:
+        with Image.open(io.BytesIO(data)) as im:
+            im.load()
+            return im.convert("RGBA")
+    except Exception:
+        return None
+
+
+def stitch(
+    tiles: Iterable[Tile],
+    get: Callable[[Tile], bytes | None],
+    tile_size: int = 256,
+) -> tuple[Image.Image, BBox, int]:
+    """Paste the tiles of one zoom level into an RGBA canvas.
+
+    Returns ``(image, bounds_3857, missing)``; missing or undecodable tiles
+    stay transparent.
+    """
+    tiles = list(tiles)
+    if not tiles:
+        raise ValueError("no tiles to stitch")
+    zs = {t.z for t in tiles}
+    if len(zs) != 1:
+        raise ValueError("stitch() needs tiles from a single zoom level")
+    xmin = min(t.x for t in tiles)
+    xmax = max(t.x for t in tiles)
+    ymin = min(t.y for t in tiles)
+    ymax = max(t.y for t in tiles)
+    decoded: dict[Tile, Image.Image] = {}
+    missing = 0
+    for t in tiles:
+        data = get(t)
+        im = decode_tile(data) if data else None
+        if im is None:
+            missing += 1
+        else:
+            decoded[t] = im
+    if decoded:
+        tile_size = next(iter(decoded.values())).width
+    canvas = Image.new(
+        "RGBA", ((xmax - xmin + 1) * tile_size, (ymax - ymin + 1) * tile_size), (0, 0, 0, 0)
+    )
+    for t, im in decoded.items():
+        if im.size != (tile_size, tile_size):
+            im = im.resize((tile_size, tile_size), Image.Resampling.BILINEAR)
+        canvas.paste(im, ((t.x - xmin) * tile_size, (t.y - ymin) * tile_size))
+    z = zs.pop()
+    tl = tile_bounds_mercator((z, xmin, ymin))
+    br = tile_bounds_mercator((z, xmax, ymax))
+    return canvas, BBox(tl.minx, br.miny, br.maxx, tl.maxy), missing
+
+
+def _clip(target: BBox, grid: BBox) -> BBox:
+    b = BBox(
+        max(target.minx, grid.minx),
+        max(target.miny, grid.miny),
+        min(target.maxx, grid.maxx),
+        min(target.maxy, grid.maxy),
+    )
+    if b.minx >= b.maxx or b.miny >= b.maxy:
+        raise ValueError("target bbox does not overlap the mosaic")
+    return b
+
+
+def crop_mercator(
+    image: Image.Image,
+    grid: BBox,
+    target: BBox,
+    *,
+    resampling: str = "bilinear",
+    snap: bool = False,
+) -> tuple[Image.Image, BBox]:
+    """Crop a Web Mercator mosaic to ``target`` (EPSG:3857 metres).
+
+    With ``snap=True`` the window is widened to whole source pixels and copied
+    without resampling. Otherwise the output covers ``target`` exactly, with
+    the source resolution rounded so that an integer number of pixels fits,
+    and is resampled.
+    """
+    res_x = (grid.maxx - grid.minx) / image.width
+    res_y = (grid.maxy - grid.miny) / image.height
+    t = _clip(target, grid)
+    if snap:
+        c0 = int(math.floor((t.minx - grid.minx) / res_x + 1e-9))
+        c1 = int(math.ceil((t.maxx - grid.minx) / res_x - 1e-9))
+        r0 = int(math.floor((grid.maxy - t.maxy) / res_y + 1e-9))
+        r1 = int(math.ceil((grid.maxy - t.miny) / res_y - 1e-9))
+        out = image.crop((c0, r0, c1, r1))
+        return out, BBox(
+            grid.minx + c0 * res_x,
+            grid.maxy - r1 * res_y,
+            grid.minx + c1 * res_x,
+            grid.maxy - r0 * res_y,
+        )
+    w = max(1, round((t.maxx - t.minx) / res_x))
+    h = max(1, round((t.maxy - t.miny) / res_y))
+    px = (t.maxx - t.minx) / w
+    py = (t.maxy - t.miny) / h
+    # Pillow maps output pixel centres through the affine matrix, in pixel-edge coordinates.
+    coeffs = (
+        px / res_x,
+        0.0,
+        (t.minx - grid.minx) / res_x,
+        0.0,
+        py / res_y,
+        (grid.maxy - t.maxy) / res_y,
+    )
+    out = image.transform((w, h), Image.Transform.AFFINE, coeffs, resample=_resampling(resampling))
+    return out, t
+
+
+def _lat_from_merc(y: float) -> float:
+    return math.degrees(2.0 * math.atan(math.exp(y / EARTH_RADIUS)) - math.pi / 2.0)
+
+
+def warp_to_lonlat(
+    image: Image.Image,
+    grid: BBox,
+    target: BBox,
+    *,
+    resampling: str = "bilinear",
+    width: int | None = None,
+) -> tuple[Image.Image, BBox]:
+    """Resample a Web Mercator mosaic to a lon/lat (EPSG:4326) grid covering ``target``.
+
+    ``target`` is ``(west, south, east, north)`` in degrees. Output pixels are
+    square in degrees; by default the width keeps the source resolution at the
+    bbox. The inverse mapping is exact per output row: longitude is linear in
+    Mercator x, and each row's latitude is projected to its source row.
+    """
+    res_x = (grid.maxx - grid.minx) / image.width
+    res_y = (grid.maxy - grid.miny) / image.height
+    gw = math.degrees(grid.minx / EARTH_RADIUS)
+    ge = math.degrees(grid.maxx / EARTH_RADIUS)
+    gs = _lat_from_merc(grid.miny)
+    gn = _lat_from_merc(grid.maxy)
+    t = _clip(target, BBox(gw, gs, ge, gn))
+    mx0, _ = lonlat_to_mercator(t.minx, 0.0)
+    mx1, _ = lonlat_to_mercator(t.maxx, 0.0)
+    if width is None:
+        width = max(1, round((mx1 - mx0) / res_x))
+    deg = (t.maxx - t.minx) / width
+    height = max(1, round((t.maxy - t.miny) / deg))
+    deg_y = (t.maxy - t.miny) / height
+    sx0 = (mx0 - grid.minx) / res_x
+    sx1 = (mx1 - grid.minx) / res_x
+
+    def src_row(lat: float) -> float:
+        _, my = lonlat_to_mercator(0.0, lat)
+        return (grid.maxy - my) / res_y
+
+    mesh = []
+    for r in range(height):
+        top = src_row(t.maxy - r * deg_y)
+        bottom = src_row(t.maxy - (r + 1) * deg_y)
+        # quad corners: upper-left, lower-left, lower-right, upper-right
+        mesh.append(((0, r, width, r + 1), (sx0, top, sx0, bottom, sx1, bottom, sx1, top)))
+    out = image.transform(
+        (width, height), Image.Transform.MESH, mesh, resample=_resampling(resampling)
+    )
+    return out, t