moringa.py

import os.path
import sys
import argparse
import OBIA.segmentation
import VHR.vhrbase
from TimeSeries import s2theia, s2planetary

def run_segmentation(img, threshold, cw, sw , out_seg,
                     n_first_iter, margin, roi, n_proc, memory,
                     remove_graph, force_parallel):
    params = OBIA.segmentation.LSGRMParams(threshold, cw, sw, n_first_iter, margin)
    OBIA.segmentation.lsgrm(img, params, out_seg, n_proc, memory, roi, remove_graph, force_parallel)
    return

def preprocess_spot67(in_fld, out_fld, dem_fld, geoid_file, skip_ps, compress,
                      clip, align_to, align_to_band, align_using_band):
    sp = VHR.vhrbase.SPOT67RasterPipeline(in_fld)
    sp.to_toa()
    sp.orthorectify(dem_fld, geoid_file)
    if clip is not None and os.path.exists(clip):
        sp.clip(clip)
    if not skip_ps:
        sp.write_outputs(out_fld, update_pipe=True, compress=compress)
        sp.pansharp()
    if align_to is not None and os.path.exists(align_to):
        sp.rigid_align(align_to, this_band=align_using_band-1, ref_band=align_to_band-1)
    sp.write_outputs(out_fld, compress=compress)
    return

def preprocess_s2(in_fld, out_fld, output_dates_file=None, roi=None, coregister_to=None, coregister_to_band=1,
                  coregister_using_band=3, provider='theia'):
    S2Processor = None
    if provider == 'theia':
        S2Processor = s2theia.S2TheiaPipeline
    elif provider == 'planetary':
        S2Processor = s2planetary.S2PlaneteryPipeline
    else:
        raise ValueError("Unsupported/non-valid provider")

    s2 = S2Processor(in_fld, temp_fld=out_fld)
    if roi is not None:
        s2.set_roi(roi)
    if output_dates_file is not None:
        s2.set_output_dates_by_file(output_dates_file)
    else:
        raise ValueError("Please provide path to a text file containing output dates.")

    s2.extract_feature_set(out_fld, store_gapfill=True, mosaicking='vrt')
    return

def main(args):
    parser = argparse.ArgumentParser(prog="moringa", add_help=False)
    subpar = parser.add_subparsers(dest="cmd")

    prepr = subpar.add_parser("preprocess_s2", help="Performs Moringa preset time series preprocessing.",
                              formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    prepr.add_argument("in_folder", type=str, help="Path to the folder containing (de-zipped) S2 images.")
    prepr.add_argument("out_folder", type=str, help="Path to the folder in which preprocessed stacks will be stored.")
    prepr.add_argument("--output_dates_file", type=str, default=None, help="Path to the text file containing output dates for temporal interpolation.")
    prepr.add_argument("--roi", type=str, default=None, help="Path to the ROI vector file.")
    prepr.add_argument("--coregister_to", type=str, default=None, help="Path to a reference image to which the stacks must be coregistered.")
    prepr.add_argument("--coregister_to_band", type=int, nargs='?', default=1, help="Band of reference image used for co-registration.")
    prepr.add_argument("--coregister_using_band", type=int, nargs='?', default=3, help="Band of current stack used for co-registration.")
    prepr.add_argument("--provider", type=str, default='theia', help="S2 image provider. Supported: 'theia', 'theial3a', 'sen2cor', 'planetary'")

    segmt = subpar.add_parser("segment", help="Performs (large scale Baatz-Shape) segmentation of an input image.",
                              formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    segmt.add_argument("img", type=str, help="Path to the image to segment.")
    segmt.add_argument("threshold", type=float, help="Threshold for the heterogeneity criterion in Baatz-Shape.")
    segmt.add_argument("outimg", type=str, help="Path to the output segmentation file (.tif, .shp, .gpkg, .gml).")
    segmt.add_argument("--cw", type=float, nargs="?", default=0.5, help="Color weight in Baatz-Shape criterion.")
    segmt.add_argument("--sw", type=float, nargs="?", default=0.5, help="Spatial weight in Baatz-Shape criterion.")
    segmt.add_argument("--n_first_iter", type=int, nargs="?", default=12, help="Number of iterations for parallel tile processing.")
    segmt.add_argument("--tile_margin", type=int, nargs="?", default=100, help="Margin for tile overlap.")
    segmt.add_argument("--roi", type=str, default=None, help="Vector file containing an ROI.")
    segmt.add_argument("--n_proc", type=int, help="Number of cores to use.")
    segmt.add_argument("--mem_limit", type=int, help="Memory limit in MB.")
    segmt.add_argument("--keep_graph", help="Keep the graph files (.bin) after segmentation.", action='store_true')
    segmt.add_argument("--force_parallel", help="Force the spot6/7 preprocess one-liner parallelization of the process even if the full graph fits in memory.", action='store_true')

    vhrprep = subpar.add_parser("preprocess_spot67", help="Perform baseline pre-processing of a SPOT6/7 scene in raster sensor.",
                                formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    vhrprep.add_argument("fld", type=str, help="Path to the root folder containing PAN and MS subfolders.")
    vhrprep.add_argument("out_fld", type=str, help="Path to the output folder for preprocessed images.")
    vhrprep.add_argument("dem_fld", type=str, help="Path to the folder containing DEM covering the scene in WGS84 projection.")
    vhrprep.add_argument("geoid", type=str, help="Path to the geoid file.")
    vhrprep.add_argument("--clip", type=str, default=None, help="Path to a vector file for clipping.")
    vhrprep.add_argument("--align_to", type=str, default=None, help="Path to a reference image to which the image must be aligned (rigid).")
    vhrprep.add_argument("--align_to_band", type=int, nargs='?', default=3, help="Band of reference image used for alignment.")
    vhrprep.add_argument("--align_using_band", type=int, nargs='?', default=1, help="Band of current image used for alignment.")
    vhrprep.add_argument("--skip_ps", help="Skip pansharpening step.", action='store_true')
    vhrprep.add_argument("--compress", help="Use lossless compresion on outputs.", action='store_true')



    if len(args) == 1:
        parser.print_help()
        sys.exit(0)

    arg = parser.parse_args()

    if arg.cmd == "preprocess_s2":
        preprocess_s2(arg.in_folder, arg.out_folder, output_dates_file=arg.output_dates_file, roi=arg.roi,
                      coregister_to=arg.coregister_to, provider=arg.provider)

    if arg.cmd == "segment":
        run_segmentation(arg.img, arg.threshold, arg.cw, arg.sw, arg.outimg, arg.n_first_iter, arg.tile_margin,
                         arg.roi, arg.n_proc, arg.mem_limit, not arg.keep_graph, arg.force_parallel)

    if arg.cmd == "preprocess_spot67":
        preprocess_spot67(arg.fld, arg.out_fld, arg.dem_fld, arg.geoid, arg.skip_ps, arg.compress,
                          arg.clip, arg.align_to, arg.align_to_band, arg.align_using_band)

    return 0

if __name__ == "__main__":
    main(sys.argv)