Slice a raster into tiles using rio-tiler, write to a pmtiles file

rio-to-pmtiles.py

import argparse
import os
import warnings

import mercantile
import rasterio
from pmtiles.tile import Compression
from pmtiles.tile import TileType
from pmtiles.tile import zxy_to_tileid
from pmtiles.writer import Writer
from rasterio.warp import transform_bounds
from rio_tiler.io import Reader
from rio_tiler.profiles import img_profiles
from tqdm import tqdm


def slice_to_pmtiles(
    input_file,
    output_file,
    min_zoom,
    max_zoom,
    image_type,
    tile_size,
    resampling_method,
    quality,
    compression_level,
):
    # Get bounding box and CRS of the file
    with rasterio.open(input_file) as src:
        bbox = list(src.bounds)
        src_crs = src.crs

        # Convert bounding box to EPSG:4326
        bbox_epsg4326 = transform_bounds(src_crs, "EPSG:4326", *bbox)

    with open(output_file, "wb") as output_file:
        pmtiles_writer = Writer(output_file)
        reader_options = {
            "resampling_method": resampling_method,
        }
        render_options = img_profiles.get(image_type.lower())
        if quality is not None:
            render_options["quality"] = quality
        if compression_level is not None:
            render_options["zlevel"] = compression_level

        with Reader(input_file, options=reader_options) as cog:
            tiles = mercantile.tiles(
                bbox[0], bbox[1], bbox[2], bbox[3], range(min_zoom, max_zoom + 1)
            )
            for tile in tqdm(tiles, unit=" tiles", desc="Processing tiles"):
                x, y, z = tile.x, tile.y, tile.z
                img = cog.tile(x, y, z, tilesize=tile_size)
                tile_content = img.render(img_format=image_type, **render_options)
                pmtiles_writer.write_tile(zxy_to_tileid(z, x, y), bytes(tile_content))

        pmtiles_writer.finalize(
            header={
                "tile_compression": Compression.NONE,
                "tile_type": TileTypeFromImageFormat(image_type),
                "min_lon_e7": int(bbox_epsg4326[0] * 1e7),
                "min_lat_e7": int(bbox_epsg4326[1] * 1e7),
                "max_lon_e7": int(bbox_epsg4326[2] * 1e7),
                "max_lat_e7": int(bbox_epsg4326[3] * 1e7),
            },
            metadata={},
        )


def ignore_warnings():
    warnings.simplefilter("ignore")


def TileTypeFromImageFormat(image_format: str) -> TileType:
    if image_format == "JPEG":
        return TileType.JPEG
    elif image_format == "PNG":
        return TileType.PNG
    elif image_format == "WEBP":
        return TileType.WEBP
    else:
        raise ValueError("Invalid image format.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Slice raster into tiles and write to PMTiles."
    )
    parser.add_argument("input_file", help="Input raster file")
    parser.add_argument("output_file", help="Output PMTiles file")
    parser.add_argument(
        "--min_zoom", required=True, type=int, help="Minimum zoom level"
    )
    parser.add_argument(
        "--max_zoom", required=True, type=int, help="Maximum zoom level"
    )
    parser.add_argument(
        "--image_type",
        choices=["JPEG", "PNG", "WEBP"],
        default="JPEG",
        help="Image type for tiles",
    )
    parser.add_argument(
        "--tile_size", choices=[256, 512], type=int, default=512, help="Tile size"
    )
    parser.add_argument(
        "--resampling_method",
        choices=[
            "nearest",
            "bilinear",
            "cubic",
            "cubic_spline",
            "lanczos",
            "average",
            "mode",
            "gauss",
            "max",
            "min",
            "med",
            "q1",
            "q3",
        ],
        default="nearest",
        help="Resampling method",
    )
    parser.add_argument("--quality", type=int, help="Image quality (0-100)")
    parser.add_argument(
        "--compression_level",
        type=int,
        choices=range(1, 10),
        help="Compression level (1-9)",
    )

    args = parser.parse_args()

    if args.quality is not None:
        if not (0 <= args.quality <= 100):
            raise ValueError("Quality must be in the range 0-100")
        if args.image_type == "PNG":
            raise ValueError("Quality parameter is not applicable for PNG format")

    if args.compression_level is not None:
        if args.image_type != "PNG":
            raise ValueError("Compression level is only applicable for PNG format")

    ignore_warnings()

    if os.path.exists(args.output_file):
        raise ValueError("Output file already exists.")

    slice_to_pmtiles(
        args.input_file,
        args.output_file,
        args.min_zoom,
        args.max_zoom,
        args.image_type,
        args.tile_size,
        args.resampling_method,
        args.quality,
        args.compression_level,
    )