generate_tiles.py

#!/usr/bin/env python3

import os
import sys
import warnings
import logging
import subprocess
import shutil
import json
from glob import glob

import tile_index

LOG = logging.getLogger(__name__)


def group_files(products, input_files):
    """Group input geotiff files by product."""
    groups = {}
    for prod in products:
        prods_files = [f for f in input_files if prod in f]
        if prods_files:
            groups[prod] = prods_files
    if len(groups) != len(products):
        warnings.warn("Not all product geotiffs were provided.")
    return groups


def remap_to_lonlat(itif, otif):
    """Remap a single geotiff by calling gdalwarp."""
    try:
        subprocess.run(['gdalwarp', '-multi', '-wo', 'NUM_THREADS=ALL_CPUS', '-t_srs', 'EPSG:4326', itif, otif], check=True)
    except subprocess.CalledProcessError:
        LOG.error("Could not remap geotiff %s -> %s" % (itif, otif))
        return None
    return otif


def remap_tifs(input_tifs, out_dir, remap_suffix):
    """Remap all input geotiffs to EPSG:4326."""
    for itif in input_tifs:
        ifn = os.path.basename(itif)
        otif = os.path.join(out_dir, ifn.replace('.tif', remap_suffix))
        otif = remap_to_lonlat(itif, otif)
        if otif is not None:
            yield otif


def link_or_copy(input_tifs, out_dir):
    """Hardlink input tifs to output directory."""
    for prod_file in input_tifs:
        out_file = os.path.join(out_dir, os.path.basename(prod_file))
        if os.path.exists(out_file):
            yield out_file
            continue
        try:
            os.link(prod_file, out_file)
        except OSError:
            # on different mounts probably?
            shutil.copy2(prod_file, out_file)
        yield out_file


def main():
    import argparse
    parser = argparse.ArgumentParser(description="Take input geotiffs and generate mapserver compatible tiles.")
    parser.add_argument('--remap', action='store_true',
                        help="Remap input geotiffs to EPSG:4326")
    parser.add_argument('--remap-suffix', default='_ll.tif',
                        help="Replace 'tif' with provided suffix when geotiffs are remapped.")
    parser.add_argument('-p', '--products', nargs="*",
                        help="Product names to group together in each "
                             "'layer'. Product name must be in the filename.")
    parser.add_argument('--shape-file', default='{product}.shp',
                        help="Shapefile filename pattern to use and placed "
                             "in the output directory. "
                             "(default: '{product}.shp')")
    parser.add_argument('--json-times', action='store_true',
                        help="Print a JSON dict to stdout containing removed "
                             "and added times (ex. '{\'times_removed\': [], "
                             "\'times_added\': []}'). Times are in the form: "
                             "YYYY-MM-DDTHH:MM:SS.")
    parser.add_argument('--shared-bbox', action='store_true',
                        help="Assume that all input geotiffs have the same "
                             "bounding box so no need to compute it for every "
                             "input.")
    parser.add_argument('out_dir',
                        help="Output path to save tile information to (ex. '/data/tiles/{product}')")
    parser.add_argument('input_files', nargs="+",
                        help="Input geotiffs to generate tiles for (separate from product lists with '--')")
    args = parser.parse_args()

    # allow for glob patterns to reduce number of files passed through command line
    args.input_files = [fn for glob_pat in args.input_files for fn in glob(glob_pat)]
    groups = group_files(args.products, args.input_files)
    all_added = set()
    all_removed = set()
    for prod, prod_files in groups.items():
        out_dir = args.out_dir.format(product=prod)
        os.makedirs(out_dir, exist_ok=True)
        shp_fn = args.shape_file.format(product=prod)
        shp_pathname = os.path.join(out_dir, shp_fn)

        if args.remap:
            # remap if needed
            prod_files = list(remap_tifs(prod_files, out_dir, args.remap_suffix))
        else:
            # hardlink if needed
            prod_files = list(link_or_copy(prod_files, out_dir))

        # get all products in the current directory
        ext = os.path.splitext(prod_files[-1])[-1]
        all_prod_files = sorted(glob(os.path.join(out_dir, '*' + prod + '*' + ext)))

        # create shape file
        LOG.info("Rebuilding shapefile index with:\n\t{}".format(", ".join(all_prod_files)))
        removed_times, added_times = tile_index.index(
            all_prod_files, shp_pathname, shared_bbox=args.shared_bbox)
        all_removed.update(removed_times)
        all_added.update(added_times)

    if args.json_times:
        # all_removed = [x.isoformat() for x in sorted(all_removed)]
        # all_added = [x.isoformat() for x in sorted(all_added)]
        all_removed = sorted(all_removed)
        all_added = sorted(all_added)
        print(json.dumps({'removed_times': all_removed, 'added_times': all_added}))


if __name__ == "__main__":
    sys.exit(main())