cloudMask.py

import getopt
import glob
import os
import platform
import subprocess
import sys

import gdal
import numpy as np

from mtdUtils import createOverlappingTIFF


def createMaskFromFmaskOutput(fin, fout):
    """
    Converts FMask output in a binary validity mask, excluding clouds and shadows.

    """
    ds = gdal.Open(fin)
    bnd = ds.GetRasterBand(1)
    arr = bnd.ReadAsArray()

    # Validity mask: keeps all that is different from Cloud or Shadow
    out = np.logical_and(arr != 2, arr != 3)
    out = np.logical_and(out, arr != 0)

    dsout = createOverlappingTIFF(fout, ds, 1, gdal.GDT_Byte)
    bndout = dsout.GetRasterBand(1)
    ret = bndout.WriteArray(out.astype(np.uint8))
    bndout.SetNoDataValue(0)

    bnd = None
    bndout = None
    ds = None
    dsout = None

    return ret


def L8_CloudMask(fld, hr):
    print 'Computing cloud mask...'

    if platform.system() == 'Linux':
        gdl_merge = ['gdal_merge.py', '-separate', '-of', 'GTiff', '-o']
        fmsk_angles = ['fmask_usgsLandsatMakeAnglesImage.py']
        fmsk_satur = ['fmask_usgsLandsatSaturationMask.py']
        fmsk_toa = ['fmask_usgsLandsatTOA.py']
        fmsk_clouds = ['fmask_usgsLandsatStacked.py']
        sh = False
    elif platform.system() == 'Windows':
        gdl_merge = ['gdal_merge', '-separate', '-of', 'GTiff', '-o']
        fmsk_angles = ['python', 'fmask_usgsLandsatMakeAnglesImage.py']
        fmsk_satur = ['python', 'fmask_usgsLandsatSaturationMask.py']
        fmsk_toa = ['python', 'fmask_usgsLandsatTOA.py']
        fmsk_clouds = ['python', 'fmask_usgsLandsatStacked.py']
        sh = True
    else:
        sys.exit("Platform not supported!")

    b1fn = ''
    b1fns = glob.glob(fld + '/*_B1*.TIF')
    for fln in b1fns:
        if 'B10' not in fln and 'B11' not in fln:
            b1fn = fln

    oli_list = [b1fn,
                glob.glob(fld + '/*_B2*.TIF')[0],
                glob.glob(fld + '/*_B3*.TIF')[0],
                glob.glob(fld + '/*_B4*.TIF')[0],
                glob.glob(fld + '/*_B5*.TIF')[0],
                glob.glob(fld + '/*_B6*.TIF')[0],
                glob.glob(fld + '/*_B7*.TIF')[0],
                glob.glob(fld + '/*_B9*.TIF')[0]]

    tirs_list = [glob.glob(fld + '/*_B10*.TIF')[0],
                 glob.glob(fld + '/*_B11*.TIF')[0]]

    mdf = ['-m', glob.glob(fld + '/*_MTL*.txt')[0]]

    oli_fn = [fld + '/oli.tif']
    tirs_fn = [fld + '/tirs.tif']
    angles_fn = [fld + '/angles.tif']
    satur_fn = [fld + '/saturation.tif']
    oli_toa_fn = [glob.glob(fld + '/*_B2*.TIF')[0].replace('_B2', '_OLI_TOA')]
    cmask_fn = [glob.glob(fld + '/*_B2*.TIF')[0].replace('_B2', '_CLOUDS')]
    mask_fn = [glob.glob(fld + '/*_B2*.TIF')[0].replace('_B2', '_MASK')]

    subprocess.call(gdl_merge + oli_fn + oli_list, shell=sh)
    subprocess.call(gdl_merge + tirs_fn + tirs_list, shell=sh)
    subprocess.call(fmsk_angles + mdf + ['-t'] + oli_fn + ['-o'] + angles_fn, shell=sh)
    subprocess.call(fmsk_toa + ['-i'] + oli_fn + mdf + ['-z'] + angles_fn + ['-o'] + oli_toa_fn, shell=sh)

    # **** Manage high resolution mask : BEGIN****
    # Solution 1 : if HR is set, perform on pansharpened
    '''
    if hr:
        pan_fn = [glob.glob(fld + '/*_B8*.TIF')[0]]
        pan_toa_fn = [pan_fn[0].replace('_B8', '_PAN_TOA')]
        ps_toa_fn = [glob.glob(fld + '/*_B2*.TIF')[0].replace('_B2', '_PS_TOA')]
        angles_pan_fn = [fld + '/angles_pan.tif']
        subprocess.call(fmsk_angles + mdf + ['-t'] + pan_fn + ['-o'] + angles_pan_fn, shell=sh)
        subprocess.call(fmsk_toa + ['-i'] + pan_fn + mdf + ['-z'] + angles_pan_fn + ['-o'] + pan_toa_fn, shell=sh)
        subprocess.call(
            ['otbcli_Superimpose', '-inr'] + pan_toa_fn + ['-inm'] + oli_toa_fn + ['-out', fld + '/ms.tif', 'uint16'],
            shell=sh)
        subprocess.call(
            ['otbcli_Superimpose', '-inr'] + pan_toa_fn + ['-inm'] + tirs_fn + ['-out', fld + '/tirs15.tif', 'uint16'],
            shell=sh)
        subprocess.call(
            ['otbcli_Pansharpening', '-inp'] + pan_toa_fn + ['-inxs', fld + '/ms.tif', '-out'] + ps_toa_fn + ['uint16',
                                                                                                              '-method',
                                                                                                              'bayes'],
            shell=sh)
        subprocess.call(fmsk_satur + ['-i'] + ps_toa_fn + mdf + ['-o'] + satur_fn, shell=sh)
        subprocess.call(
            fmsk_clouds + ['-t'] + [fld + '/tirs15.tif'] + ['-a'] + ps_toa_fn + mdf + ['-z'] + angles_pan_fn + [
                '-s'] + satur_fn + [
                '-o'] + cmask_fn, shell=sh)
        os.remove(satur_fn[0])
        os.remove(fld + '/ms.tif')
        os.remove(fld + '/tirs15.tif')
    else:
        subprocess.call(fmsk_satur + ['-i'] + oli_fn + mdf + ['-o'] + satur_fn, shell=sh)
        subprocess.call(
            fmsk_clouds + ['-t'] + tirs_fn + ['-a'] + oli_toa_fn + mdf + ['-z'] + angles_fn + ['-s'] + satur_fn + [
                '-o'] + cmask_fn, shell=sh)
        os.remove(satur_fn[0])
    '''
    # --- Very selective behaviour, switching to...
    # Solution 2: a simple resample (superimpose) of the 30m validity mask
    subprocess.call(fmsk_satur + ['-i'] + oli_fn + mdf + ['-o'] + satur_fn, shell=sh)
    subprocess.call(
        fmsk_clouds + ['-t'] + tirs_fn + ['-a'] + oli_toa_fn + mdf + ['-z'] + angles_fn + ['-s'] + satur_fn + [
            '-o'] + cmask_fn, shell=sh)
    createMaskFromFmaskOutput(cmask_fn[0], mask_fn[0])
    os.remove(satur_fn[0])
    if hr:
        pan_fn = [glob.glob(fld + '/*_B8*.TIF')[0]]
        mask_hr_fn = [glob.glob(fld + '/*_B2*.TIF')[0].replace('_B2', '_MASK_HR')]
        subprocess.call(
            ['otbcli_Superimpose', '-inr'] + pan_fn + ['-inm'] + mask_fn + ['-out'] + mask_hr_fn + [
                '-interpolator',
                'nn'],
            shell=sh)
        os.remove(mask_fn[0])
    # **** Manage high resolution mask : END ****

    os.remove(oli_fn[0])
    os.remove(tirs_fn[0])
    os.remove(angles_fn[0])
    for f in glob.glob(fld + '/*.aux.xml'):
        os.remove(f)


