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