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

# Source Code for Module osgeo.utils.gdal_polygonize

      1  #!/usr/bin/env python3
      2  # -*- coding: utf-8 -*-
      3  # ******************************************************************************
      4  #  $Id$
      5  #
      6  #  Project:  GDAL Python Interface
      7  #  Purpose:  Application for converting raster data to a vector polygon layer.
      8  #  Author:   Frank Warmerdam, warmerdam@pobox.com
      9  #
     10  # ******************************************************************************
     11  #  Copyright (c) 2008, Frank Warmerdam
     12  #  Copyright (c) 2009-2013, Even Rouault <even dot rouault at spatialys.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  import os.path
     34  import sys
     35
     36  from osgeo import gdal
     37  from osgeo import ogr
     38
     39
     40 -def Usage():
     41      print("""
     42  gdal_polygonize [-8] [-nomask] [-mask filename] raster_file [-b band|mask]
     43                  [-q] [-f ogr_format] out_file [layer] [fieldname]
     44  """)
     45      sys.exit(1)
     46
     47
     48 -def DoesDriverHandleExtension(drv, ext):
     49      exts = drv.GetMetadataItem(gdal.DMD_EXTENSIONS)
     50      return exts is not None and exts.lower().find(ext.lower()) >= 0
     51
     52
     53 -def GetExtension(filename):
     54      ext = os.path.splitext(filename)[1]
     55      if ext.startswith('.'):
     56          ext = ext[1:]
     57      return ext
     58
     59
     60 -def GetOutputDriversFor(filename):
     61      drv_list = []
     62      ext = GetExtension(filename)
     63      for i in range(gdal.GetDriverCount()):
     64          drv = gdal.GetDriver(i)
     65          if (drv.GetMetadataItem(gdal.DCAP_CREATE) is not None or
     66              drv.GetMetadataItem(gdal.DCAP_CREATECOPY) is not None) and \
     67             drv.GetMetadataItem(gdal.DCAP_VECTOR) is not None:
     68              if ext and DoesDriverHandleExtension(drv, ext):
     69                  drv_list.append(drv.ShortName)
     70              else:
     71                  prefix = drv.GetMetadataItem(gdal.DMD_CONNECTION_PREFIX)
     72                  if prefix is not None and filename.lower().startswith(prefix.lower()):
     73                      drv_list.append(drv.ShortName)
     74
     75      return drv_list
     76
     77
     78 -def GetOutputDriverFor(filename):
     79      drv_list = GetOutputDriversFor(filename)
     80      ext = GetExtension(filename)
     81      if not drv_list:
     82          if not ext:
     83              return 'ESRI Shapefile'
     84          else:
     85              raise Exception("Cannot guess driver for %s" % filename)
     86      elif len(drv_list) > 1:
     87          print("Several drivers matching %s extension. Using %s" % (ext if ext else '', drv_list[0]))
     88      return drv_list[0]
     89
     90
     91 -def main(argv):
     92      frmt = None
     93      options = []
     94      quiet_flag = 0
     95      src_filename = None
     96      src_band_n = 1
     97
     98      dst_filename = None
     99      dst_layername = None
    100      dst_fieldname = None
    101      dst_field = -1
    102
    103      mask = 'default'
    104
    105      gdal.AllRegister()
    106      argv = gdal.GeneralCmdLineProcessor(argv)
    107      if argv is None:
    108          sys.exit(0)
    109
    110      # Parse command line arguments.
    111      i = 1
    112      while i < len(argv):
    113          arg = argv[i]
    114
    115          if arg == '-f' or arg == '-of':
    116              i = i + 1
    117              frmt = argv[i]
    118
    119          elif arg == '-q' or arg == '-quiet':
    120              quiet_flag = 1
    121
    122          elif arg == '-8':
    123              options.append('8CONNECTED=8')
    124
    125          elif arg == '-nomask':
    126              mask = 'none'
    127
    128          elif arg == '-mask':
    129              i = i + 1
    130              mask = argv[i]
    131
    132          elif arg == '-b':
    133              i = i + 1
    134              if argv[i].startswith('mask'):
    135                  src_band_n = argv[i]
    136              else:
    137                  src_band_n = int(argv[i])
    138
    139          elif src_filename is None:
    140              src_filename = argv[i]
    141
    142          elif dst_filename is None:
    143              dst_filename = argv[i]
    144
    145          elif dst_layername is None:
    146              dst_layername = argv[i]
    147
    148          elif dst_fieldname is None:
    149              dst_fieldname = argv[i]
    150
    151          else:
    152              Usage()
    153
    154          i = i + 1
    155
    156      if src_filename is None or dst_filename is None:
    157          Usage()
    158
    159      if frmt is None:
    160          frmt = GetOutputDriverFor(dst_filename)
    161
    162      if dst_layername is None:
    163          dst_layername = 'out'
    164
    165      # =============================================================================
    166      #   Verify we have next gen bindings with the polygonize method.
    167      # =============================================================================
    168      try:
    169          gdal.Polygonize
    170      except AttributeError:
    171          print('')
    172          print('gdal.Polygonize() not available.  You are likely using "old gen"')
    173          print('bindings or an older version of the next gen bindings.')
    174          print('')
    175          sys.exit(1)
    176
    177      # =============================================================================
    178      # Open source file
    179      # =============================================================================
    180
    181      src_ds = gdal.Open(src_filename)
    182
    183      if src_ds is None:
    184          print('Unable to open %s' % src_filename)
    185          sys.exit(1)
    186
    187      if src_band_n == 'mask':
    188          srcband = src_ds.GetRasterBand(1).GetMaskBand()
    189          # Workaround the fact that most source bands have no dataset attached
    190          options.append('DATASET_FOR_GEOREF=' + src_filename)
    191      elif isinstance(src_band_n, str) and src_band_n.startswith('mask,'):
    192          srcband = src_ds.GetRasterBand(int(src_band_n[len('mask,'):])).GetMaskBand()
    193          # Workaround the fact that most source bands have no dataset attached
    194          options.append('DATASET_FOR_GEOREF=' + src_filename)
    195      else:
    196          srcband = src_ds.GetRasterBand(src_band_n)
    197
    198      if mask == 'default':
    199          maskband = srcband.GetMaskBand()
    200      elif mask == 'none':
    201          maskband = None
    202      else:
    203          mask_ds = gdal.Open(mask)
    204          maskband = mask_ds.GetRasterBand(1)
    205
    206      # =============================================================================
    207      #       Try opening the destination file as an existing file.
    208      # =============================================================================
    209
    210      try:
    211          gdal.PushErrorHandler('CPLQuietErrorHandler')
    212          dst_ds = ogr.Open(dst_filename, update=1)
    213          gdal.PopErrorHandler()
    214      except:
    215          dst_ds = None
    216
    217      # =============================================================================
    218      #   Create output file.
    219      # =============================================================================
    220      if dst_ds is None:
    221          drv = ogr.GetDriverByName(frmt)
    222          if not quiet_flag:
    223              print('Creating output %s of format %s.' % (dst_filename, frmt))
    224          dst_ds = drv.CreateDataSource(dst_filename)
    225
    226      # =============================================================================
    227      #       Find or create destination layer.
    228      # =============================================================================
    229      try:
    230          dst_layer = dst_ds.GetLayerByName(dst_layername)
    231      except:
    232          dst_layer = None
    233
    234      if dst_layer is None:
    235
    236          srs = src_ds.GetSpatialRef()
    237          dst_layer = dst_ds.CreateLayer(dst_layername, geom_type=ogr.wkbPolygon, srs=srs)
    238
    239          if dst_fieldname is None:
    240              dst_fieldname = 'DN'
    241
    242          fd = ogr.FieldDefn(dst_fieldname, ogr.OFTInteger)
    243          dst_layer.CreateField(fd)
    244          dst_field = 0
    245      else:
    246          if dst_fieldname is not None:
    247              dst_field = dst_layer.GetLayerDefn().GetFieldIndex(dst_fieldname)
    248              if dst_field < 0:
    249                  print("Warning: cannot find field '%s' in layer '%s'" % (dst_fieldname, dst_layername))
    250
    251      # =============================================================================
    252      # Invoke algorithm.
    253      # =============================================================================
    254
    255      if quiet_flag:
    256          prog_func = None
    257      else:
    258          prog_func = gdal.TermProgress_nocb
    259
    260      result = gdal.Polygonize(srcband, maskband, dst_layer, dst_field, options,
    261                               callback=prog_func)
    262
    263      srcband = None
    264      src_ds = None
    265      dst_ds = None
    266      mask_ds = None
    267
    268      return result
    269
    270
    271  if __name__ == '__main__':
    272      sys.exit(main(sys.argv))
    273