collect-terrain-tiles.py

#!/usr/bin/env python
from __future__ import print_function
from math import log, tan, pi
from itertools import product
from argparse import ArgumentParser
from os.path import join, splitext

import tempfile, shutil, urllib, io, sys, subprocess

tile_url = 'https://elevation-tiles-prod.s3.amazonaws.com/geotiff/{z}/{x}/{y}.tif'

def mercator(lat, lon, zoom):
    ''' Convert latitude, longitude to z/x/y tile coordinate at given zoom.
    '''
    # convert to radians
    x1, y1 = lon * pi/180, lat * pi/180

    # project to mercator
    x2, y2 = x1, log(tan(0.25 * pi + 0.5 * y1))

    # transform to tile space
    tiles, diameter = 2 ** zoom, 2 * pi
    x3, y3 = int(tiles * (x2 + pi) / diameter), int(tiles * (pi - y2) / diameter)

    return zoom, x3, y3

def tiles(zoom, lat1, lon1, lat2, lon2):
    ''' Convert geographic bounds into a list of tile coordinates at given zoom.
    '''
    # convert to geographic bounding box
    minlat, minlon = min(lat1, lat2), min(lon1, lon2)
    maxlat, maxlon = max(lat1, lat2), max(lon1, lon2)

    # convert to tile-space bounding box
    _, xmin, ymin = mercator(maxlat, minlon, zoom)
    _, xmax, ymax = mercator(minlat, maxlon, zoom)

    # generate a list of tiles
    xs, ys = range(xmin, xmax+1), range(ymin, ymax+1)
    tiles = [(zoom, x, y) for (y, x) in product(ys, xs)]

    return tiles

def download(output_path, tiles, verbose=True):
    ''' Download list of tiles to a temporary directory and return its name.
    '''
    dir = tempfile.mkdtemp(prefix='collected-')
    _, ext = splitext(output_path)

    merge_geotiff = bool(ext.lower() in ('.tif', '.tiff', '.geotiff'))

    try:
        files = []

        for (z, x, y) in tiles:
            response = urllib.urlopen(tile_url.format(z=z, x=x, y=y))
            if response.getcode() != 200:
                raise RuntimeError('No such tile: {}'.format((z, x, y)))
            if verbose:
                print('Downloaded', response.url, file=sys.stderr)

            with io.open(join(dir, '{}-{}-{}.tif'.format(z, x, y)), 'wb') as file:
                file.write(response.read())
                files.append(file.name)

        if merge_geotiff:
            if verbose:
                print('Combining', len(files), 'into', output_path, '...', file=sys.stderr)
            temp_tif = join(dir, 'temp.tif')
            subprocess.check_call(['gdal_merge.py', '-o', temp_tif] + files)
            shutil.move(temp_tif, output_path)
        else:
            if verbose:
                print('Moving', dir, 'to', output_path, '...', file=sys.stderr)
            shutil.move(dir, output_path)

    finally:
        if merge_geotiff:
            shutil.rmtree(dir)

parser = ArgumentParser(description='''Collect Mapzen elevation tiles into a
single GeoTIFF or directory. If output_path ends in ".tif", ".tiff", or
".geotiff", gdal_merge.py will be called to merge all downloaded tiles into a
single image. Otherwise, they are collected into the named directory.''')

parser.add_argument('--testing', action='store_const', const=True, default=False,
                    help='If set, run unit tests and bail out.')

parser.add_argument('--bounds', metavar='DEG', type=float, nargs=4,
                    default=(37.8434, -122.3193, 37.7517, -122.0927),
                    help='Geographic bounds given as two lat/lon pairs. Defaults to Oakland, CA hills.')

parser.add_argument('--zoom', type=int, default=12,
                    help='Map zoom level given as integer. Defaults to 12.')

parser.add_argument('output_path', help='Output GeoTIFF filename or local directory name.')

if __name__ == '__main__':
    args = parser.parse_args()

    download(args.output_path, tiles(args.zoom, *args.bounds))