Package osgeo :: Package utils :: Module gdal_calc |
[hide private]
[frames] | no frames]

# Source Code for Module osgeo.utils.gdal_calc

      1  #!/usr/bin/env python3
      2  # -*- coding: utf-8 -*-
      3  # ******************************************************************************
      4  #
      5  #  Project:  GDAL
      6  #  Purpose:  Command line raster calculator with numpy syntax
      7  #  Author:   Chris Yesson, chris.yesson@ioz.ac.uk
      8  #
      9  # ******************************************************************************
     10  #  Copyright (c) 2010, Chris Yesson <chris.yesson@ioz.ac.uk>
     11  #  Copyright (c) 2010-2011, Even Rouault <even dot rouault at spatialys.com>
     12  #  Copyright (c) 2016, Piers Titus van der Torren <pierstitus@gmail.com>
     13  #
     14  #  Permission is hereby granted, free of charge, to any person obtaining a
     15  #  copy of this software and associated documentation files (the "Software"),
     16  #  to deal in the Software without restriction, including without limitation
     17  #  the rights to use, copy, modify, merge, publish, distribute, sublicense,
     18  #  and/or sell copies of the Software, and to permit persons to whom the
     19  #  Software is furnished to do so, subject to the following conditions:
     20  #
     21  #  The above copyright notice and this permission notice shall be included
     22  #  in all copies or substantial portions of the Software.
     23  #
     24  #  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
     25  #  OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
     26  #  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
     27  #  THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
     28  #  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
     29  #  FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
     30  #  DEALINGS IN THE SOFTWARE.
     31  # ******************************************************************************
     32
     33  ################################################################
     34  # Command line raster calculator with numpy syntax. Use any basic arithmetic supported by numpy arrays such as +-*\ along with logical operators such as >.  Note that all files must have the same dimensions, but no projection checking is performed.  Use gdal_calc.py --help for list of options.
     35
     36  # example 1 - add two files together
     37  # gdal_calc.py -A input1.tif -B input2.tif --outfile=result.tif --calc="A+B"
     38
     39  # example 2 - average of two layers
     40  # gdal_calc.py -A input.tif -B input2.tif --outfile=result.tif --calc="(A+B)/2"
     41
     42  # example 3 - set values of zero and below to null
     43  # gdal_calc.py -A input.tif --outfile=result.tif --calc="A*(A>0)" --NoDataValue=0
     44
     45  # example 4 - using logical operator to keep a range of values from input:
     46  # gdal_calc.py -A input.tif --outfile=result.tif --calc="A*logical_and(A>100,A<150)"
     47
     48  # example 5 - work with multiple bands:
     49  # gdal_calc.py -A input.tif --A_band=1 -B input.tif --B_band=2 --outfile=result.tif --calc="(A+B)/2" --calc="A*logical_and(A>100,A<150)"
     50  ################################################################
     51
     52  from optparse import OptionParser, OptionConflictError, Values
     53  import os
     54  import os.path
     55  import sys
     56  import shlex
     57
     58  import numpy
     59
     60  from osgeo import gdal
     61  from osgeo import gdalnumeric
     62
     63
     64  # create alphabetic list for storing input layers
     65  AlphaList =  ["A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M",
     66                "N", "O", "P", "Q", "R", "S", "T", "U", "V", "W", "X", "Y", "Z"]
     67
     68  # set up some default nodatavalues for each datatype
     69  DefaultNDVLookup = {'Byte': 255, 'UInt16': 65535, 'Int16': -32767, 'UInt32': 4294967293, 'Int32': -2147483647, 'Float32': 3.402823466E+38, 'Float64': 1.7976931348623158E+308}
     70
     71
     72 -def DoesDriverHandleExtension(drv, ext):
     73      exts = drv.GetMetadataItem(gdal.DMD_EXTENSIONS)
     74      return exts is not None and exts.lower().find(ext.lower()) >= 0
     75
     76
     77 -def GetExtension(filename):
     78      ext = os.path.splitext(filename)[1]
     79      if ext.startswith('.'):
     80          ext = ext[1:]
     81      return ext
     82
     83
     84 -def GetOutputDriversFor(filename):
     85      drv_list = []
     86      ext = GetExtension(filename)
     87      for i in range(gdal.GetDriverCount()):
     88          drv = gdal.GetDriver(i)
     89          if (drv.GetMetadataItem(gdal.DCAP_CREATE) is not None or
     90              drv.GetMetadataItem(gdal.DCAP_CREATECOPY) is not None) and \
     91             drv.GetMetadataItem(gdal.DCAP_RASTER) is not None:
     92              if ext and DoesDriverHandleExtension(drv, ext):
     93                  drv_list.append(drv.ShortName)
     94              else:
     95                  prefix = drv.GetMetadataItem(gdal.DMD_CONNECTION_PREFIX)
     96                  if prefix is not None and filename.lower().startswith(prefix.lower()):
     97                      drv_list.append(drv.ShortName)
     98
     99      # GMT is registered before netCDF for opening reasons, but we want
    100      # netCDF to be used by default for output.
    101      if ext.lower() == 'nc' and not drv_list and \
    102         drv_list[0].upper() == 'GMT' and drv_list[1].upper() == 'NETCDF':
    103          drv_list = ['NETCDF', 'GMT']
    104
    105      return drv_list
    106
    107
    108 -def GetOutputDriverFor(filename):
    109      drv_list = GetOutputDriversFor(filename)
    110      ext = GetExtension(filename)
    111      if not drv_list:
    112          if not ext:
    113              return 'GTiff'
    114          else:
    115              raise Exception("Cannot guess driver for %s" % filename)
    116      elif len(drv_list) > 1:
    117          print("Several drivers matching %s extension. Using %s" % (ext if ext else '', drv_list[0]))
    118      return drv_list[0]
    119
    120  ################################################################
    121
    122
    123 -def doit(opts, args):
    124      # pylint: disable=unused-argument
    125
    126      if opts.debug:
    127          print("gdal_calc.py starting calculation %s" % (opts.calc))
    128
    129      # set up global namespace for eval with all functions of gdalnumeric
    130      global_namespace = dict([(key, getattr(gdalnumeric, key))
    131                               for key in dir(gdalnumeric) if not key.startswith('__')])
    132
    133      if not opts.calc:
    134          raise Exception("No calculation provided.")
    135      elif not opts.outF:
    136          raise Exception("No output file provided.")
    137
    138      if opts.format is None:
    139          opts.format = GetOutputDriverFor(opts.outF)
    140
    141      ################################################################
    142      # fetch details of input layers
    143      ################################################################
    144
    145      # set up some lists to store data for each band
    146      myFiles = []
    147      myBands = []
    148      myAlphaList = []
    149      myDataType = []
    150      myDataTypeNum = []
    151      myNDV = []
    152      DimensionsCheck = None
    153
    154      # loop through input files - checking dimensions
    155      for myI, myF in opts.input_files.items():
    156          if not myI.endswith("_band"):
    157              # check if we have asked for a specific band...
    158              if "%s_band" % myI in opts.input_files:
    159                  myBand = opts.input_files["%s_band" % myI]
    160              else:
    161                  myBand = 1
    162
    163              myFile = gdal.Open(myF, gdal.GA_ReadOnly)
    164              if not myFile:
    165                  raise IOError("No such file or directory: '%s'" % myF)
    166
    167              myFiles.append(myFile)
    168              myBands.append(myBand)
    169              myAlphaList.append(myI)
    170              myDataType.append(gdal.GetDataTypeName(myFile.GetRasterBand(myBand).DataType))
    171              myDataTypeNum.append(myFile.GetRasterBand(myBand).DataType)
    172              myNDV.append(myFile.GetRasterBand(myBand).GetNoDataValue())
    173              # check that the dimensions of each layer are the same
    174              if DimensionsCheck:
    175                  if DimensionsCheck != [myFile.RasterXSize, myFile.RasterYSize]:
    176                      raise Exception("Error! Dimensions of file %s (%i, %i) are different from other files (%i, %i).  Cannot proceed" %
    177                                      (myF, myFile.RasterXSize, myFile.RasterYSize, DimensionsCheck[0], DimensionsCheck[1]))
    178              else:
    179                  DimensionsCheck = [myFile.RasterXSize, myFile.RasterYSize]
    180
    181              if opts.debug:
    182                  print("file %s: %s, dimensions: %s, %s, type: %s" % (myI, myF, DimensionsCheck[0], DimensionsCheck[1], myDataType[-1]))
    183
    184      # process allBands option
    185      allBandsIndex = None
    186      allBandsCount = 1
    187      if opts.allBands:
    188          if len(opts.calc) > 1:
    189              raise Exception("Error! --allBands implies a single --calc")
    190          try:
    191              allBandsIndex = myAlphaList.index(opts.allBands)
    192          except ValueError:
    193              raise Exception("Error! allBands option was given but Band %s not found.  Cannot proceed" % (opts.allBands))
    194          allBandsCount = myFiles[allBandsIndex].RasterCount
    195          if allBandsCount <= 1:
    196              allBandsIndex = None
    197      else:
    198          allBandsCount = len(opts.calc)
    199
    200      ################################################################
    201      # set up output file
    202      ################################################################
    203
    204      # open output file exists
    205      if os.path.isfile(opts.outF) and not opts.overwrite:
    206          if allBandsIndex is not None:
    207              raise Exception("Error! allBands option was given but Output file exists, must use --overwrite option!")
    208          if len(opts.calc) > 1:
    209              raise Exception("Error! multiple calc options were given but Output file exists, must use --overwrite option!")
    210          if opts.debug:
    211              print("Output file %s exists - filling in results into file" % (opts.outF))
    212          myOut = gdal.Open(opts.outF, gdal.GA_Update)
    213          if [myOut.RasterXSize, myOut.RasterYSize] != DimensionsCheck:
    214              raise Exception("Error! Output exists, but is the wrong size.  Use the --overwrite option to automatically overwrite the existing file")
    215          myOutB = myOut.GetRasterBand(1)
    216          myOutNDV = myOutB.GetNoDataValue()
    217          myOutType = gdal.GetDataTypeName(myOutB.DataType)
    218
    219      else:
    220          # remove existing file and regenerate
    221          if os.path.isfile(opts.outF):
    222              os.remove(opts.outF)
    223          # create a new file
    224          if opts.debug:
    225              print("Generating output file %s" % (opts.outF))
    226
    227          # find data type to use
    228          if not opts.type:
    229              # use the largest type of the input files
    230              myOutType = gdal.GetDataTypeName(max(myDataTypeNum))
    231          else:
    232              myOutType = opts.type
    233
    234          # create file
    235          myOutDrv = gdal.GetDriverByName(opts.format)
    236          myOut = myOutDrv.Create(
    237              opts.outF, DimensionsCheck[0], DimensionsCheck[1], allBandsCount,
    238              gdal.GetDataTypeByName(myOutType), opts.creation_options)
    239
    240          # set output geo info based on first input layer
    241          myOut.SetGeoTransform(myFiles[0].GetGeoTransform())
    242          myOut.SetProjection(myFiles[0].GetProjection())
    243
    244          if opts.NoDataValue is not None:
    245              myOutNDV = opts.NoDataValue
    246          else:
    247              myOutNDV = DefaultNDVLookup[myOutType]
    248
    249          for i in range(1, allBandsCount + 1):
    250              myOutB = myOut.GetRasterBand(i)
    251              myOutB.SetNoDataValue(myOutNDV)
    252              # write to band
    253              myOutB = None
    254
    255      if opts.debug:
    256          print("output file: %s, dimensions: %s, %s, type: %s" % (opts.outF, myOut.RasterXSize, myOut.RasterYSize, myOutType))
    257
    258      ################################################################
    259      # find block size to chop grids into bite-sized chunks
    260      ################################################################
    261
    262      # use the block size of the first layer to read efficiently
    263      myBlockSize = myFiles[0].GetRasterBand(myBands[0]).GetBlockSize()
    264      # find total x and y blocks to be read
    265      nXBlocks = (int)((DimensionsCheck[0] + myBlockSize[0] - 1) / myBlockSize[0])
    266      nYBlocks = (int)((DimensionsCheck[1] + myBlockSize[1] - 1) / myBlockSize[1])
    267      myBufSize = myBlockSize[0] * myBlockSize[1]
    268
    269      if opts.debug:
    270          print("using blocksize %s x %s" % (myBlockSize[0], myBlockSize[1]))
    271
    272      # variables for displaying progress
    273      ProgressCt = -1
    274      ProgressMk = -1
    275      ProgressEnd = nXBlocks * nYBlocks * allBandsCount
    276
    277      ################################################################
    278      # start looping through each band in allBandsCount
    279      ################################################################
    280
    281      for bandNo in range(1, allBandsCount + 1):
    282
    283          ################################################################
    284          # start looping through blocks of data
    285          ################################################################
    286
    287          # store these numbers in variables that may change later
    288          nXValid = myBlockSize[0]
    289          nYValid = myBlockSize[1]
    290
    291          # loop through X-lines
    292          for X in range(0, nXBlocks):
    293
    294              # in case the blocks don't fit perfectly
    295              # change the block size of the final piece
    296              if X == nXBlocks - 1:
    297                  nXValid = DimensionsCheck[0] - X * myBlockSize[0]
    298
    299              # find X offset
    300              myX = X * myBlockSize[0]
    301
    302              # reset buffer size for start of Y loop
    303              nYValid = myBlockSize[1]
    304              myBufSize = nXValid * nYValid
    305
    306              # loop through Y lines
    307              for Y in range(0, nYBlocks):
    308                  ProgressCt += 1
    309                  if 10 * ProgressCt / ProgressEnd % 10 != ProgressMk and not opts.quiet:
    310                      ProgressMk = 10 * ProgressCt / ProgressEnd % 10
    311                      from sys import version_info
    312                      if version_info >= (3, 0, 0):
    313                          exec('print("%d.." % (10*ProgressMk), end=" ")')
    314                      else:
    315                          exec('print 10*ProgressMk, "..",')
    316
    317                  # change the block size of the final piece
    318                  if Y == nYBlocks - 1:
    319                      nYValid = DimensionsCheck[1] - Y * myBlockSize[1]
    320                      myBufSize = nXValid * nYValid
    321
    322                  # find Y offset
    323                  myY = Y * myBlockSize[1]
    324
    325                  # create empty buffer to mark where nodata occurs
    326                  myNDVs = None
    327
    328                  # make local namespace for calculation
    329                  local_namespace = {}
    330
    331                  # fetch data for each input layer
    332                  for i, Alpha in enumerate(myAlphaList):
    333
    334                      # populate lettered arrays with values
    335                      if allBandsIndex is not None and allBandsIndex == i:
    336                          myBandNo = bandNo
    337                      else:
    338                          myBandNo = myBands[i]
    339                      myval = gdalnumeric.BandReadAsArray(myFiles[i].GetRasterBand(myBandNo),
    340                                                          xoff=myX, yoff=myY,
    341                                                          win_xsize=nXValid, win_ysize=nYValid)
    342                      if myval is None:
    343                          raise Exception('Input block reading failed')
    344
    345                      # fill in nodata values
    346                      if myNDV[i] is not None:
    347                          if myNDVs is None:
    348                              myNDVs = numpy.zeros(myBufSize)
    349                              myNDVs.shape = (nYValid, nXValid)
    350                          myNDVs = 1 * numpy.logical_or(myNDVs == 1, myval == myNDV[i])
    351
    352                      # add an array of values for this block to the eval namespace
    353                      local_namespace[Alpha] = myval
    354                      myval = None
    355
    356                  # try the calculation on the array blocks
    357                  calc = opts.calc[bandNo-1 if len(opts.calc) > 1 else 0]
    358                  try:
    359                      myResult = eval(calc, global_namespace, local_namespace)
    360                  except:
    361                      print("evaluation of calculation %s failed" % (calc))
    362                      raise
    363
    364                  # Propagate nodata values (set nodata cells to zero
    365                  # then add nodata value to these cells).
    366                  if myNDVs is not None:
    367                      myResult = ((1 * (myNDVs == 0)) * myResult) + (myOutNDV * myNDVs)
    368                  elif not isinstance(myResult, numpy.ndarray):
    369                      myResult = numpy.ones((nYValid, nXValid)) * myResult
    370
    371                  # write data block to the output file
    372                  myOutB = myOut.GetRasterBand(bandNo)
    373                  if gdalnumeric.BandWriteArray(myOutB, myResult, xoff=myX, yoff=myY) != 0:
    374                      raise Exception('Block writing failed')
    375
    376      gdal.ErrorReset()
    377      myOut.FlushCache()
    378      myOut = None
    379      if gdal.GetLastErrorMsg() != '':
    380          raise Exception('Dataset writing failed')
    381
    382      if not opts.quiet:
    383          print("100 - Done")
    384
    385  ################################################################
    386
    387
    388 -def Calc(calc, outfile, NoDataValue=None, type=None, format=None, creation_options=None, allBands='', overwrite=False, debug=False, quiet=False, **input_files):
    389      """ Perform raster calculations with numpy syntax.
    390      Use any basic arithmetic supported by numpy arrays such as +-* along with logical
    391      operators such as >. Note that all files must have the same dimensions, but no projection checking is performed.
    392
    393      Keyword arguments:
    394          [A-Z]: input files
    395          [A_band - Z_band]: band to use for respective input file
    396
    397      Examples:
    398      add two files together:
    399          Calc("A+B", A="input1.tif", B="input2.tif", outfile="result.tif")
    400
    401      average of two layers:
    402          Calc(calc="(A+B)/2", A="input1.tif", B="input2.tif", outfile="result.tif")
    403
    404      set values of zero and below to null:
    405          Calc(calc="A*(A>0)", A="input.tif", A_band=2, outfile="result.tif", NoDataValue=0)
    406
    407      work with two bands:
    408          Calc(["(A+B)/2", "A*(A>0)"], A="input.tif", A_band=1, B="input.tif", B_band=2, outfile="result.tif", NoDataValue=0)
    409      """
    410      opts = Values()
    411      opts.input_files = input_files
    412      # Single calc value compatibility
    413      # (type is overridden in the parameter list)
    414      if isinstance(calc, list):
    415          opts.calc = calc
    416      else:
    417          opts.calc = [calc]
    418      opts.outF = outfile
    419      opts.NoDataValue = NoDataValue
    420      opts.type = type
    421      opts.format = format
    422      opts.creation_options = [] if creation_options is None else creation_options
    423      opts.allBands = allBands
    424      opts.overwrite = overwrite
    425      opts.debug = debug
    426      opts.quiet = quiet
    427
    428      doit(opts, None)
    429
    430
    431 -def store_input_file(option, opt_str, value, parser):
    432      # pylint: disable=unused-argument
    433      if not hasattr(parser.values, 'input_files'):
    434          parser.values.input_files = {}
    435      parser.values.input_files[opt_str.lstrip('-')] = value
    436
    437
    438 -def add_alpha_args(parser, argv):
    439      # limit the input file options to the ones in the argument list
    440      given_args = set([a[1] for a in argv if a[1:2] in AlphaList] + ['A'])
    441      for myAlpha in given_args:
    442          try:
    443              parser.add_option("-%s" % myAlpha, action="callback", callback=store_input_file, type=str, help="input gdal raster file, you can use any letter (A-Z)", metavar='filename')
    444              parser.add_option("--%s_band" % myAlpha, action="callback", callback=store_input_file, type=int, help="number of raster band for file %s (default 1)" % myAlpha, metavar='n')
    445          except OptionConflictError:
    446              pass
    447
    448
    449 -def main(argv):
    450      usage = """usage: %prog --calc=expression --outfile=out_filename [-A filename]
    451                      [--A_band=n] [-B...-Z filename] [other_options]"""
    452      parser = OptionParser(usage)
    453
    454      # define options
    455      parser.add_option("--calc", dest="calc", action="append", help="calculation in gdalnumeric syntax using +-/* or any numpy array functions (i.e. log10()). May appear multiple times to produce a multi-band file", metavar="expression")
    456      add_alpha_args(parser, argv)
    457
    458      parser.add_option("--outfile", dest="outF", help="output file to generate or fill", metavar="filename")
    459      parser.add_option("--NoDataValue", dest="NoDataValue", type=float, help="output nodata value (default datatype specific value)", metavar="value")
    460      parser.add_option("--type", dest="type", help="output datatype, must be one of %s" % list(DefaultNDVLookup.keys()), metavar="datatype")
    461      parser.add_option("--format", dest="format", help="GDAL format for output file", metavar="gdal_format")
    462      parser.add_option(
    463          "--creation-option", "--co", dest="creation_options", default=[], action="append",
    464          help="Passes a creation option to the output format driver. Multiple "
    465          "options may be listed. See format specific documentation for legal "
    466          "creation options for each format.", metavar="option")
    467      parser.add_option("--allBands", dest="allBands", default="", help="process all bands of given raster (A-Z)", metavar="[A-Z]")
    468      parser.add_option("--overwrite", dest="overwrite", action="store_true", help="overwrite output file if it already exists")
    469      parser.add_option("--debug", dest="debug", action="store_true", help="print debugging information")
    470      parser.add_option("--quiet", dest="quiet", action="store_true", help="suppress progress messages")
    471      parser.add_option("--optfile", dest="optfile", metavar="optfile", help="Read the named file and substitute the contents into the command line options list.")
    472
    473      (opts, args) = parser.parse_args(argv[1:])
    474
    475      if not hasattr(opts, "input_files"):
    476          opts.input_files = {}
    477
    478      if opts.optfile:
    479          with open(opts.optfile, 'r') as f:
    480              ofargv = [x for line in f for x in shlex.split(line, comments=True)]
    481          # Avoid potential recursion.
    482          parser.remove_option('--optfile')
    483          add_alpha_args(parser, ofargv)
    484          ofopts, ofargs = parser.parse_args(ofargv)
    485          # Let options given directly override the optfile.
    486          input_files = getattr(ofopts, 'input_files', {})
    487          input_files.update(opts.input_files)
    488          ofopts.__dict__.update({k: v for k, v in vars(opts).items() if v})
    489          opts = ofopts
    490          opts.input_files = input_files
    491          args = args + ofargs
    492
    493      if len(argv) == 1:
    494          parser.print_help()
    495          sys.exit(1)
    496      elif not opts.calc:
    497          print("No calculation provided. Nothing to do!")
    498          parser.print_help()
    499          sys.exit(1)
    500      elif not opts.outF:
    501          print("No output file provided. Cannot proceed.")
    502          parser.print_help()
    503          sys.exit(1)
    504      else:
    505          try:
    506              doit(opts, args)
    507          except IOError as e:
    508              print(e)
    509              sys.exit(1)
    510
    511
    512  if __name__ == '__main__':
    513      sys.exit(main(sys.argv))
    514