cog_translate.py

"""Create geotiff."""

import os
import multiprocessing

import click

import rasterio
from rasterio.io import MemoryFile
from rasterio.rio import options
from rasterio.enums import Resampling
from rasterio.shutil import copy


@click.command()
@options.file_in_arg
@options.file_out_arg
@click.option(
    "--blocksize",
    help="Internal blocksize.",
    type=int,
    default=256,
)
@click.option(
    "--overview-level",
    type=int,
    help="Overview level.",
    default=5,
)
@click.option(
    "--overview-resampling",
    help="Resampling algorithm.",
    type=click.Choice(
        [it.name for it in Resampling if it.value in [0, 1, 2, 3, 4, 5, 6, 7]]
    ),
    default="nearest",
)
@click.option(
    "--threads",
    help="threads",
    type=int,
    default=lambda: os.environ.get("MAX_THREADS", multiprocessing.cpu_count() * 5),
)
def cog_translate(
    input,
    output,
    blocksize,
    overview_level,
    overview_resampling,
    threads,
):
    """
    Create Cloud Optimized Geotiff.

    Parameters
    ----------
    src_path : str or PathLike object
        A dataset path or URL. Will be opened in "r" mode.
    dst_path : str or Path-like object
        An output dataset path or or PathLike object.
        Will be opened in "w" mode.
    blocksize : int, optional (default: 256)
        Internal tile size.
    overview_level : int, optional (default: 5)
        COGEO overview (decimation) level
    overview_resampling : str, optional (default: "nearest")
         Resampling algorithm for overviews
    """
    profile = {
        "driver": "GTiff",
        "interleave": "pixel",
        "tiled": True,
        "blockxsize": blocksize,
        "blockysize": blocksize,
        "compress": "DEFLATE",
        # OPTIONAL
        # "predictor": 2,
        # "zlevel": 9
    }

    config = dict(
        NUM_THREADS=threads,
        GDAL_TIFF_OVR_BLOCKSIZE=os.environ.get("GDAL_TIFF_OVR_BLOCKSIZE", blocksize),
    )

    with rasterio.Env(**config):
        with rasterio.open(input) as src_dst:

            meta = src_dst.meta
            meta.update(**profile)
            meta.pop("compress", None)

            with MemoryFile() as memfile:
                with memfile.open(**meta) as mem:
                    wind = list(mem.block_windows(1))
                    for ij, w in wind:
                        mem.write(src_dst.read(window=w), window=w)

                    overviews = [2 ** j for j in range(1, overview_level + 1)]

                    mem.build_overviews(overviews, Resampling[overview_resampling])

                    mem.set_band_description(1, src_dst.descriptions[0])
                    tags = src_dst.tags()
                    tags.update(dict(
                        OVR_RESAMPLING_ALG=Resampling[overview_resampling].name.upper()
                    ))
                    mem.update_tags(**tags)

                    copy(mem, output, copy_src_overviews=True, **profile)


if __name__ == '__main__':
    cog_translate()

python cog_translate.py --help
Usage: cog_translate.py [OPTIONS] INPUT OUTPUT

  Create Cloud Optimized Geotiff.

Options:
  --blocksize INTEGER             Internal blocksize.
  --overview-level INTEGER        Overview level.
  --overview-resampling [nearest|bilinear|cubic|cubic_spline|lanczos|average|mode|gauss]
                                  Resampling algorithm.
  --threads INTEGER               threads
  --help                          Show this message and exit.

Example:
python cog_translate.py LC08_L1TP_043035_20151227_20170224_01_A1_pixel_qa.TIF LC08_L1TP_043035_20151227_20170224_01_A1_pixel_qa_cog.TIF