genMosaicScript.py 10.18 KiB
import getopt
import glob
import os
import platform
import subprocess
import sys
import ogr
import gdal
import warnings
from shutil import copyfile
from datetime import datetime,timedelta
from mtdUtils import checkSRS, setNoDataValue, getNoDataValue

def main(argv):

    if platform.system() == 'Linux':
        sh = False
    elif platform.system() == 'Windows':
        sh = True
    else:
        sys.exit("Platform not supported!")

    try:
        opts, args = getopt.getopt(argv, '', ['onstack','extent=','force'])
    except getopt.GetoptError as err:
        print(str(err))

    imroot = args[0]
    dirs = [imroot + '/' + n for n in os.listdir(imroot) if os.path.isdir(imroot + '/' + n)]

    ext_fn = None
    force = False
    onstack = False

    for opt,val in opts:
        if opt == '--onstack':
            onstack = True
        if opt == '--extent':
            ext_fn = val
        if opt == '--force':
            force = True

    if onstack and (ext_fn is None or not os.path.exists(ext_fn)):
        sys.exit("Missing or invalid extent file.")

    sdt = []
    max_delay = {}
    max_delay['S2'] = 3
    max_delay['L8'] = 9

    sen = None

    for d in dirs:
        qq = os.path.basename(d).split('_')
        if qq[0][0:2] in ["S2","L8"]:
            if sen is None:
                sen = qq[0][0:2]
            sdt.append(tuple(qq[-2:] + [d]))

    sdt.sort(key=lambda l: l[0])

    toMosaic = []
    cur_date = datetime.strptime(sdt[0][0],'%Y%m%d')
    i = 1
    locToMosaic = [sdt[0][2]]
    while i < len(sdt):
        dt = datetime.strptime(sdt[i][0],'%Y%m%d')
        if (dt - cur_date) <= timedelta(days = max_delay[sen]):
            locToMosaic.append(sdt[i][2])
        else:
            toMosaic.append(locToMosaic)
            locToMosaic = [sdt[i][2]]
            cur_date = dt
        i += 1
    toMosaic.append(locToMosaic)

    mosd = imroot + '/MOSAICS/'
    env_fn = mosd + 'MOSAIC_envelope.tif'

    ext_srs_epsg = None
    if ext_fn is not None:
        ds_ext = ogr.Open(ext_fn)
        ly_ext = ds_ext.GetLayer(0)
        ext_srs = ly_ext.GetSpatialRef()
        ext_srs.AutoIdentifyEPSG()
        ext_srs_epsg = 'EPSG:' + ext_srs.GetAuthorityCode(None)

    if not os.path.exists(mosd):
        os.mkdir(mosd)
    for mos in toMosaic:
        if sen == 'S2':
            obd = os.path.basename(mos[0])[0:-5] + 'MOSAIC'
            if 'THEIA' in obd:
                obd = obd[0:13] + obd[15:]
            od = mosd + obd
            if not os.path.exists(od):
                os.mkdir(od)
            else:
                if not force:
                    warnings.warn('S2: directory ' + od + 'exists, skipping.')
                    continue
            if onstack:
                fns = []
                for d in mos:
                    if os.path.exists(d + '/tileInfo.json') and not os.path.exists(od + '/tileInfo.json'):
                        copyfile(d + '/tileInfo.json', od + '/tileInfo.json')
                    else:
                        mtd = glob.glob(d + '/*_MTD_ALL.xml')
                        if len(mtd) == 1 and not os.path.exists(od + '/' + 'MOSAIC_MTD_ALL.xml'):
                            copyfile(mtd[0], od + '/' + 'MOSAIC_MTD_ALL.xml')
                    #msk_fn = glob.glob(d + '/' + '*MASK*.tif')[0]
                    lst = glob.glob(d + '/' + '*MASK*.tif')
                    if 'GAPMASK' in lst[0]:
                        msk_fn = lst[1]
                    else:
                        msk_fn = lst[0]
                    stack_fn = [ff for ff in glob.glob(d + '/' + '*STACK*.tif') if '_MASKED' not in ff][0]
                    #stack_fn =  glob.glob(d + '/' + 'STACK*.tif')[0]
                    mstack_fn = stack_fn.replace('STACK','STACK_MASKED')
                    subprocess.call(['otbcli_ManageNoData', '-in', stack_fn, '-mode', 'apply', '-mode.apply.mask', msk_fn, '-out', mstack_fn, 'uint16'],shell=sh)
                    if ext_srs_epsg is None:
                        ext_srs_epsg = checkSRS(mstack_fn)
                    elif checkSRS(mstack_fn) != ext_srs_epsg:
                        mstack_fn_rep = mstack_fn.replace('.tif', '_rep.tif')
                        cmd = ['gdalwarp', '-t_srs', ext_srs_epsg, mstack_fn, mstack_fn_rep]
                        subprocess.call(cmd, shell=sh)
                        os.remove(mstack_fn)
                        os.rename(mstack_fn_rep,mstack_fn)
                    fns.append(mstack_fn)
                ofn = od + '/' + os.path.basename(fns[0]).replace('STACK_MASKED','STACK')
                mfn = od + '/' + os.path.basename(fns[0]).replace('STACK_MASKED','MASK')
                gmfn = od + '/' + os.path.basename(fns[0]).replace('STACK_MASKED', 'GAPMASK')
                cmd = ['otbcli_Mosaic', '-il'] + fns + ['-out', ofn, 'uint16']
                subprocess.call(cmd, shell=sh)
                cmd = ['otbcli_ManageNoData', '-in', ofn, '-out', mfn, 'uint8', '-mode', 'buildmask']
                subprocess.call(cmd, shell=sh)
                if not os.path.exists(env_fn):
                    cmd = ['otbcli_Rasterization', '-in', ext_fn, '-out', env_fn, 'uint8',
                           '-im', ofn, '-mode.binary.foreground', '1']
                    subprocess.call(cmd,shell=sh)
                cmd = ['otbcli_BandMathX', '-il', env_fn, mfn, '-exp', '{im1b1 != 0 && im2b1 == 0}', '-out', gmfn, 'uint8']
                subprocess.call(cmd, shell=sh)
                for f in fns:
                    os.remove(f)
            else:
                #Following does not work for THEIA images!!
                bnds = ['B01','B02','B03','B04','B05','B06','B07','B08','B8A','B09','B10','B11','B12']
                for b in bnds:
                    fns = []
                    for d in mos:
                        # TO DO: manage different projections for band-by-band mode
                        if not os.path.exists(od + '/tileInfo.json'):
                            copyfile(d + '/tileInfo.json', od + '/tileInfo.json')
                        fns.append(glob.glob(d + '/' + b + '*.tif')[0])
                    ofn = os.path.basename(fns[0])
                    cmd = ['otbcli_Mosaic', '-il'] + fns + ['-out', od + '/' + ofn, 'uint16']
                    subprocess.call(cmd, shell=sh)
        elif sen == "L8":
            od = mosd + os.path.basename(mos[0])[0:-6] + 'MOSAIC'
            if not os.path.exists(od):
                os.mkdir(od)
            else:
                if not force:
                    warnings.warn('L8: directory ' + od + 'exists, skipping.')
                    continue
            if onstack:
                fns = []
                for d in mos:
                    if not os.path.exists(od + '/MOSAIC_MTL.txt'):
                        mtl = glob.glob(d + '/*_MTL*.txt')[0]
                        copyfile(mtl, od + '/MOSAIC_MTL.txt')
                    msk_fn = glob.glob(d + '/' + '*_MASK_HR_*.TIF')[0]
                    stack_fn = [ff for ff in glob.glob(d + '/' + '*_PS_TOA_*.TIF') if '_MASKED' not in ff][0]
                    mstack_fn = stack_fn.replace('PS_TOA', 'PS_TOA_MASKED')
                    ndv = getNoDataValue(stack_fn)
                    setNoDataValue(stack_fn,0)
                    cmd = ['otbcli_ManageNoData', '-in', stack_fn, '-mode', 'apply', '-mode.apply.mask', msk_fn,
                     '-out', mstack_fn, 'uint16']
                    subprocess.call(cmd, shell=sh)
                    setNoDataValue(stack_fn,ndv)
                    if ext_srs_epsg is None:
                        ext_srs_epsg = checkSRS(mstack_fn)
                    elif checkSRS(mstack_fn) != ext_srs_epsg:
                        mstack_fn_rep = mstack_fn.replace('.TIF', '_rep.TIF')
                        cmd = ['gdalwarp', '-t_srs', ext_srs_epsg, mstack_fn, mstack_fn_rep]
                        subprocess.call(cmd, shell=sh)
                        os.remove(mstack_fn)
                        os.rename(mstack_fn_rep, mstack_fn)
                    fns.append(mstack_fn)
                ofn = od + '/' + os.path.basename(fns[0]).replace('PS_TOA_MASKED', 'PS_TOA')
                mfn = od + '/' + os.path.basename(fns[0]).replace('PS_TOA_MASKED', 'MASK_HR')
                gmfn = od + '/' + os.path.basename(fns[0]).replace('PS_TOA_MASKED', 'GAPMASK')
                cmd = ['otbcli_Mosaic', '-il'] + fns + ['-out', ofn, 'uint16']
                subprocess.call(cmd, shell=sh)
                #cmd = ['otbcli_ManageNoData', '-in', ofn, '-out', mfn, '-mode', 'buildmask']
                #No-data checked in B4 only (avoid problems with B9 in TOA)
                cmd = ['otbcli_BandMathX', '-il', ofn, '-out', mfn, 'uint8', '-exp', '{im1b4 != ' + str(getNoDataValue(ofn)) + '}']
                subprocess.call(cmd, shell=sh)
                if not os.path.exists(env_fn):
                    cmd = ['otbcli_Rasterization', '-in', ext_fn, '-out', env_fn, 'uint8',
                           '-im', ofn, '-mode.binary.foreground', '1']
                    subprocess.call(cmd, shell=sh)
                cmd = ['otbcli_BandMathX', '-il', env_fn, mfn, '-exp', '{im1b1 != 0 && im2b1 == 0}', '-out', gmfn,
                       'uint8']
                subprocess.call(cmd, shell=sh)
                for f in fns:
                    os.remove(f)
            else:
                bnds = ['B1', 'B2', 'B3', 'B4', 'B5', 'B6', 'B7', 'B8', 'B9', 'B10', 'B11']
                for b in bnds:
                    fns = []
                    for d in mos:
                        # TO DO: manage different projections for band-by-band mode
                        if not os.path.exists(od + '/MOSAIC_MTL.txt'):
                            mtl = glob.glob(d + '/*_MTL*.txt')[0]
                            copyfile(mtl, od + '/MOSAIC_MTL.txt')
                        fns.append(glob.glob(d + '/' + b + '*.TIF')[0])
                    ofn = os.path.basename(fns[0])
                    cmd = ['otbcli_Mosaic', '-il'] + fns + ['-out', od + '/' + ofn, 'uint16']
                    subprocess.call(cmd, shell=sh)
        else:
            sys.exit('Not implemented yet')

if __name__ == '__main__':
    if len(sys.argv) < 2:
        sys.exit(
            'Usage: python [--onstack] [--extent <extent shapefile>] [--force] genMosaicScript.py <images root>\nWarning! Extent is mandatory if --onstack is set, and must be the same used for tile clipping.')
    else:
        main(sys.argv[1:])