basic.py

import os
import json
import pickle
import glob
import sys
from Workflows.operations import *
from Learning.ObjectBased import ObjectBasedClassifier
from Postprocessing import Report, MapFormatting

def unroll_file_list(lst):
    out_lst = []
    for f in lst:
        out_lst.extend(sorted(glob.glob(f)))
    return out_lst

def check_step(step, val):
    if not val:
        print("[MORINGA-ERR] : ***** ERROR ON STEP {} *****".format(step))
        sys.exit(step)

def process_timeseries(oroot, d, ts_lst_pkl):
    ts_lst = []
    for ts in d['timeseries']:
        if ts['type'] == 's2':
            print('[MORINGA-INFO] : Preprocessing {} from {}'.format(ts['type'], ts['provider']))
            ots = os.path.join(oroot, 'timeseries/' + ts['type'] + ts['provider'])
            os.makedirs(ots, exist_ok=True)
            ts_lst.append(preprocess_s2(ts['path'],
                                        ots,
                                        roi=d['roi'],
                                        output_dates_file=ts['output_dates_file'],
                                        provider=ts['provider']))
        elif ts['type'] == 's1':
            print('[MORINGA-INFO] : Preprocessing {} from {}'.format(ts['type'], ts['provider']))
            ots = os.path.join(oroot, 'timeseries/' + ts['type'] + ts['provider'])
            os.makedirs(ots, exist_ok=True)
            dem, geoid = None, None
            if ts['provider'] == 'native':
                dem, geoid = d['dem']['folder'], d['dem']['geoid']
                assert(os.path.exists(dem) and os.path.exists(geoid))
            ts_lst.append(preprocess_s1(ts['path'],
                                        roi=ts['ref_img'],
                                        out_fld=ots,
                                        dem_fld=dem,
                                        geoid=geoid,
                                        provider=ts['provider']
                                        ))
        elif ts['type'] == 'planetmosaics':
            print('[MORINGA-INFO] : Preprocessing {}'.format(ts['type']))
            ots = os.path.join(oroot, 'timeseries/' + ts['type'])
            os.makedirs(ots, exist_ok=True)
            ts_lst.append(preprocess_planet(ts['path'],
                                            ots))
        elif ts['type'].startswith('landsat'):
            print('[MORINGA-INFO] : Preprocessing {}'.format(ts['type']))
            ots = os.path.join(oroot, 'timeseries/' + ts['type'])
            os.makedirs(ots, exist_ok=True)
            ts_lst.append(preprocess_landsat(ts['path'],
                                             ots,
                                             roi=d['roi'],
                                             output_dates_file=ts['output_dates_file'],
                                             which=ts['type']))

        else:
            raise ValueError('TimeSeries type not yet supported.')
    with open(ts_lst_pkl, 'wb') as ts_save:
        pickle.dump(ts_lst, ts_save)
    return os.path.exists(ts_lst_pkl)

def perform_segmentation(ofn, d):
    print('[MORINGA-INFO] : Performing segmentation')
    os.makedirs(os.path.dirname(ofn), exist_ok=True)
    run_segmentation(d['segmentation']['src'],
                     d['segmentation']['th'],
                     d['segmentation']['cw'],
                     d['segmentation']['sw'],
                     ofn,
                     n_first_iter=d['segmentation']['n_first_iter'],
                     margin=d['segmentation']['margin'],
                     roi=d['roi'],
                     n_proc=d['segmentation']['n_proc'],
                     light=d['segmentation']['light_mode'])
    return os.path.exists(ofn)

def train_valid_workflow(seg, ts_lst_pkl, d, m_file):
    assert (os.path.exists(seg))
    assert (os.path.exists(ts_lst_pkl))
    print('[MORINGA-INFO] : Running Training/Validation Workflow')
    with open(ts_lst_pkl, 'rb') as ts_save:
        ts_lst = pickle.load(ts_save)
    obc = ObjectBasedClassifier(seg,
                                ts_lst,
                                unroll_file_list(d['userfeat']),
                                reference_data=d['ref_db']['path'],
                                ref_class_field=d['ref_db']['fields'])

    obc.gen_k_folds(5, class_field=d['ref_db']['fields'][-1],
                    augment=d['training']['augment_if_missing'], 
                    min_samples_per_class=d['training']['augment_if_missing'])

    if 'export_training_base' in d['training'].keys() and d['training']['export_training_base'] is True:
        obc.save_training_base('{}/_side/training_base.pkl'.format(os.path.join(d['output_path'], d['chain_name'])))
        print('[MORINGA-INFO] : Training base export completed.')

    ok = True

    for i,cf in enumerate(d['ref_db']['fields']):
        if d['training']['classifier'] == 'rf':
            m, s, r = obc.train_RF(d['training']['parameters']['n_trees'], class_field=cf, return_true_vs_pred=True)
            m_dict = {'model': m, 'results': r, 'summary': s,
                      'perc2':obc.training_base['perc2'], 'perc98':obc.training_base['perc98']}
            os.makedirs(os.path.dirname(m_file[i]), exist_ok=True)
            with open(m_file[i], 'wb') as mf:
                pickle.dump(m_dict, mf)
            ok = ok and os.path.exists(m_file[i])
    return ok

