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

# Source Code for Module osgeo.utils.pct2rgb

      1  #!/usr/bin/env python3
      2  # ******************************************************************************
      3  #  $Id$
      4  #
      5  #  Name:     pct2rgb
      6  #  Project:  GDAL Python Interface
      7  #  Purpose:  Utility to convert paletted images into RGB (or RGBA) images.
      8  #  Author:   Frank Warmerdam, warmerdam@pobox.com
      9  #
     10  # ******************************************************************************
     11  #  Copyright (c) 2001, Frank Warmerdam
     12  #  Copyright (c) 2009-2010, 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
     38  progress = gdal.TermProgress_nocb
     39
     40  try:
     41      import numpy as Numeric
     42      Numeric.arrayrange = Numeric.arange
     43  except ImportError:
     44      import Numeric
     45
     46
     47 -def Usage():
     48      print('Usage: pct2rgb.py [-of format] [-b <band>] [-rgba] source_file dest_file')
     49      sys.exit(1)
     50
     51
     52 -def DoesDriverHandleExtension(drv, ext):
     53      exts = drv.GetMetadataItem(gdal.DMD_EXTENSIONS)
     54      return exts is not None and exts.lower().find(ext.lower()) >= 0
     55
     56
     57 -def GetExtension(filename):
     58      ext = os.path.splitext(filename)[1]
     59      if ext.startswith('.'):
     60          ext = ext[1:]
     61      return ext
     62
     63
     64 -def GetOutputDriversFor(filename):
     65      drv_list = []
     66      ext = GetExtension(filename)
     67      for i in range(gdal.GetDriverCount()):
     68          drv = gdal.GetDriver(i)
     69          if (drv.GetMetadataItem(gdal.DCAP_CREATE) is not None or
     70              drv.GetMetadataItem(gdal.DCAP_CREATECOPY) is not None) and \
     71             drv.GetMetadataItem(gdal.DCAP_RASTER) is not None:
     72              if ext and DoesDriverHandleExtension(drv, ext):
     73                  drv_list.append(drv.ShortName)
     74              else:
     75                  prefix = drv.GetMetadataItem(gdal.DMD_CONNECTION_PREFIX)
     76                  if prefix is not None and filename.lower().startswith(prefix.lower()):
     77                      drv_list.append(drv.ShortName)
     78
     79      # GMT is registered before netCDF for opening reasons, but we want
     80      # netCDF to be used by default for output.
     81      if ext.lower() == 'nc' and not drv_list and \
     82         drv_list[0].upper() == 'GMT' and drv_list[1].upper() == 'NETCDF':
     83          drv_list = ['NETCDF', 'GMT']
     84
     85      return drv_list
     86
     87
     88 -def GetOutputDriverFor(filename):
     89      drv_list = GetOutputDriversFor(filename)
     90      ext = GetExtension(filename)
     91      if not drv_list:
     92          if not ext:
     93              return 'GTiff'
     94          else:
     95              raise Exception("Cannot guess driver for %s" % filename)
     96      elif len(drv_list) > 1:
     97          print("Several drivers matching %s extension. Using %s" % (ext if ext else '', drv_list[0]))
     98      return drv_list[0]
     99
    100
    101 -def main(argv):
    102      frmt = None
    103      src_filename = None
    104      dst_filename = None
    105      out_bands = 3
    106      band_number = 1
    107
    108      gdal.AllRegister()
    109      argv = gdal.GeneralCmdLineProcessor(argv)
    110      if argv is None:
    111          sys.exit(0)
    112
    113      # Parse command line arguments.
    114      i = 1
    115      while i < len(argv):
    116          arg = argv[i]
    117
    118          if arg == '-of' or arg == '-f':
    119              i = i + 1
    120              frmt = argv[i]
    121
    122          elif arg == '-b':
    123              i = i + 1
    124              band_number = int(argv[i])
    125
    126          elif arg == '-rgba':
    127              out_bands = 4
    128
    129          elif src_filename is None:
    130              src_filename = argv[i]
    131
    132          elif dst_filename is None:
    133              dst_filename = argv[i]
    134
    135          else:
    136              Usage()
    137
    138          i = i + 1
    139
    140      if dst_filename is None:
    141          Usage()
    142
    143      # ----------------------------------------------------------------------------
    144      # Open source file
    145
    146      src_ds = gdal.Open(src_filename)
    147      if src_ds is None:
    148          print('Unable to open %s ' % src_filename)
    149          sys.exit(1)
    150
    151      src_band = src_ds.GetRasterBand(band_number)
    152
    153      # ----------------------------------------------------------------------------
    154      # Ensure we recognise the driver.
    155
    156      if frmt is None:
    157          frmt = GetOutputDriverFor(dst_filename)
    158
    159      dst_driver = gdal.GetDriverByName(frmt)
    160      if dst_driver is None:
    161          print('"%s" driver not registered.' % frmt)
    162          sys.exit(1)
    163
    164      # ----------------------------------------------------------------------------
    165      # Build color table.
    166
    167      ct = src_band.GetRasterColorTable()
    168
    169      ct_size = ct.GetCount()
    170      lookup = [Numeric.arrayrange(ct_size),
    171                Numeric.arrayrange(ct_size),
    172                Numeric.arrayrange(ct_size),
    173                Numeric.ones(ct_size) * 255]
    174
    175      if ct is not None:
    176          for i in range(ct_size):
    177              entry = ct.GetColorEntry(i)
    178              for c in range(4):
    179                  lookup[c][i] = entry[c]
    180
    181      # ----------------------------------------------------------------------------
    182      # Create the working file.
    183
    184      if frmt == 'GTiff':
    185          tif_filename = dst_filename
    186      else:
    187          tif_filename = 'temp.tif'
    188
    189      gtiff_driver = gdal.GetDriverByName('GTiff')
    190
    191      tif_ds = gtiff_driver.Create(tif_filename,
    192                                   src_ds.RasterXSize, src_ds.RasterYSize, out_bands)
    193
    194
    195      # ----------------------------------------------------------------------------
    196      # We should copy projection information and so forth at this point.
    197
    198      tif_ds.SetProjection(src_ds.GetProjection())
    199      tif_ds.SetGeoTransform(src_ds.GetGeoTransform())
    200      if src_ds.GetGCPCount() > 0:
    201          tif_ds.SetGCPs(src_ds.GetGCPs(), src_ds.GetGCPProjection())
    202
    203      # ----------------------------------------------------------------------------
    204      # Do the processing one scanline at a time.
    205
    206      progress(0.0)
    207      for iY in range(src_ds.RasterYSize):
    208          src_data = src_band.ReadAsArray(0, iY, src_ds.RasterXSize, 1)
    209
    210          for iBand in range(out_bands):
    211              band_lookup = lookup[iBand]
    212
    213              dst_data = Numeric.take(band_lookup, src_data)
    214              tif_ds.GetRasterBand(iBand + 1).WriteArray(dst_data, 0, iY)
    215
    216          progress((iY + 1.0) / src_ds.RasterYSize)
    217
    218
    219      tif_ds = None
    220
    221      # ----------------------------------------------------------------------------
    222      # Translate intermediate file to output format if desired format is not TIFF.
    223
    224      if tif_filename != dst_filename:
    225          tif_ds = gdal.Open(tif_filename)
    226          dst_driver.CreateCopy(dst_filename, tif_ds)
    227          tif_ds = None
    228
    229          gtiff_driver.Delete(tif_filename)
    230
    231
    232  if __name__ == '__main__':
    233      sys.exit(main(sys.argv))
    234