landsat8/image_process.py at master · apburnes/landsat8 · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Created on Thu Nov  3 11:56:48 2011

The main program. The program takes 3 files as input:

    #. The red band (any GDAL-compatible format will do)
    #. the near-infrarred band (again as above)
    #. An output file name

dditionally, one can specify the output format. By default, we use  GeoTIFF.


@author: Jose Gómez-Dans (UCL/NCEO) - j.gomez-dans@ucl.ac.uk
"""
import os
import sys
import optparse

import numpy as np

from osgeo import gdal

def calculate_ndvi ( red_filename, nir_filename ):
    """
    A function to calculate the Normalised Difference Vegetation Index
    from red and near infrarred reflectances. The reflectance data ought to
    be present on two different files, specified by the varaibles
    `red_filename` and `nir_filename`. The file format ought to be
    recognised by GDAL
    """

    g_red = gdal.Open ( red_filename )
    red = g_red.ReadAsArray()
    g_nir = gdal.Open ( nir_filename )
    nir = g_nir.ReadAsArray()
    if ( g_red.RasterXSize != g_nir.RasterXSize ) or \
            ( g_red.RasterYSize != g_nir.RasterYSize ):
        print "ERROR: Input datasets do't match!"
        print "\t Red data shape is %dx%d" % ( red.shape )
        print "\t NIR data shape is %dx%d" % ( nir.shape )

        sys.exit ( -1 )
    passer = np.logical_and ( red > 1, nir > 1 )
    ndvi = np.where ( passer,  (1.*nir - 1.*red ) / ( 1.*nir + 1.*red ), -999 )
    return ndvi

def save_raster ( output_name, raster_data, dataset, driver="GTiff" ):
    """
    A function to save a 1-band raster using GDAL to the file indicated
    by ``output_name``. It requires a GDAL-accesible dataset to collect
    the projection and geotransform.
    """

    # Open the reference dataset
    g_input = gdal.Open ( dataset )
    # Get the Geotransform vector
    geo_transform = g_input.GetGeoTransform ()
    x_size = g_input.RasterXSize # Raster xsize
    y_size = g_input.RasterYSize # Raster ysize
    srs = g_input.GetProjectionRef () # Projection
    # Need a driver object. By default, we use GeoTIFF
    if driver == "GTiff":
        driver = gdal.GetDriverByName ( driver )
        dataset_out = driver.Create ( output_name, x_size, y_size, 1, \
                gdal.GDT_Float32, ['TFW=YES', \
                'COMPRESS=LZW', 'TILED=YES'] )
    else:
        driver = gdal.GetDriverByName ( driver )
        dataset_out = driver.Create ( output_name, x_size, y_size, 1, \
                gdal.GDT_Float32 )

    dataset_out.SetGeoTransform ( geo_transform )
    dataset_out.SetProjection ( srs )
    dataset_out.GetRasterBand ( 1 ).WriteArray ( \
            raster_data.astype(np.float32) )
    dataset_out.GetRasterBand ( 1 ).SetNoDataValue ( float(-999) )
    dataset_out = None

if __name__ == "__main__":

    arg_parser = optparse.OptionParser()
    arg_parser.add_option( '-r', '--red', dest="red_fname", \
            help="The RED data" )
    arg_parser.add_option( '-n', '--nir', dest="nir_fname", \
            help="The NIR data" )
    arg_parser.add_option( '-o', '--output', dest="out_fname", \
            help="The output dataset" )
    arg_parser.add_option( '-f', '--format', dest="out_format", \
            default="GTiff", help="Output format" )
    options, extra_junk = arg_parser.parse_args ()

    if not os.path.exists ( options.red_fname ):
        print "ERROR: The red filename %s does not exist" % options.red_fname
        sys.exit ( -1 )
    if not os.path.exists ( options.nir_fname ):
        print "ERROR: The nir filename %s does not exist" % options.nir_fname
        sys.exit ( -1 )
    if os.path.exists ( options.out_fname):
        print "ERROR: The output filename %s does already exist" % \
            options.out_fname
        print "\t Select a different one, or delete the file."
        sys.exit ( -1 )


    c_ndvi = calculate_ndvi ( options.red_fname, options.nir_fname )
    save_raster ( options.out_fname, c_ndvi, options.red_fname, \
            driver=options.out_format )