[gdal-dev] 4 band GeoTIFF to 3 band GeoTIFF

[abhay menon]

Skipped content of type multipart/alternative-------------- next part --------------
@python "%FWTOOLS_DIR%\bin\4bandto3band.py" %*
-------------- next part --------------
#!/usr/bin/env python
#******************************************************************************
#  Name:    4bandto3band
#  Purpose:  Utility to convert 4 band image to 3 band images.
#  Author:   Abhay Menon,abhay.menon at gmail.com
#******************************************************************************

import gdal
import sys
import Numeric
import os.path
import string

def Usage():
    print 'Usage: 4bandto3band.py [-of format] source_file dest_file [-worldfile]'
    sys.exit(1)

# =============================================================================
# 	Mainline
# =============================================================================

format = 'GTiff'
src_filename = None
dst_filename = None
out_bands = 3
band_number = 1
worldfile=0

gdal.AllRegister()
argv = gdal.GeneralCmdLineProcessor( sys.argv )
if argv is None:
    sys.exit( 0 )

# Parse command line arguments.
i = 1
while i < len(argv):
    arg = argv[i]

    if arg == '-of':
        i = i + 1
        format = argv[i]

    elif arg == '-worldfile':
       worldfile=1
		
    elif src_filename is None:
        src_filename = argv[i]

    elif dst_filename is None:
        dst_filename = argv[i]

    else:
        Usage()

    i = i + 1

if dst_filename is None:
    Usage()
    
# ----------------------------------------------------------------------------
# Open source file

src_ds = gdal.Open( src_filename )
if src_ds is None:
    print 'Unable to open ', src_filename
    sys.exit(1)

if worldfile is not None:
	geotransform = src_ds.GetGeoTransform()
	if not geotransform is None:
		wrldfile = dst_filename
		wrldfile = os.path.splitext(wrldfile)[0] + '.wld'
		ulx = geotransform[0]
		uly = geotransform[3]
		ulx = ulx + (geotransform[1]*0.5)
		uly = uly + (geotransform[5]*0.5)
		fileHandle = open ( wrldfile, 'w' )
		fileHandle.write ( "%.16f\n%.16f\n%.16f\n%.16f\n%.16f\n%.16f" % (geotransform[1],geotransform[2],geotransform[4],geotransform[5],ulx,uly) )
		fileHandle.close() 
	
src_band = src_ds.GetRasterBand(band_number)

# ----------------------------------------------------------------------------
# Ensure we recognise the driver.

dst_driver = gdal.GetDriverByName(format)
if dst_driver is None:
    print '"%s" driver not registered.' % format
    sys.exit(1)

# ----------------------------------------------------------------------------
# Build color table.

lookup = [ Numeric.arrayrange(256), 
           Numeric.arrayrange(256), 
           Numeric.arrayrange(256), 
           Numeric.ones(256)*255 ]

ct = src_band.GetRasterColorTable()

if ct is not None:
    for i in range(min(256,ct.GetCount())):
        entry = ct.GetColorEntry(i)
        for c in range(3):
            lookup[c][i] = entry[c]

# ----------------------------------------------------------------------------
# Create the working file.

if format == 'GTiff':
    tif_filename = dst_filename
else:
    tif_filename = 'temp.tif'

gtiff_driver = gdal.GetDriverByName( 'GTiff' )
						
tif_ds = gtiff_driver.Create( tif_filename,
                              src_ds.RasterXSize, src_ds.RasterYSize, out_bands )


# ----------------------------------------------------------------------------
# We should copy projection information and so forth at this point.

tif_ds.SetProjection( src_ds.GetProjection() )
tif_ds.SetGeoTransform( src_ds.GetGeoTransform() )

# ----------------------------------------------------------------------------
# Do the processing one scanline at a time. 
print '\nWriting Band 1'
gdal.TermProgress( 0.0 )
for iY in range(src_ds.RasterYSize):
    src_data = src_band.ReadAsArray(0,iY,src_ds.RasterXSize,1)

    dst_data = Numeric.take(lookup[0],src_data)
    tif_ds.GetRasterBand(1).WriteArray(dst_data,0,iY)

    gdal.TermProgress( (iY+1.0) / src_ds.RasterYSize )

print '\nWriting Band 2'
gdal.TermProgress( 0.0 )	

src_band = src_ds.GetRasterBand(2)

lookup = [ Numeric.arrayrange(256), 
           Numeric.arrayrange(256), 
           Numeric.arrayrange(256), 
           Numeric.ones(256)*255 ]

ct = src_band.GetRasterColorTable()

if ct is not None:
    for i in range(min(256,ct.GetCount())):
        entry = ct.GetColorEntry(i)
        for c in range(3):
            lookup[c][i] = entry[c]

for iY in range(src_ds.RasterYSize):
    src_data = src_band.ReadAsArray(0,iY,src_ds.RasterXSize,1)

    dst_data = Numeric.take(lookup[0],src_data)
    tif_ds.GetRasterBand(2).WriteArray(dst_data,0,iY)

    gdal.TermProgress( (iY+1.0) / src_ds.RasterYSize )
   
print '\nWriting Band 3'
gdal.TermProgress( 0.0 )	

src_band = src_ds.GetRasterBand(3)

lookup = [ Numeric.arrayrange(256), 
           Numeric.arrayrange(256), 
           Numeric.arrayrange(256), 
           Numeric.ones(256)*255 ]

ct = src_band.GetRasterColorTable()

if ct is not None:
    for i in range(min(256,ct.GetCount())):
        entry = ct.GetColorEntry(i)
        for c in range(3):
            lookup[c][i] = entry[c]

for iY in range(src_ds.RasterYSize):
    src_data = src_band.ReadAsArray(0,iY,src_ds.RasterXSize,1)

    dst_data = Numeric.take(lookup[0],src_data)
    tif_ds.GetRasterBand(3).WriteArray(dst_data,0,iY)

    gdal.TermProgress( (iY+1.0) / src_ds.RasterYSize )
	
tif_ds = None

# ----------------------------------------------------------------------------
# Translate intermediate file to output format if desired format is not TIFF.

if tif_filename <> dst_filename:
    tif_ds = gdal.Open( tif_filename )
    dst_driver.CreateCopy( dst_filename, tif_ds )
    tif_ds = None

    gtiff_driver.Delete( tif_filename )