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