Python wrapper around `gdaldem color-relief`to produce a colourised image of any raster format with gdal.

gdal2image.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Basically a wrapper around 
gdaldem color-relief which allows
automation of images and specifying
some nice color features etc
"""
import argparse
import matplotlib
import os 
import gdal 
import matplotlib.cm
import numpy as np
import csv
from PIL import Image

# some acceptable constants
N_COLOUR_STEPS = 10 # generally irrelevant
MIN_PERCENTILE = 5.0
MAX_PERCENTILE = 95.0


if __name__ == "__main__":
    """
    get options
    """
    helpTxt = """Make image of a gdal file as a png. It's really just a wrapper around
    `gdaldem color-relief` ;)"""
    parser = argparse.ArgumentParser(description=helpTxt)
    parser.add_argument('input_file', default='input_file.vrt', type=str,
                            help="""input file readable by gdal """)
    parser.add_argument('output_file', default='out.vrt',type=str,
                            help="""output filename """)
    parser.add_argument('--cmap', default="viridis",
                            help="""a matplotlib colormap. """)
    parser.add_argument('--vmin', default=None, type=float,
                            help="""minimum value mapped. """)
    parser.add_argument('--vmax', default=None, type=float,
                            help="""maximum value mapped. """)
    parser.add_argument('--band', default=1,type=int,
                            help="""band to produce image of. """)
    parser.add_argument('--display', default=False, type=bool,
                            help="""display the image with the native image viewer. """)
    options = parser.parse_args()

    # get the input file
    srcDS = gdal.Open(options.input_file)

    # Figure out a good set of lower-upper bounds...
    band = srcDS.GetRasterBand(options.band)
    # get range of the data
    min, max, mean, std = band.GetStatistics(0,1)
    # compute a histogram
    hist = np.array(band.GetHistogram(min, max)).astype(float)
    bins = np.linspace(min, max, 256)
    chist = np.cumsum(hist)
    # approx the pdf
    total = np.sum(hist)
    phist = hist/total 
    co  = 100.0/total
    percentiles = np.cumsum(hist * co)
    # get bounds orse some defaults
    if options.vmin==None:
        # specify lower bound as 5% percentile
        idx_min = np.argmin(np.abs(percentiles-MIN_PERCENTILE))
        options.vmin = bins[idx_min]
    if options.vmax==None:
        # specify upper bound as 95% percentile
        idx_max = np.argmin(np.abs(percentiles-MAX_PERCENTILE))
        options.vmax = bins[idx_max]

    """
    make the colormap file by
    sampling along the colormap at N samples
    """
    cmap = matplotlib.cm.get_cmap(options.cmap)
    grid_cls = []
    for i, j in zip(np.linspace(options.vmin,
                                 options.vmax, 
                                 N_COLOUR_STEPS), 
                    np.linspace(0, 1, N_COLOUR_STEPS)):
        rgba = list(cmap(j,bytes=True))
        rgba.insert(0, i)
        grid_cls.append(rgba)

    # write the necessary color map file
    with open('.cmap.txt', 'w') as f:
        wr = csv.writer(f, delimiter=" ")
        wr.writerows(grid_cls)

    # That's it!
    gdal.DEMProcessing(options.output_file, srcDS, colorFilename=".cmap.txt",
                format="PNG", processing='color-relief')

    # do we want to view it too?
    if options.display:
        #import webbrowser
        #webbrowser.open(options.output_file)
        Image.open(options.output_file).show()