rasterio/scripts/rio

#!/usr/bin/env python

"""Rasterio command line interface"""

import argparse
import functools
import json
import logging
import os.path
import pprint
import sys
import warnings

import click

import rasterio
import rasterio.features
import rasterio.tool
import rasterio.warp

warnings.simplefilter('default')

def configure_logging(verbosity):
    log_level = max(10, 30 - 10*verbosity)
    logging.basicConfig(stream=sys.stderr, level=log_level)

# The CLI command group.
@click.group(help="Rasterio command line interface.")
@click.option('--verbose', '-v', count=True, help="Increase verbosity.")
@click.option('--quiet', '-q', count=True, help="Decrease verbosity.")
@click.pass_context
def cli(ctx, verbose, quiet):
    verbosity = verbose - quiet
    configure_logging(verbosity)
    ctx.obj = {}
    ctx.obj['verbosity'] = verbosity

# Commands are below.
#
# Command bodies less than ~20 lines, e.g. info() below, can go in this
# module. Longer ones, e.g. insp() shall call functions imported from
# rasterio.tool.

# Info command.
@cli.command(short_help="Print information about a data file.")
@click.argument('src_path', type=click.Path(exists=True))
@click.option('--meta', 'aspect', flag_value='meta', default=True, 
    help="Show data file structure (default).")
@click.option('--tags', 'aspect', flag_value='tags', 
    help="Show data file tags.")
@click.option('--indent', default=2, type=int, 
              help="Indentation level for pretty printed output")
@click.option('--namespace', help="Select a tag namespace.")
@click.pass_context
def info(ctx, src_path, aspect, indent, namespace):
    verbosity = ctx.obj['verbosity']
    logger = logging.getLogger('rio')
    try:
        with rasterio.drivers(CPL_DEBUG=verbosity>2):
            with rasterio.open(src_path) as src:
                if aspect == 'meta':
                    pprint.pprint(src.meta, indent=indent)
                elif aspect == 'tags':
                    pprint.pprint(src.tags(ns=namespace), indent=indent)
        sys.exit(0)
    except Exception:
        logger.exception("Failed. Exception caught")
        sys.exit(1)

# Insp command.
@cli.command(short_help="Open a data file and start an interpreter.")
@click.argument('src_path', type=click.Path(exists=True))
@click.option('--mode', type=click.Choice(['r', 'r+']), default='r', help="File mode (default 'r').")
@click.pass_context
def insp(ctx, src_path, mode):
    verbosity = ctx.obj['verbosity']
    logger = logging.getLogger('rio')
    try:
        with rasterio.drivers(CPL_DEBUG=verbosity>2):
            with rasterio.open(src_path, mode) as src:
                rasterio.tool.main(
                    "Rasterio %s Interactive Inspector (Python %s)\n"
                    'Type "src.meta", "src.read_band(1)", or "help(src)" '
                    'for more information.' %  (
                        rasterio.__version__,
                        '.'.join(map(str, sys.version_info[:3]))),
                    src)
        sys.exit(0)
    except Exception:
        logger.exception("Failed. Exception caught")
        sys.exit(1)

# Bounds command.
@cli.command(short_help="Write bounding boxes to stdout as GeoJSON.",
    help="""
Write bounding boxes to stdout as GeoJSON for use with, e.g., geojsonio

  $ rio bounds *.tif | geojsonio
""")
@click.argument('src_paths', nargs=-1, type=click.Path(exists=True))
@click.option('--precision', type=int, default=-1,
              help="Decimal precision of coordinates.")
@click.option('--indent', default=None, type=int, 
              help="Indentation level for JSON output")
@click.option('--projected/--geographic', default=False,
              help="Output in projected coordinates (default is geographic).")
@click.pass_context
def bounds(ctx, src_paths, precision, indent, projected):
    verbosity = ctx.obj['verbosity']
    logger = logging.getLogger('rio')
    try:
        collection = {'type': 'FeatureCollection'}
        features = []
        col_xs = []
        col_ys = []
        with rasterio.drivers(CPL_DEBUG=verbosity>2):
            for i, path in enumerate(src_paths):
                with rasterio.open(path) as src:
                    bounds = src.bounds
                    xs = [bounds[0], bounds[2]]
                    ys = [bounds[1], bounds[3]]
                    if not projected:
                        xs, ys = rasterio.warp.transform(
                            src.crs, {'init': 'epsg:4326'}, xs, ys)
                    if precision >= 0:
                        xs = [round(v, precision) for v in xs]
                        ys = [round(v, precision) for v in ys]
                    col_xs.extend(xs)
                    col_ys.extend(ys)
                    features.append(
                        {'type': 'Feature', 'geometry': 
                            {'type': 'Polygon', 'coordinates': [
                                [[xs[0], ys[0]], [xs[1], ys[0]], [xs[1], ys[1]],
                                 [xs[0], ys[1]], [xs[0], ys[0]] ]]},
                        'properties': {'id': str(i), 'title': path} })
        collection['bbox'] = [
            min(col_xs), min(col_ys), max(col_xs), max(col_ys)]
        collection['features'] = features
        print(json.dumps(collection, indent=indent, sort_keys=True))
        sys.exit(0)
    except Exception:
        logger.exception("Failed. Exception caught")
        sys.exit(1)

