s2planetary.py

from TimeSeries.s2sen2cor import *
import planetary_computer as PC
from pystac_client import Client
import rasterio
import rasterio.mask
from Common.geotools import get_query_bbox
from pyproj import Transformer as T
from shapely.geometry import Polygon
import tqdm
from datetime import *
import geopandas as gpd
from scipy.interpolate import interp1d


def fetch(shp, dt, output_fld, auth=None, band_list=None):

    bbox, i_bbox, shp_srs = get_query_bbox(shp, return_all=True)

    if auth is not None:
        os.environ['PC_SDK_SUBSCRIPTION_KEY'] = auth
    api = Client.open('https://planetarycomputer.microsoft.com/api/stac/v1', modifier=PC.sign_inplace)
    res = api.search(collections="sentinel-2-l2a", bbox=bbox, datetime=dt)
    lst = ['B02','B03','B04','B05','B06','B07','B08','B8A','B11','B12','SCL']
    if band_list is not None:
        lst = band_list

    ofiles = []
    prg = tqdm.tqdm(total=len(res.item_collection()) * len(lst), desc="Fetching from Planetary")
    for item in res.items():
        with rasterio.open(item.assets['B02'].href) as ds:
            img_srs = ds.crs.to_epsg()
        if shp_srs == img_srs:
            o_bbox = [20 * round(x / 20) for x in i_bbox]
        else:
            tr = T.from_crs(shp_srs, img_srs, always_xy=True)
            x_min, y_min = tr.transform(i_bbox[0], i_bbox[1])
            x_max, y_max = tr.transform(i_bbox[2], i_bbox[3])
            o_bbox = [20 * round(x / 20) for x in [x_min, y_min, x_max, y_max]]

        for a in lst:
            ofn = os.path.join(output_fld, '/'.join(item.assets[a].get_absolute_href().split('?')[0].split('/')[10:]))
            ofiles.append(ofn)
            os.makedirs(os.path.dirname(ofn),exist_ok=True)
            with rasterio.open(item.assets[a].href) as img:
                out_img, out_geot = rasterio.mask.mask(img,
                                                       [Polygon.from_bounds(o_bbox[0], o_bbox[1],
                                                                            o_bbox[2], o_bbox[3])],
                                                       crop=True)
                out_meta = img.meta
            out_meta.update({"driver": "GTiff",
                             "height": out_img.shape[1],
                             "width": out_img.shape[2],
                             "transform": out_geot})
            with rasterio.open(ofn, "w", **out_meta) as dest:
                dest.write(out_img)
            prg.update()
    prg.close()

    if band_list is None:
        return S2PlaneteryPipeline(output_fld)
    else:
        warnings.warn("Queried for a non-default band list. Skipping pipeline setup.")
        return ofiles

def split_dt_monthly(dt):
    sdt, edt = [datetime.strptime(x,'%Y-%m-%d') for x in dt.split('/')]
    dts = [sdt]
    cdt = sdt
    while cdt < edt:
        cdt += timedelta(weeks=4)
        dts.append(cdt)
    if dts[-1] > edt:
        dts = dts[:-1]
    if dts[-1] < edt:
        dts.append(edt+timedelta(1))
    out = []
    for i in range(len(dts)-1):
        out.append(datetime.strftime(dts[i], '%Y-%m-%d') + '/' + datetime.strftime(dts[i+1]-timedelta(1), '%Y-%m-%d'))
    return out
    