def S2_CloudMask(fld, hr):
    print 'Computing cloud mask...'

    if hr:
        res = '10'
    else:
        res = '20'

    gdl_vrt = ['gdalbuildvrt', '-resolution', 'user', '-tr', res, res, '-separate']

    if platform.system() == 'Linux':
        fmsk_angles = ['fmask_sentinel2makeAnglesImage.py']
        fmsk_clouds = ['fmask_sentinel2Stacked.py']
        sh = False
    elif platform.system() == 'Windows':
        fmsk_angles = ['python', 'fmask_sentinel2makeAnglesImage.py']
        fmsk_clouds = ['python', 'fmask_sentinel2Stacked.py']
        sh = True
    else:
        sys.exit("Platform not supported!")

    bnd_list = [glob.glob(fld + '/*B01*.tif')[0],
                glob.glob(fld + '/*B02*.tif')[0],
                glob.glob(fld + '/*B03*.tif')[0],
                glob.glob(fld + '/*B04*.tif')[0],
                glob.glob(fld + '/*B05*.tif')[0],
                glob.glob(fld + '/*B06*.tif')[0],
                glob.glob(fld + '/*B07*.tif')[0],
                glob.glob(fld + '/*B08*.tif')[0],
                glob.glob(fld + '/*B8A*.tif')[0],
                glob.glob(fld + '/*B09*.tif')[0],
                glob.glob(fld + '/*B10*.tif')[0],
                glob.glob(fld + '/*B11*.tif')[0],
                glob.glob(fld + '/*B12*.tif')[0]]

    mdf = ['-i', fld + '/metadata.xml']

    vrt_fn = [fld + '/allbands.vrt']
    angles_fn = [fld + '/angles.tif']
    cmask_fn = [glob.glob(fld + '/*B02*.tif')[0].replace('B02', 'CLOUDS')]
    mask_tmp_fn = [glob.glob(fld + '/*B02*.tif')[0].replace('B02', 'MASK_tmp')]
    mask_fn = [glob.glob(fld + '/*B02*.tif')[0].replace('B02', 'MASK')]

    subprocess.call(gdl_vrt + vrt_fn + bnd_list, shell=sh)
    subprocess.call(fmsk_angles + mdf + ['-o'] + angles_fn, shell=sh)
    subprocess.call(fmsk_clouds + ['-a'] + vrt_fn + ['-z'] + angles_fn + ['-o'] + cmask_fn, shell=sh)
    createMaskFromFmaskOutput(cmask_fn[0], mask_tmp_fn[0])

    if hr:
        refbnd = bnd_list[1]
    else:
        refbnd = bnd_list[4]

    subprocess.call(
        ['otbcli_Superimpose', '-inr', refbnd, '-inm'] + mask_tmp_fn + ['-out'] + mask_fn + ['uint8', '-interpolator',
                                                                                             'nn'],
        shell=sh)

    os.remove(vrt_fn[0])
    os.remove(angles_fn[0])
    os.remove(mask_tmp_fn[0])
    for f in glob.glob(fld + '/*.aux.xml'):
        os.remove(f)


def main(argv):
    try:
        opts, args = getopt.getopt(argv, 'h', ['HR'])
    except getopt.GetoptError as err:
        print str(err)

    hr = False

    fld = args[0]
    for opt, val in opts:
        if opt == '--HR':
            hr = True

    # check if S2 or L8 (search for metadata file)
    mdfL8 = glob.glob(fld + '/*_MTL*.txt')
    mdfS2 = glob.glob(fld + '/tileInfo.json')
    if len(mdfL8) > 0:
        L8_CloudMask(fld, hr)
    elif len(mdfS2) > 0:
        S2_CloudMask(fld, hr)
    else:
        sys.exit("No supported Landsat 8 or Sentinel-2 data in " + fld)
    return 0


if __name__ == '__main__':
    if len(sys.argv) < 2:
        sys.exit(
            'Usage: python cloudMask.py [--HR] <image folder>')
    main(sys.argv[1:])