def transform_geom(transformer, g, precision=-1):
    if g['type'] == 'Point':
        x, y = g['coordinates']
        xp, yp = transformer([x], [y])
        if precision >= 0:
            xp = [round(v, precision) for v in xp]
            yp = [round(v, precision) for v in yp]
        new_coords = list(zip(xp, yp))[0]
    if g['type'] in ['LineString', 'MultiPoint']:
        xp, yp = transformer(*zip(g['coordinates']))
        if precision >= 0:
            xp = [round(v, precision) for v in xp]
            yp = [round(v, precision) for v in yp]
        new_coords = list(zip(xp, yp))
    elif g['type'] in ['Polygon', 'MultiLineString']:
        new_coords = []
        for piece in g['coordinates']:
            xp, yp = transformer(*zip(*piece))
            if precision >= 0:
                xp = [round(v, precision) for v in xp]
                yp = [round(v, precision) for v in yp]
            new_coords.append(list(zip(xp, yp)))
    elif g['type'] == 'MultiPolygon':
        parts = g['coordinates']
        new_coords = []
        for part in parts:
            inner_coords = []
            for ring in part:
                xp, yp = transformer(*zip(*ring))
                if precision >= 0:
                    xp = [round(v, precision) for v in xp]
                    yp = [round(v, precision) for v in yp]
                inner_coords.append(list(zip(xp, yp)))
            new_coords.append(inner_coords)
    g['coordinates'] = new_coords
    return g

# Shapes command.
@cli.command(short_help="Write the shapes of features.")
@click.argument('src_path', type=click.Path(exists=True))
@click.option('--precision', type=int, default=-1,
              help="Decimal precision of coordinates.")
@click.option('--indent', default=None, type=int, 
              help="Indentation level for JSON output")
@click.option('--projected/--geographic', default=False,
              help="Output in projected coordinates (default is geographic).")
@click.option('--json-seq/--json-obj', default=False,
              help="Write a JSON sequence (default is object).")
@click.pass_context
def shapes(ctx, src_path, precision, indent, projected, json_seq):
    verbosity = ctx.obj['verbosity']
    logger = logging.getLogger('rio')
    def to_feature(i, geom):
        return {
            'type': 'Feature', 
            'id': str(i), 
            'properties': {'val': i},
            'geometry': geom }
    try:
        collection = {'type': 'FeatureCollection'}
        with rasterio.drivers(CPL_DEBUG=verbosity>2):
            with rasterio.open(src_path) as src:
                bounds = src.bounds
                xs = [bounds[0], bounds[2]]
                ys = [bounds[1], bounds[3]]
                if not projected:
                    xs, ys = rasterio.warp.transform(
                        src.crs, {'init': 'epsg:4326'}, xs, ys)
                if precision >= 0:
                    xs = [round(v, precision) for v in xs]
                    ys = [round(v, precision) for v in ys]
                collection['bbox'] = (min(xs), min(ys), max(xs), max(ys))
                if not projected:
                    transformer = functools.partial(
                        rasterio.warp.transform,
                        src.crs,
                        {'init': 'epsg:4326'})
                    feature_gen = (to_feature(i,
                                    transform_geom(
                                        transformer, g, precision)) for g, i
                                in rasterio.features.shapes(
                                    src.read(1), transform=src.affine))
                else:
                    feature_gen = (to_feature(i, g) for g, i
                                in rasterio.features.shapes(
                                    src.read(1), transform=src.affine))
                if json_seq:
                    sys.stdout.write(
                        json.dumps(collection, indent=indent, sort_keys=True) 
                        + "\n")
                    for feature in feature_gen:
                        sys.stdout.write(
                            json.dumps(feature, indent=indent, sort_keys=True)
                        + "\n")
                else:
                    collection['features'] = list(feature_gen)
                    sys.stdout.write(
                        json.dumps(collection, indent=indent, sort_keys=True))
        sys.exit(0)
    except Exception:
        logger.exception("Failed. Exception caught")
        sys.exit(1)

# Transform command.
@cli.command(short_help="Transform coordinates.")
@click.argument('input', type=click.File('rb'))
@click.option('--src_crs', default='EPSG:4326', help="Source CRS.")
@click.option('--dst_crs', default='EPSG:4326', help="Destination CRS.")
@click.option('--interleaved', 'mode', flag_value='interleaved', default=True)
@click.option('--precision', type=int, default=-1,
              help="Decimal precision of coordinates.")
@click.pass_context
def transform(ctx, input, src_crs, dst_crs, mode, precision):
    verbosity = ctx.obj['verbosity']
    logger = logging.getLogger('rio')
    try:
        if mode == 'interleaved':
            coords = json.loads(input.read().decode('utf-8'))
            xs = coords[::2]
            ys = coords[1::2]
        else:
            raise ValueError("Invalid input type '%s'" % input_type)
        with rasterio.drivers(CPL_DEBUG=verbosity>2):
            if src_crs.startswith('EPSG'):
                src_crs = {'init': src_crs}
            elif os.path.exists(src_crs):
                with rasterio.open(src_crs) as f:
                    src_crs = f.crs
            if dst_crs.startswith('EPSG'):
                dst_crs = {'init': dst_crs}
            elif os.path.exists(dst_crs):
                with rasterio.open(dst_crs) as f:
                    dst_crs = f.crs
            xs, ys = rasterio.warp.transform(src_crs, dst_crs, xs, ys)
            if precision >= 0:
                xs = [round(v, precision) for v in xs]
                ys = [round(v, precision) for v in ys]
        if mode == 'interleaved':
            result = [0]*len(coords)
            result[::2] = xs
            result[1::2] = ys
        print(json.dumps(result))
        sys.exit(0)
    except Exception:
        logger.exception("Failed. Exception caught")
        sys.exit(1)


if __name__ == '__main__':
    cli()