Prepare Landsat image code

prepare_landsat_image.py

import recipy

import sys
from LandsatUtils import extract_and_process_uncorrected, BOAReflectance, mask_all_bands
from GDALUtils import image_like
from PyFMask.fmask import run_FMask
import gdal
import logging

import os
from glob import glob


def prepare(filename, path, cloud_prob=22.5):
    """
    Prepares a Landsat image for use with the HOTBAR algorithm.

    Parameters:
        - filename: full path to .tar.gz file downloaded from EarthExplorer
                    (or equivalent)
        - path: path to store extracted and processed files

    Performs the following steps:

    1. Extract .tar.gz
    2. Layerstack bands
    3. Process to radiance
    4. Perform cloud masking using the FMask method
    5. Process to Bottom of Atmosphere reflectance (BOAp)

    Supports Landsat 5, 7 and 8 imagery.
    """

    # Extract .tar.gz, stack and process to radiance
    prepare1(filename, path)

    # Do FMask cloud masking
    cloud_mask(path, cloud_prob=cloud_prob)

    # Process to BOAp
    prepare2(path)


def cloud_mask(path, cloud_prob=22.5):
    """
    Uses the FMask algorithm to mask clouds from the Landsat imagery in the
    given folder.

    Assumes that the original MTL file and original bands are in the folder,
    along with a Merged_Radiance.tif file created by prepare1().

    Uses the default FMask parameters as suggested in the original paper: 22.5
    percent, with 3 pixel dilations.
    """
    # Run the FMask algorithm to produce a cloud mask from the raw Landsat data
    # (this just requires the path to the MTL file)
    mtl_filename = glob(os.path.join(path, '*_MTL.txt'))[0]

    run_FMask(mtl_filename, cldprob=cloud_prob)

    # Take the mask and apply it to the Merged_Radiance.tif file created by
    # the prepare1 function
    mask_all_bands(os.path.join(path, 'Merged_Radiance.tif'),
                   os.path.join(path, 'fmask_final_mask'),
                   os.path.join(path, 'Merged_Radiance_Masked.tif'),
                   mask_value_to_delete=255, nodata=-9999)


def prepare1(filename, path):
    """
    Takes a Landsat .tar.gz file downloaded from EarthExplorer (or similar),
    extracts the files, stacks them, and processes to radiance.

    The next stage is cloud masking, and the file should be saved in the same
    folder but with the name Merged_Radiance_Masked.tif. If there are no clouds,
    then don't create any new file, and it will use Merged_Radiance.tif

    """
    extract_and_process_uncorrected(filename, path, "Merged_Radiance.tif")


def prepare2(path):
    """
    Takes a cloud-masked version of the output from prepare1 and processes it
    to BOAp ready for use in the PhDAlgorithm procedure.
    """

    filename = os.path.join(path, 'Merged_Radiance_Masked.tif')

    if not os.path.exists(filename):
        filename = os.path.join(path, 'Merged_Radiance.tif')

    logging.debug(filename)
    metadata_filename = glob(os.path.join(path, "*MTL.txt"))[0]

    b = BOAReflectance(metadata_filename)
    b.create_lut()

    input_im = gdal.Open(filename)
    output_im = image_like(input_im,
                           os.path.join(path, 'Merged_BOAp.tif'),
                           datatype=float)

    logging.debug(filename)
    logging.debug(metadata_filename)
    logging.debug(input_im.RasterCount)

    for i in range(1, input_im.RasterCount + 1):
        in_band = input_im.GetRasterBand(i).ReadAsArray()

        out_band = output_im.GetRasterBand(i)

        boap_band = b.correct_band(in_band, i)

        out_band.WriteArray(boap_band)

    del input_im
    del output_im

if __name__ == '__main__':
    if len(sys.argv) == 2:
        filename = sys.argv[1]
        path = filename.replace('.tar.gz', '')
    elif len(sys.argv) == 3:
        filename = sys.argv[1]
        path = sys.argv[2]
    else:
        print('Usage: prepare_landsat_image.py mtl_filename [extraction_path]')
        sys.exit()

    prepare(filename, path)