pansharpen.py

#!/usr/bin/env python

# pansharpen.py pan.tif rgb.tif output.tif

# Takes a pan and an rgb image of identical spatial extent,
# makes the rgb one the same size, and writes out the combination.
# First draft: inelegant, inefficient, and generally not to be trusted.

# This works for me with Landsat 8 data:
# $ gdal_merge.py -separate $SCENE/*B{4,3,2}.TIF -o $SCENE/rgb.tif
# $ mkdir done
# $ python ~/Desktop/pansharpen.py $SCENE/*B8.TIF $SCENE/rgb.tif done/$SCENE.tif

# Then you may want to scale that to something subjectively reasonable, e.g. with:
# $ convert -channel B -gamma 0.97 \
#     -channel RGB -sigmoidal-contrast 40x15% done/$SCENE.tif done/$SCENE-pretty.tif
# But each scene will need its own trial and error. Note that `convert` strips geotags.

'''
todo:
- toa scaling built in
- mask using bqa (to avoid ugly edges and allow swath mosaics)
- less incredibly RAM-hungry
- rasterio (should help with RAM)
'''

import numpy as np
from sys import argv, exit
from osgeo import gdal, gdalconst
import os
os.environ["GDAL_CACHEMAX"]="1024"

pan, rgb = map(gdal.Open, argv[1:3])

h, w = pan.RasterYSize, pan.RasterXSize
proj = pan.GetProjection()
gt = pan.GetGeoTransform()

from_proj = rgb.GetProjection()
#from_gt = rgb.GetGeoTransform()

output = gdal.GetDriverByName('GTiff').Create(
    argv[3],
    w, h,
    3,
    gdal.GDT_UInt16,
    options = ['PHOTOMETRIC=RGB'])

output.SetGeoTransform(gt)
output.SetProjection(proj)
print 'here comes reproj'
gdal.ReprojectImage(
    rgb, output,
    from_proj, proj,
    gdalconst.GRA_Bilinear) # gdalconst.GRA_Lanczos, GRA_Bilinear

print 'that was reproj'
r, g, b = [ 
    output.GetRasterBand(n).ReadAsArray(
        0, 0, w, h
    ).astype(np.float32)
    for n in (1, 2, 3) ]

np.seterr(invalid='ignore', divide='ignore')
print 'here comes the ratio'
ratio = pan.GetRasterBand(1).ReadAsArray(
    0, 0, w, h
    ).astype(np.float32) / ((r + b + g)/3
    ).astype(np.float32)

print 'here comes the multiplication'

# todo: loop this up:
print 'red'
r *= ratio
del r

print 'green'
g *= ratio
del g

print 'blue'
b *= ratio
del b

print 'here comes the writing'

output.GetRasterBand(1).WriteArray(r)
output.GetRasterBand(2).WriteArray(g)
output.GetRasterBand(3).WriteArray(b)