def sample(shp, dt, band_list, field_name, temp_dir = '/tmp', strategy='all', auth = None):
    fd = dt.split('/')[0]
    fd = datetime.strptime(fd,'%Y-%m-%d')
    ed = fd + timedelta(5)
    ndt = datetime.strftime(fd, '%Y-%m-%d') + '/' + datetime.strftime(ed, '%Y-%m-%d')
    #fn = fetch(shp, ndt, temp_dir, band_list=[band_list[0]])
    fn = ['/work/temp/S2B_MSIL2A_20210104T103329_N0212_R108_T30PVT_20210121T134941.SAFE/GRANULE/L2A_T30PVT_A020011_20210104T104033/IMG_DATA/R10m/T30PVT_20210104T103329_B04_10m.tif',
          '/work/temp/S2B_MSIL2A_20210104T103329_N0212_R108_T30PUT_20210105T133826.SAFE/GRANULE/L2A_T30PUT_A020011_20210104T104033/IMG_DATA/R10m/T30PUT_20210104T103329_B04_10m.tif']
    
    mos = otb.Registry.CreateApplication('Mosaic')
    mos.SetParameterStringList('il', fn)
    mos.Execute()
    sts = otb.Registry.CreateApplication('PolygonClassStatistics')
    sts.SetParameterInputImage('in', mos.GetParameterOutputImage('out'))
    sts.SetParameterString('vec', shp)
    sts.SetParameterString('field', field_name)
    sts.SetParameterString('out', os.path.join(temp_dir, 'stats.xml'))
    sts.ExecuteAndWriteOutput()

    stats_fn = os.path.join(temp_dir, 'stats.xml')
    samples_fn = os.path.join(temp_dir, 'samples.sqlite')

    mos = otb.Registry.CreateApplication('Mosaic')
    mos.SetParameterStringList('il', fn)
    mos.Execute()
    smp = otb.Registry.CreateApplication('SampleSelection')
    smp.SetParameterInputImage('in', mos.GetParameterOutputImage('out'))
    smp.SetParameterString('vec', shp)
    smp.SetParameterString('instats', stats_fn)
    smp.SetParameterString('field', field_name)
    if isinstance(strategy, str):
        smp.SetParameterString('strategy', strategy)
    elif isinstance(strategy, int):
        smp.SetParameterString('strategy','constant')
        smp.SetParameterInt('strategy.constant.nb', strategy)
    smp.SetParameterString('out', samples_fn)
    smp.ExecuteAndWriteOutput()

    ds = gpd.read_file(samples_fn)
    coords = list(zip(list(ds.geometry.x),list(ds.geometry.y)))
    bbox, _, _ = get_query_bbox(shp, return_all=True)
    if auth is not None:
        os.environ['PC_SDK_SUBSCRIPTION_KEY'] = auth
    
    all_arr = []
    dates = []
    for dti in split_dt_monthly(dt):

        api = Client.open('https://planetarycomputer.microsoft.com/api/stac/v1', modifier=PC.sign_inplace)
        res = api.search(collections="sentinel-2-l2a", bbox=bbox, datetime=dti, sortby='datetime')
        found = res.item_collection()

        if 'SCL' not in band_list:
            band_list.append('SCL')
        arr = {x : np.empty((ds.shape[0],len(found))) for x in band_list}
        i = 0
        prg = tqdm.tqdm(total=len(found), desc="Sampling from Planetary on period {}".format(dti))
        for item in found:
            dates.append(item.datetime)
            for band in band_list:
                with rasterio.open(item.assets[band].href) as src:
                    arr[band][:,i] = np.array([x for x in src.sample(coords)]).flatten()
            i += 1
            prg.update()
        prg.close()
        api = None
        all_arr.append(arr)

    
    days = np.array([(x-dates[0]).days for x in dates])
    final_arr = {}
    for b in band_list:
        final_arr[b] = np.hstack([x[b] for x in all_arr])
    
    band_list.remove('SCL')
    final_arr['SCL'] = np.invert(np.isin(final_arr['SCL'],[0,1,3,8,9,10]))
    for i in tqdm.tqdm(range(ds.shape[0]), desc="Interpolating time series"):
        x = days[final_arr['SCL'][i,:]]
        for b in band_list:
            y = final_arr[b][i,final_arr['SCL'][i,:]]
            f = interp1d(x,y, fill_value='extrapolate')
            final_arr[b][i,:] = f(days)
    
    _, ii = np.unique(days, return_index=True)
    for b in band_list:
        final_arr[b] = final_arr[b][:,ii]

    #[os.remove(x) for x in fn]
    os.remove(stats_fn)
    os.remove(samples_fn)

    return final_arr, dates[ii], list(ds[field_name.lower()])

class S2PlanetaryTilePipeline(S2Sen2CorTilePipeline):
    NAME = 'S2-L2A-SEN2COR-PLANETARY'
    PTRN_10m = ['GRANULE/*/IMG_DATA/R10m/*_B02_10m.tif',
                'GRANULE/*/IMG_DATA/R10m/*_B03_10m.tif',
                'GRANULE/*/IMG_DATA/R10m/*_B04_10m.tif',
                'GRANULE/*/IMG_DATA/R10m/*_B08_10m.tif']
    PTRN_20m = ['GRANULE/*/IMG_DATA/R20m/*_B05_20m.tif',
                'GRANULE/*/IMG_DATA/R20m/*_B06_20m.tif',
                'GRANULE/*/IMG_DATA/R20m/*_B07_20m.tif',
                'GRANULE/*/IMG_DATA/R20m/*_B8A_20m.tif',
                'GRANULE/*/IMG_DATA/R20m/*_B11_20m.tif',
                'GRANULE/*/IMG_DATA/R20m/*_B12_20m.tif']
    PTRN_msk = ['GRANULE/*/IMG_DATA/R20m/*_SCL_20m.tif']
    MERG_msk = ['min']
    PTRN_ful = PTRN_10m[0:3] + PTRN_20m[0:3] + [PTRN_10m[3]] + PTRN_20m[3:]

class S2PlaneteryPipeline(S2Sen2CorPipeline):

    S2TilePipeline = S2PlanetaryTilePipeline
    _check = S2TilePipeline._check
    _tile_id = S2TilePipeline._tile_id