panchunk.py

# Beta-quality Landsat 8-oriented windowed pansharpener by Charlie Loyd

# python3 panchunk.py $SCENE/*B{4,3,2,8}.TIF ${SCENE}-pansharp.tif

# Then you may want something like this to brighten it up:
# convert -channel B -gamma 0.96 -channel RGB -sigmoidal-contrast 40,14% ${SCENE}-pansharp.tif ${SCENE}-pretty.tif
# But that will strip geo tags.

import asyncio
import rasterio as rio
import numpy as np
from rasterio.warp import reproject, RESAMPLING

# todo: figure out why proper windows weren't working
blocksize=2048

matht = np.float32
np.seterr(invalid='ignore', divide='ignore')

def half_window(window):
	((fr, tr), (fc, tc)) = window
	return ((int(fr/2), int(tr/2)), (int(fc/2), int(tc/2)))

def make_windows(w, h):
	for x in range(int(w/blocksize+1)):
	 	for y in range(int(h/blocksize+1)):
	 		yield ( # todo: tighten up the slop here
	 			(y*blocksize, min(h-1, (y*blocksize)+blocksize+2)),
				(x*blocksize, min(w-1, (x*blocksize)+blocksize+2))
			)

def main(red, green, blue, pan, out):
	r, g, b, p = (rio.open(f) for f in [red, green, blue, pan])
	rgb_o = (r, g, b)
	
	color_meta, pan_meta = b.meta, p.meta
	w, h = pan_meta['width'], pan_meta['height']
	
	pan_meta.update(photometric='rgb', count=3)
	
	with rio.open(out, 'w', **pan_meta) as dest:
		for wind in make_windows(w, h):
			#print(wind, half_window(wind))
			pan = p.read_band(1, window=wind).astype(matht)
			rgb = list(x.read_band(1, window=half_window(wind)).astype(matht) for x in rgb_o)
			
			for i, c in enumerate(rgb):
				x = np.empty(pan.shape)
				reproject(
					c, x,
					src_transform=color_meta['transform'],
					src_crs=color_meta['crs'],
					dst_transform=pan_meta['transform'],
					dst_crs=pan_meta['crs'],
					resampling=RESAMPLING.bilinear) # hmmm
					
				rgb[i] = x
				
			ratio = (pan / ((rgb[0] + rgb[1] + rgb[2])/3))
			
			for i, data in enumerate(rgb):
				done = (data * ratio).astype(np.uint16)
				dest.write_band(i+1, done, window=wind)
				
	print('that was cool and fun')

if __name__ == '__main__':
	import sys
	main(*sys.argv[1:])