main.py

import argparse
from feature_extraction import segment, merge, extract_features
from prediction import predict

def main(steps):
    image_dir = 'data/images'
    patient_info = 'data/patient_info.csv'
    seg_dir = 'data/segmentations'
    feature_dir = 'data'
    features = 'data/features.csv'
    model_path = 'data/model.json'
    pred_dir = 'data'
    if 'segment' in steps:
        segment(image_dir, seg_dir)
    if 'merge' in steps:
        merge(seg_dir)
    if 'extract' in steps:
        extract_features(image_dir, seg_dir, feature_dir)
    if 'predict' in steps:
        predict(features, patient_info, model_path, pred_dir)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--feature-extraction', action='store_true',
                        help='All feature extraction steps (segment, merge, extract)')
    parser.add_argument('--segment', action='store_true')
    parser.add_argument('--merge', action='store_true')
    parser.add_argument('--extract', action='store_true')
    parser.add_argument('--predict', action='store_true')
    args = vars(parser.parse_args())
    steps = []
    feature_extraction_steps = ['segment', 'merge', 'extract']
    for step in feature_extraction_steps:
        if args['feature_extraction'] or args[step]:
            steps.append(step)
    for step in ['predict']:
        if args[step]:
            steps.append(step)
    main(steps)