central-coast-aerial-survey/generate_heatmaps.py

"""Generate heatmaps using ArcGIS.

This script must be run from the version of python installed with ArcGIS, e.g.
C:/Python27/ArcGIS10.4/python generate_heatmaps.py
"""

import os
import re
import sys
import argparse
from glob import glob

import arcpy

# Checkout Spatial Analyst toolbox to prevent licence errors
arcpy.CheckOutExtension('Spatial')

# Overwrite duplicate files
arcpy.env.overwriteOutput = True


def parse_yaml(yaml_name):
    """Parse yaml file manually (when 'yaml' module is unavailable)

    Args:
        yaml_name:  path to input yaml file

    Returns:
        dict of objects from yaml file

    """
    with open(yaml_name, 'r') as f:
        yaml = f.read()

    # Fit wrapped strings on one line
    yaml = re.sub('\n\s', '', yaml)

    # Parse yaml file
    params = {}
    for line in yaml.split('\n'):
        if line:
            key, val = line.split(':')
            params[key] = val.strip()

    return params


def process(yaml_name):
    params = parse_yaml(yaml_name)
    beach = params['BEACH']
    base_name = os.path.splitext(os.path.basename(params['INPUT LAS']))[0]

    # Set all paths relative to current script location
    pwd = os.getcwd()
    input_tif_dir = os.path.join(pwd, params['TIF DEM FOLDER'])
    mxd_name = os.path.join(pwd, params['MAP DOCUMENT'])
    symb_lyr_name = os.path.join(pwd, params['SYMBOLOGY LAYER'])
    clip_poly = os.path.join(pwd, params['HEATMAP CROP POLY'])
    output_tif_dir = os.path.join(pwd, params['TIF HEATMAP FOLDER'])
    previous_date = params['PREVIOUS SURVEY DATE']

    # Check if previous survey date was provided
    try:
        int(previous_date)
    except ValueError:
        raise ValueError('No previous survey date provided')

    # Set paths for raster files
    current_raster = os.path.join(input_tif_dir, base_name + '.tif')
    previous_base_name = re.sub('\d+', previous_date, base_name)
    previous_raster = os.path.join(input_tif_dir, previous_base_name + '.tif')
    heatmap_raster = os.path.join(output_tif_dir, base_name + '.tif')

    print('processing {}'.format(beach))

    # Save heatmap in memory
    heatmap_memory = 'in_memory/' + base_name

    # Subtract the previous raster from the current raster
    arcpy.Minus_3d(
        in_raster_or_constant1=current_raster,
        in_raster_or_constant2=previous_raster,
        out_raster=heatmap_memory)

    # Clip the heatmap
    arcpy.Clip_management(
        in_raster=heatmap_memory,
        rectangle='350000 6290000 360000 6300000',
        out_raster=heatmap_raster,
        in_template_dataset=clip_poly,
        nodata_value='-9.999000e+003',
        clipping_geometry='ClippingGeometry',
        maintain_clipping_extent='NO_MAINTAIN_EXTENT')


def main():
    example_text = """examples:

    # Process single survey at specific beach
    C:/Python27/ArcGIS10.4/python generate_heatmaps.py survey-1-avoca.yaml

    # Process single survey at multiple beaches
    C:/Python27/ArcGIS10.4/python generate_heatmaps.py survey-1-avoca.yaml survey-1-pearl.yaml

    # Process all surveys at specific beach
    C:/Python27/ArcGIS10.4/python generate_heatmaps.py *avoca.yaml

    # Process all beaches for specific survey date
    C:/Python27/ArcGIS10.4/python generate_heatmaps.py survey-1*.yaml
    """

    # Set up command line arguments
    parser = argparse.ArgumentParser(
        epilog=example_text,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('input', help='path to yaml file(s)', nargs='*')

    # Print usage if no arguments are provided
    if len(sys.argv) == 1:
        parser.print_help(sys.stderr)
        sys.exit(1)

    # Parse arguments
    args = parser.parse_args()
    yaml_files = []
    [yaml_files.extend(glob(f)) for f in args.input]

    for yaml_file in yaml_files:
        process(yaml_file)


if __name__ == '__main__':
    main()