def classify(seg, ts_lst_pkl, m_files, d, map_files):
    assert (os.path.exists(seg))
    assert (os.path.exists(ts_lst_pkl))
    for m_file in m_files:
        assert (os.path.exists(m_file))
    print('[MORINGA-INFO] : Performing classification')
    with open(ts_lst_pkl, 'rb') as ts_save:
        ts_lst = pickle.load(ts_save)
    obc = ObjectBasedClassifier(seg,
                                ts_lst,
                                unroll_file_list(d['userfeat']))
    models = []
    for m_file in m_files:
        with open(m_file, 'rb') as mf:
            m_dict = pickle.load(mf)
        models.append(m_dict['model'])
    perc = [m_dict['perc2'], m_dict['perc98']]
    obc.classify(models, perc=perc, output_files=map_files)
    return all([os.path.exists(x) for x in map_files])

def report(map_files, m_files, d, report_files):
    print('[MORINGA-INFO] : Generating report(s)')
    for map_file, palette_fn, m_file, report_file in zip(map_files, d['map_output']['palette_files'], m_files, report_files):
        assert os.path.exists(map_file)
        assert os.path.exists(m_file)
        os.makedirs(report_file + '_figures', exist_ok=True)
        with open(m_file, 'rb') as mf:
            m_dict = pickle.load(mf)
        # For pdf report
        of = Report.generate_report_figures(
            map_file,
            palette_fn,
            m_dict['results'],
            m_dict['summary'],
            report_file + '_figures',
            d['chain_name'])
        Report.generate_pdf(of, report_file + '.pdf', d['chain_name'])
        Report.generate_text_report(m_dict['results'], m_dict['summary'], palette_fn, report_file + '.txt', d['chain_name'])
    return all([os.path.exists(x + '.pdf') for x in report_files] +
               [os.path.exists(x + '.txt') for x in report_files]) 

def basic(cfg, runlevel=1, single_step=False):
    os.environ['OTB_LOGGER_LEVEL'] = 'CRITICAL'
    with open(cfg,'r') as f:
        d = json.load(f)

    oroot = os.path.join(d['output_path'], d['chain_name'])
    oside = os.path.join(oroot, '_side')
    os.makedirs(oside, exist_ok=True)

    step = runlevel

    # Preprocess timeseries
    ts_lst_pkl = os.path.join(oside, 'time_series_list.pkl')
    if step == 1:
        print("[MORINGA-INFO] : ***** BEGIN STEP {} *****".format(step))
        check_step(step, process_timeseries(oroot, d, ts_lst_pkl))
        step += 1
        if single_step:
            return

    # Segmentation
    seg = os.path.join(oroot, 'segmentation/{}_obj_layer.tif'.format(d['chain_name']))
    if step == 2:
        print("[MORINGA-INFO] : ***** BEGIN STEP {} *****".format(step))
        check_step(step, perform_segmentation(seg, d))
        step += 1
        if single_step:
            return

    # Training/Validation Workflow
    m_files = []
    for cf in d['ref_db']['fields']:
        m_files.append(os.path.join(oroot, 'model/model_{}.pkl'.format(cf)))
    if step == 3:
        print("[MORINGA-INFO] : ***** BEGIN STEP {} *****".format(step))
        check_step(step, train_valid_workflow(seg, ts_lst_pkl, d, m_files))
        step += 1
        if single_step:
            return

    # Classification
    map_files = []
    for cf in d['ref_db']['fields']:
        map_files.append(os.path.join(oroot, 'maps/{}_map_{}.tif'.format(d['chain_name'],cf)))
    if step == 4:
        print("[MORINGA-INFO] : ***** BEGIN STEP {} *****".format(step))
        check_step(step, classify(seg, ts_lst_pkl, m_files, d, map_files))
        for m,p in zip(map_files, d['map_output']['palette_files']):
            MapFormatting.create_qgs_style(m,p)
        step += 1
        if single_step:
            return

    # Report
    report_fn = []
    for cf in d['ref_db']['fields']:
        report_fn.append(os.path.join(oroot, 'reports/{}_report_{}'.format(d['chain_name'],cf)))
    if step == 5:
        print("[MORINGA-INFO] : ***** BEGIN STEP {} *****".format(step))
        check_step(step, report(map_files, m_files, d, report_fn))

    print("[MORINGA-INFO] : ***** PROCESS FINISHED *****".format(step))
    return