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

# Source Code for Module osgeo.utils.rgb2pct

      1  #!/usr/bin/env python3
      2  # ******************************************************************************
      3  #  $Id$
      4  #
      5  #  Name:     rgb2pct
      6  #  Project:  GDAL Python Interface
      7  #  Purpose:  Application for converting an RGB image to a pseudocolored image.
      8  #  Author:   Frank Warmerdam, warmerdam@pobox.com
      9  #
     10  # ******************************************************************************
     11  #  Copyright (c) 2001, Frank Warmerdam
     12  #
     13  #  Permission is hereby granted, free of charge, to any person obtaining a
     14  #  copy of this software and associated documentation files (the "Software"),
     15  #  to deal in the Software without restriction, including without limitation
     16  #  the rights to use, copy, modify, merge, publish, distribute, sublicense,
     17  #  and/or sell copies of the Software, and to permit persons to whom the
     18  #  Software is furnished to do so, subject to the following conditions:
     19  #
     20  #  The above copyright notice and this permission notice shall be included
     21  #  in all copies or substantial portions of the Software.
     22  #
     23  #  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
     24  #  OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
     25  #  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
     26  #  THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
     27  #  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
     28  #  FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
     29  #  DEALINGS IN THE SOFTWARE.
     30  # ******************************************************************************
     31
     32  import os.path
     33  import sys
     34
     35  from osgeo import gdal
     36
     37
     38 -def Usage():
     39      print('Usage: rgb2pct.py [-n colors | -pct palette_file] [-of format] source_file dest_file')
     40      sys.exit(1)
     41
     42
     43 -def DoesDriverHandleExtension(drv, ext):
     44      exts = drv.GetMetadataItem(gdal.DMD_EXTENSIONS)
     45      return exts is not None and exts.lower().find(ext.lower()) >= 0
     46
     47
     48 -def GetExtension(filename):
     49      ext = os.path.splitext(filename)[1]
     50      if ext.startswith('.'):
     51          ext = ext[1:]
     52      return ext
     53
     54
     55 -def GetOutputDriversFor(filename):
     56      drv_list = []
     57      ext = GetExtension(filename)
     58      for i in range(gdal.GetDriverCount()):
     59          drv = gdal.GetDriver(i)
     60          if (drv.GetMetadataItem(gdal.DCAP_CREATE) is not None or
     61              drv.GetMetadataItem(gdal.DCAP_CREATECOPY) is not None) and \
     62             drv.GetMetadataItem(gdal.DCAP_RASTER) is not None:
     63              if ext and DoesDriverHandleExtension(drv, ext):
     64                  drv_list.append(drv.ShortName)
     65              else:
     66                  prefix = drv.GetMetadataItem(gdal.DMD_CONNECTION_PREFIX)
     67                  if prefix is not None and filename.lower().startswith(prefix.lower()):
     68                      drv_list.append(drv.ShortName)
     69
     70      # GMT is registered before netCDF for opening reasons, but we want
     71      # netCDF to be used by default for output.
     72      if ext.lower() == 'nc' and not drv_list and \
     73         drv_list[0].upper() == 'GMT' and drv_list[1].upper() == 'NETCDF':
     74          drv_list = ['NETCDF', 'GMT']
     75
     76      return drv_list
     77
     78
     79 -def GetOutputDriverFor(filename):
     80      drv_list = GetOutputDriversFor(filename)
     81      ext = GetExtension(filename)
     82      if not drv_list:
     83          if not ext:
     84              return 'GTiff'
     85          else:
     86              raise Exception("Cannot guess driver for %s" % filename)
     87      elif len(drv_list) > 1:
     88          print("Several drivers matching %s extension. Using %s" % (ext if ext else '', drv_list[0]))
     89      return drv_list[0]
     90
     91
     92 -def main(argv):
     93      color_count = 256
     94      frmt = None
     95      src_filename = None
     96      dst_filename = None
     97      pct_filename = None
     98
     99      gdal.AllRegister()
    100      argv = gdal.GeneralCmdLineProcessor(argv)
    101      if argv is None:
    102          sys.exit(0)
    103
    104      # Parse command line arguments.
    105      i = 1
    106      while i < len(argv):
    107          arg = argv[i]
    108
    109          if arg == '-of' or arg == '-f':
    110              i = i + 1
    111              frmt = argv[i]
    112
    113          elif arg == '-n':
    114              i = i + 1
    115              color_count = int(argv[i])
    116
    117          elif arg == '-pct':
    118              i = i + 1
    119              pct_filename = argv[i]
    120
    121          elif src_filename is None:
    122              src_filename = argv[i]
    123
    124          elif dst_filename is None:
    125              dst_filename = argv[i]
    126
    127          else:
    128              Usage()
    129
    130          i = i + 1
    131
    132      if dst_filename is None:
    133          Usage()
    134
    135      # Open source file
    136
    137      src_ds = gdal.Open(src_filename)
    138      if src_ds is None:
    139          print('Unable to open %s' % src_filename)
    140          sys.exit(1)
    141
    142      if src_ds.RasterCount < 3:
    143          print('%s has %d band(s), need 3 for inputs red, green and blue.'
    144                % (src_filename, src_ds.RasterCount))
    145          sys.exit(1)
    146
    147      # Ensure we recognise the driver.
    148
    149      if frmt is None:
    150          frmt = GetOutputDriverFor(dst_filename)
    151
    152      dst_driver = gdal.GetDriverByName(frmt)
    153      if dst_driver is None:
    154          print('"%s" driver not registered.' % frmt)
    155          sys.exit(1)
    156
    157      # Generate palette
    158
    159      ct = gdal.ColorTable()
    160      if pct_filename is None:
    161          err = gdal.ComputeMedianCutPCT(src_ds.GetRasterBand(1),
    162                                         src_ds.GetRasterBand(2),
    163                                         src_ds.GetRasterBand(3),
    164                                         color_count, ct,
    165                                         callback=gdal.TermProgress_nocb)
    166      else:
    167          pct_ds = gdal.Open(pct_filename)
    168          ct = pct_ds.GetRasterBand(1).GetRasterColorTable().Clone()
    169
    170      # Create the working file.  We have to use TIFF since there are few formats
    171      # that allow setting the color table after creation.
    172
    173      if format == 'GTiff':
    174          tif_filename = dst_filename
    175      else:
    176          import tempfile
    177          tif_filedesc, tif_filename = tempfile.mkstemp(suffix='.tif')
    178
    179      gtiff_driver = gdal.GetDriverByName('GTiff')
    180
    181      tif_ds = gtiff_driver.Create(tif_filename,
    182                                   src_ds.RasterXSize, src_ds.RasterYSize, 1)
    183
    184      tif_ds.GetRasterBand(1).SetRasterColorTable(ct)
    185
    186      # ----------------------------------------------------------------------------
    187      # We should copy projection information and so forth at this point.
    188
    189      tif_ds.SetProjection(src_ds.GetProjection())
    190      tif_ds.SetGeoTransform(src_ds.GetGeoTransform())
    191      if src_ds.GetGCPCount() > 0:
    192          tif_ds.SetGCPs(src_ds.GetGCPs(), src_ds.GetGCPProjection())
    193
    194      # ----------------------------------------------------------------------------
    195      # Actually transfer and dither the data.
    196
    197      err = gdal.DitherRGB2PCT(src_ds.GetRasterBand(1),
    198                               src_ds.GetRasterBand(2),
    199                               src_ds.GetRasterBand(3),
    200                               tif_ds.GetRasterBand(1),
    201                               ct,
    202                               callback=gdal.TermProgress_nocb)
    203
    204      tif_ds = None
    205
    206      if tif_filename != dst_filename:
    207          tif_ds = gdal.Open(tif_filename)
    208          dst_driver.CreateCopy(dst_filename, tif_ds)
    209          tif_ds = None
    210
    211          os.close(tif_filedesc)
    212          gtiff_driver.Delete(tif_filename)
    213
    214      return err
    215
    216
    217  if __name__ == '__main__':
    218      sys.exit(main(sys.argv))
    219