prepareS2THEIA.py

import getopt
import os
import subprocess
import sys
import glob
import gdal
import osr
import shutil
import platform
import stat
import numpy as np

from mtdUtils import randomword, getRasterExtentAsShapefile

def genScript(argv):
    try:
        opts, args = getopt.getopt(argv, 'o:c:t:',['srs='])
    except getopt.GetoptError as err:
        print str(err)

    # Get sub-folders names
    dirs = [args[0] + '/' + dn for dn in os.listdir(args[0]) if os.path.isdir(args[0] + '/' + dn)]

    # Parse options
    od = os.getcwd()
    clipping = False
    clip_shp = None
    single_tile = ''
    srs=None
    for opt, val in opts:
        if opt == '-o':
            od = val
        elif opt == '-c':
            clip_shp = val
            clipping = True
        elif opt == '-t':
            single_tile = val
        elif opt == '--srs':
            srs = str(val)
    if srs is None :
        ds = gdal.Open(glob.glob(dirs[0] + '/*_B2.tif')[0])
        prj = ds.GetProjectionRef()
        srs = osr.SpatialReference(prj)
        srs = 'EPSG:'+srs.GetAuthorityCode('PROJCS')

    if platform.system() == 'Windows':
        f = open('S2THEIAPreparation.bat', 'w')
        f.write("@echo off\n")
    else:
        f = open('S2THEIAPreparation.sh', 'w')

    for dir in dirs:
        if single_tile != '' and single_tile not in dir:
            continue
        else:
            mtd = glob.glob(dir + '/*_MTD_ALL.xml')
            if len(mtd) > 0:
                #cmd = ['python', 'prepareS2THEIA.py'] + argv[:-1] + [dir]
                cmd = ['python', 'prepareS2THEIA.py', '-o', '\"' + od + '\"']
                if clip_shp is not None:
                    cmd += ['-c', '\"' + clip_shp + '\"']
                cmd += ['--srs', srs]
                cmd += ['\"' + dir + '\"']
                f.write(' '.join(cmd) + '\n')
    f.close()

    if platform.system() == 'Linux':
        st = os.stat('S2THEIAPreparation.sh')
        os.chmod('S2THEIAPreparation.sh', st.st_mode | stat.S_IEXEC)

    return

def prepare(argv):
    try:
        opts, args = getopt.getopt(argv, 'o:c:', ['srs='])
    except getopt.GetoptError as err:
        print str(err)

    # Parse options
    od = os.getcwd()
    clipping = False
    clip_shp = None
    srs = None

    for opt, val in opts:
        if opt == '-o':
            od = val
        elif opt == '-c':
            clip_shp = val
            clipping = True
        elif opt == '--srs':
            srs = str(val)

    src_dir = args[0]
    tile_date = os.path.basename(src_dir)[11:19]
    tile_id = os.path.basename(src_dir).split('_')[3][1:]

    if os.environ.get('GDAL_DATA') == None:
        sys.exit('Please set GDAL_DATA environment variable!')

    print('Processing ' + src_dir + '...')

    # Create output directories
    if not os.path.exists(od):
        os.mkdir(od)
    seqnum = 0
    pd = od
    od = pd + '/S2-L2A-THEIA_' + str(seqnum) + '_' + tile_date + '_' + tile_id
    while os.path.exists(od):
        seqnum += 1
        od = pd + '/S2-L2A-THEIA_' + str(seqnum) + '_' + tile_date + '_' + tile_id
    os.mkdir(od)

    curdir = os.getcwd()
    os.chdir(od)

    # Prepare ROI (full extent if not clipping)
    tmpname = randomword(16)
    if not clipping:
        clip_shp = od + '/' + tmpname + '.shp'
        ext_ref = glob.glob(src_dir + '/*_FRE_B2.tif')[0]
        getRasterExtentAsShapefile(ext_ref, clip_shp)
        ext = 'full'
    else:
        ext = os.path.splitext(os.path.basename(clip_shp))[0]

    clip_shp_val = clip_shp
    # Generate and run clip commands
    clip_cmd = 'gdalwarp -q -of GTiff -ot Int16 -dstnodata -10000 -srcnodata -10000 -cutline \"' + clip_shp_val + '\" -crop_to_cutline '
    if srs != None :
        clip_cmd += '-t_srs {} '.format(srs)
    fns = ['B2', 'B3', 'B4', 'B5', 'B6', 'B7', 'B8', 'B8A', 'B11', 'B12']

    for fn_pre in fns:
        fn = glob.glob(src_dir + '/*_FRE_' + fn_pre + '.tif')[0]
        out_fn = od + '/' + os.path.splitext(os.path.basename(fn))[0] + '_' + ext + '.tif'
        subprocess.call(clip_cmd + '\"' + fn + '\" \"' + out_fn + '\"', shell=True)

    # clip_cmd = 'gdalwarp -q -of GTiff -ot UInt16 -dstnodata 0 -srcnodata 0 -cutline \"' + clip_shp_val + '\" -crop_to_cutline '
    clip_cmd = 'gdalwarp -q -of GTiff -ot UInt16 -cutline \"' + clip_shp_val + '\" -crop_to_cutline '
    if srs != None :
        clip_cmd += '-t_srs {} '.format(srs)
    fns = ['CLM_R1', 'CLM_R2']
    for fn_pre in fns:
        fn = glob.glob(src_dir + '/MASKS/*_' + fn_pre + '.tif')[0]
        #out_tmp_fn = od + '/' + os.path.splitext(os.path.basename(fn))[0] + '_tmp.tif'
        #cmd = ['otbcli_BandMath', '-il', fn, '-exp', 'im1b1+1', '-out', out_tmp_fn, 'uint16']
        #subprocess.call(' '.join(cmd), shell=True)
        out_fn = od + '/' + os.path.splitext(os.path.basename(fn))[0] + '_' + ext + '.tif'
        #subprocess.call(clip_cmd + '\"' + out_tmp_fn + '\" \"' + out_fn + '\"', shell=True)
        subprocess.call(clip_cmd + '\"' + fn + '\" \"' + out_fn + '\"', shell=True)
        #os.remove(out_tmp_fn)

    mtd_fn = glob.glob(src_dir + '/*_MTD_ALL.xml')[0]
    shutil.copyfile(mtd_fn, od + '/' + os.path.basename(mtd_fn))

    #check if the tile are in the wanted area or not
    ds = gdal.Open(glob.glob(os.path.join(od,'*B4*.tif'))[0])
    ly = ds.GetRasterBand(1)
    nbval = len(np.unique(ly.ReadAsArray()))
    ds = None
    if nbval == 1 :
        print(" {}_{} doesn't cover requested area, deleting folder".format(tile_date,tile_id))
        shutil.rmtree(od)

    os.chdir(curdir)


if __name__ == '__main__':
    if len(sys.argv) < 3:
        sys.exit(
            'Usage: python prepareS2THEIA.py [-o <output dir>] [-c <clip_shp>] [--srs <EPSG:4326] [-t tile] <original THEIA tile folder / parent folder>\n'
            'If parent folder is given, searches for all available images and creates a preparation batch file.')
    else:
        mtd = glob.glob(sys.argv[-1] + '/*_MTD_ALL.xml')
        if len(mtd) > 0:
            prepare(sys.argv[1:])
        else:
            genScript(sys.argv[1:])