Files
landsat-util/landsat/image.py
2016-02-18 14:49:25 -05:00

534 lines
18 KiB
Python

# Pansharpened Image Process using Rasterio
# Landsat Util
# License: CC0 1.0 Universal
import os
import tarfile
import glob
from copy import copy
import subprocess
from shutil import copyfile
from os.path import join, isdir
import numpy
import rasterio
from rasterio.coords import disjoint_bounds
from rasterio.warp import reproject, RESAMPLING, transform, transform_bounds
from skimage import transform as sktransform
from skimage.util import img_as_ubyte
from skimage.exposure import rescale_intensity
from polyline.codec import PolylineCodec
from mixins import VerbosityMixin
from utils import get_file, check_create_folder, exit, adjust_bounding_box
from decorators import rasterio_decorator
class FileDoesNotExist(Exception):
""" Exception to be used when the file does not exist. """
pass
class BoundsDoNotOverlap(Exception):
""" Exception for when bounds do not overlap with the image """
pass
class BaseProcess(VerbosityMixin):
"""
Image procssing class
To initiate the following parameters must be passed:
:param path:
Path of the image.
:type path:
String
:param bands:
The band sequence for the final image. Must be a python list. (optional)
:type bands:
List
:param dst_path:
Path to the folder where the image should be stored. (optional)
:type dst_path:
String
:param verbose:
Whether the output should be verbose. Default is False.
:type verbose:
boolean
:param force_unzip:
Whether to force unzip the tar file. Default is False
:type force_unzip:
boolean
"""
def __init__(self, path, bands=None, dst_path=None, verbose=False, force_unzip=False, bounds=None):
self.projection = {'init': 'epsg:3857'}
self.dst_crs = {'init': u'epsg:3857'}
self.scene = get_file(path).split('.')[0]
self.bands = bands if isinstance(bands, list) else [4, 3, 2]
self.clipped = False
# Landsat source path
self.src_path = path.replace(get_file(path), '')
# Build destination folder if doesn't exist
self.dst_path = dst_path if dst_path else os.getcwd()
self.dst_path = check_create_folder(join(self.dst_path, self.scene))
self.verbose = verbose
# Path to the unzipped folder
self.scene_path = join(self.src_path, self.scene)
# Unzip files
if self._check_if_zipped(path):
self._unzip(join(self.src_path, get_file(path)), join(self.src_path, self.scene), self.scene, force_unzip)
if (bounds):
self.bounds = bounds
self.scene_path = self.clip()
self.clipped = True
self.bands_path = []
for band in self.bands:
self.bands_path.append(join(self.scene_path, self._get_full_filename(band)))
def _get_boundaries(self, src, shape):
self.output("Getting boundaries", normal=True, arrow=True)
output = {'ul': {'x': [0, 0], 'y': [0, 0]}, # ul: upper left
'ur': {'x': [0, 0], 'y': [0, 0]}, # ur: upper right
'll': {'x': [0, 0], 'y': [0, 0]}, # ll: lower left
'lr': {'x': [0, 0], 'y': [0, 0]}} # lr: lower right
output['ul']['x'][0] = src['affine'][2]
output['ul']['y'][0] = src['affine'][5]
output['ur']['x'][0] = output['ul']['x'][0] + self.pixel * src['shape'][1]
output['ur']['y'][0] = output['ul']['y'][0]
output['ll']['x'][0] = output['ul']['x'][0]
output['ll']['y'][0] = output['ul']['y'][0] - self.pixel * src['shape'][0]
output['lr']['x'][0] = output['ul']['x'][0] + self.pixel * src['shape'][1]
output['lr']['y'][0] = output['ul']['y'][0] - self.pixel * src['shape'][0]
output['ul']['x'][1], output['ul']['y'][1] = transform(src['crs'], self.projection,
[output['ul']['x'][0]],
[output['ul']['y'][0]])
output['ur']['x'][1], output['ur']['y'][1] = transform(src['crs'], self.projection,
[output['ur']['x'][0]],
[output['ur']['y'][0]])
output['ll']['x'][1], output['ll']['y'][1] = transform(src['crs'], self.projection,
[output['ll']['x'][0]],
[output['ll']['y'][0]])
output['lr']['x'][1], output['lr']['y'][1] = transform(src['crs'], self.projection,
[output['lr']['x'][0]],
[output['lr']['y'][0]])
dst_corner_ys = [output[k]['y'][1][0] for k in output.keys()]
dst_corner_xs = [output[k]['x'][1][0] for k in output.keys()]
y_pixel = abs(max(dst_corner_ys) - min(dst_corner_ys)) / shape[0]
x_pixel = abs(max(dst_corner_xs) - min(dst_corner_xs)) / shape[1]
return (min(dst_corner_xs), x_pixel, 0.0, max(dst_corner_ys), 0.0, -y_pixel)
def _read_bands(self):
""" Reads a band with rasterio """
bands = []
try:
for i, band in enumerate(self.bands):
bands.append(rasterio.open(self.bands_path[i]).read_band(1))
except IOError as e:
exit(e.message, 1)
return bands
def _warp(self, proj_data, bands, new_bands):
self.output("Projecting", normal=True, arrow=True)
for i, band in enumerate(bands):
self.output("band %s" % self.bands[i], normal=True, color='green', indent=1)
reproject(band, new_bands[i], src_transform=proj_data['transform'], src_crs=proj_data['crs'],
dst_transform=proj_data['dst_transform'], dst_crs=self.dst_crs, resampling=RESAMPLING.nearest,
num_threads=2)
def _unzip(self, src, dst, scene, force_unzip=False):
""" Unzip tar files """
self.output("Unzipping %s - It might take some time" % scene, normal=True, arrow=True)
try:
# check if file is already unzipped, skip
if isdir(dst) and not force_unzip:
self.output('%s is already unzipped.' % scene, normal=True, color='green', indent=1)
return
else:
tar = tarfile.open(src, 'r')
tar.extractall(path=dst)
tar.close()
except tarfile.ReadError:
check_create_folder(dst)
subprocess.check_call(['tar', '-xf', src, '-C', dst])
def _get_full_filename(self, band):
base_file = '%s_B%s.*' % (self.scene, band)
try:
return glob.glob(join(self.scene_path, base_file))[0].split('/')[-1]
except IndexError:
raise FileDoesNotExist('%s does not exist' % '%s_B%s.*' % (self.scene, band))
def _check_if_zipped(self, path):
""" Checks if the filename shows a tar/zip file """
filename = get_file(path).split('.')
if filename[-1] in ['bz', 'bz2', 'gz']:
return True
return False
def _read_metadata(self):
output = {}
try:
with open(self.scene_path + '/' + self.scene + '_MTL.txt', 'rU') as mtl:
lines = mtl.readlines()
for line in lines:
if 'REFLECTANCE_ADD_BAND_3' in line:
output['REFLECTANCE_ADD_BAND_3'] = float(line.replace('REFLECTANCE_ADD_BAND_3 = ', ''))
if 'REFLECTANCE_MULT_BAND_3' in line:
output['REFLECTANCE_MULT_BAND_3'] = float(line.replace('REFLECTANCE_MULT_BAND_3 = ', ''))
if 'REFLECTANCE_ADD_BAND_4' in line:
output['REFLECTANCE_ADD_BAND_4'] = float(line.replace('REFLECTANCE_ADD_BAND_4 = ', ''))
if 'REFLECTANCE_MULT_BAND_4' in line:
output['REFLECTANCE_MULT_BAND_4'] = float(line.replace('REFLECTANCE_MULT_BAND_4 = ', ''))
if 'CLOUD_COVER' in line:
output['CLOUD_COVER'] = float(line.replace('CLOUD_COVER = ', ''))
return output
except IOError:
return output
def _get_image_data(self):
src = rasterio.open(self.bands_path[-1])
# Get pixel size from source
self.pixel = src.affine[0]
# Only collect src data that is needed and delete the rest
image_data = {
'transform': src.transform,
'crs': src.crs,
'affine': src.affine,
'shape': src.shape,
'dst_transform': None
}
image_data['dst_transform'] = self._get_boundaries(image_data, image_data['shape'])
return image_data
def _generate_new_bands(self, shape):
new_bands = []
for i in range(0, 3):
new_bands.append(numpy.empty(shape, dtype=numpy.uint16))
return new_bands
@rasterio_decorator
def _write_to_file(self, new_bands, **kwargs):
# Read coverage from QBA
coverage = self._calculate_cloud_ice_perc()
self.output("Final Steps", normal=True, arrow=True)
suffix = 'bands_%s' % "".join(map(str, self.bands))
output_file = join(self.dst_path, self._filename(suffix=suffix))
output = rasterio.open(output_file, 'w', **kwargs)
for i, band in enumerate(new_bands):
# Color Correction
band = self._color_correction(band, self.bands[i], 0, coverage)
output.write_band(i + 1, img_as_ubyte(band))
new_bands[i] = None
self.output("Writing to file", normal=True, color='green', indent=1)
return output_file
def _color_correction(self, band, band_id, low, coverage):
if self.bands == [4, 5]:
return band
else:
self.output("Color correcting band %s" % band_id, normal=True, color='green', indent=1)
p_low, cloud_cut_low = self._percent_cut(band, low, 100 - (coverage * 3 / 4))
temp = numpy.zeros(numpy.shape(band), dtype=numpy.uint16)
cloud_divide = 65000 - coverage * 100
mask = numpy.logical_and(band < cloud_cut_low, band > 0)
temp[mask] = rescale_intensity(band[mask], in_range=(p_low, cloud_cut_low), out_range=(256, cloud_divide))
temp[band >= cloud_cut_low] = rescale_intensity(band[band >= cloud_cut_low],
out_range=(cloud_divide, 65535))
return temp
def _percent_cut(self, color, low, high):
return numpy.percentile(color[numpy.logical_and(color > 0, color < 65535)], (low, high))
def _calculate_cloud_ice_perc(self):
self.output('Calculating cloud and snow coverage from QA band', normal=True, arrow=True)
a = rasterio.open(join(self.scene_path, self._get_full_filename('QA'))).read_band(1)
count = 0
snow = [56320, 39936, 31744, 28590, 26656, 23552]
cloud = [61440, 59424, 57344, 53248, 28672, 36896, 36864, 24576]
for item in cloud:
count += numpy.extract(a == item, a).size
for item in snow:
count += numpy.extract(a == item, a).size * 2
perc = numpy.true_divide(count, a.size) * 100
self.output('cloud/snow coverage: %s' % round(perc, 2), indent=1, normal=True, color='green')
return perc
def _filename(self, name=None, suffix=None, prefix=None):
""" File name generator for processed images """
filename = ''
if prefix:
filename += str(prefix) + '_'
if name:
filename += str(name)
else:
filename += str(self.scene)
if suffix:
filename += '_' + str(suffix)
if self.clipped:
bounds = [tuple(self.bounds[0:2]), tuple(self.bounds[2:4])]
polyline = PolylineCodec().encode(bounds)
filename += '_clipped_' + polyline
filename += '.TIF'
return filename
@rasterio_decorator
def clip(self):
""" Clip images based on bounds provided
Implementation is borrowed from
https://github.com/brendan-ward/rasterio/blob/e3687ce0ccf8ad92844c16d913a6482d5142cf48/rasterio/rio/convert.py
"""
self.output("Clipping", normal=True)
# create new folder for clipped images
path = check_create_folder(join(self.scene_path, 'clipped'))
try:
temp_bands = copy(self.bands)
temp_bands.append('QA')
for i, band in enumerate(temp_bands):
band_name = self._get_full_filename(band)
band_path = join(self.scene_path, band_name)
self.output("Band %s" % band, normal=True, color='green', indent=1)
with rasterio.open(band_path) as src:
bounds = transform_bounds(
{
'proj': 'longlat',
'ellps': 'WGS84',
'datum': 'WGS84',
'no_defs': True
},
src.crs,
*self.bounds
)
if disjoint_bounds(bounds, src.bounds):
bounds = adjust_bounding_box(src.bounds, bounds)
window = src.window(*bounds)
out_kwargs = src.meta.copy()
out_kwargs.update({
'driver': 'GTiff',
'height': window[0][1] - window[0][0],
'width': window[1][1] - window[1][0],
'transform': src.window_transform(window)
})
with rasterio.open(join(path, band_name), 'w', **out_kwargs) as out:
out.write(src.read(window=window))
# Copy MTL to the clipped folder
copyfile(join(self.scene_path, self.scene + '_MTL.txt'), join(path, self.scene + '_MTL.txt'))
return path
except IOError as e:
exit(e.message, 1)
class Simple(BaseProcess):
@rasterio_decorator
def run(self):
""" Executes the image processing.
:returns:
(String) the path to the processed image
"""
self.output('Image processing started for bands %s' % '-'.join(map(str, self.bands)), normal=True, arrow=True)
bands = self._read_bands()
image_data = self._get_image_data()
new_bands = self._generate_new_bands(image_data['shape'])
self._warp(image_data, bands, new_bands)
# Bands are no longer needed
del bands
rasterio_options = {
'driver': 'GTiff',
'width': image_data['shape'][1],
'height': image_data['shape'][0],
'count': 3,
'dtype': numpy.uint8,
'nodata': 0,
'transform': image_data['dst_transform'],
'photometric': 'RGB',
'crs': self.dst_crs
}
return self._write_to_file(new_bands, **rasterio_options)
class PanSharpen(BaseProcess):
def __init__(self, path, bands=None, **kwargs):
if bands:
bands.append(8)
else:
bands = [4, 3, 2, 8]
self.band8 = bands.index(8)
super(PanSharpen, self).__init__(path, bands, **kwargs)
@rasterio_decorator
def run(self):
""" Executes the pansharpen image processing.
:returns:
(String) the path to the processed image
"""
self.output('PanSharpened Image processing started for bands %s' % '-'.join(map(str, self.bands)),
normal=True, arrow=True)
bands = self._read_bands()
image_data = self._get_image_data()
new_bands = self._generate_new_bands(image_data['shape'])
bands[:3] = self._rescale(bands[:3])
new_bands.append(numpy.empty(image_data['shape'], dtype=numpy.uint16))
self._warp(image_data, bands, new_bands)
# Bands are no longer needed
del bands
# Calculate pan band
pan = self._pansize(new_bands)
del self.bands[self.band8]
del new_bands[self.band8]
rasterio_options = {
'driver': 'GTiff',
'width': image_data['shape'][1],
'height': image_data['shape'][0],
'count': 3,
'dtype': numpy.uint8,
'nodata': 0,
'transform': image_data['dst_transform'],
'photometric': 'RGB',
'crs': self.dst_crs
}
return self._write_to_file(new_bands, pan, **rasterio_options)
@rasterio_decorator
def _write_to_file(self, new_bands, pan, **kwargs):
# Read coverage from QBA
coverage = self._calculate_cloud_ice_perc()
self.output("Final Steps", normal=True, arrow=True)
suffix = 'bands_%s_pan' % "".join(map(str, self.bands))
output_file = join(self.dst_path, self._filename(suffix=suffix))
output = rasterio.open(output_file, 'w', **kwargs)
for i, band in enumerate(new_bands):
# Color Correction
band = numpy.multiply(band, pan)
band = self._color_correction(band, self.bands[i], 0, coverage)
output.write_band(i + 1, img_as_ubyte(band))
new_bands[i] = None
self.output("Writing to file", normal=True, color='green', indent=1)
return output_file
def _pansize(self, bands):
self.output('Calculating Pan Ratio', normal=True, arrow=True)
m = numpy.add(bands[0], bands[1])
m = numpy.add(m, bands[2])
pan = numpy.multiply(numpy.nan_to_num(numpy.true_divide(1, m)), bands[self.band8])
return pan
def _rescale(self, bands):
""" Rescale bands """
self.output("Rescaling", normal=True, arrow=True)
for key, band in enumerate(bands):
self.output("band %s" % self.bands[key], normal=True, color='green', indent=1)
bands[key] = sktransform.rescale(band, 2)
bands[key] = (bands[key] * 65535).astype('uint16')
return bands
if __name__ == '__main__':
p = PanSharpen('/Users/ajdevseed/Desktop/LC81950282014159LGN00')
p.run()