open-kbp/main.py at master · malonec3/open-kbp · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
import shutil

from provided_code.data_loader import DataLoader
from provided_code.dose_evaluation_class import EvaluateDose
from provided_code.general_functions import get_paths, make_directory_and_return_path
from provided_code.network_functions import PredictionModel

if __name__ == '__main__':
    # Define project directories
    # TODO: Must define the path of where the data is stored.
    primary_directory = '/Users/aaronbabier/Documents/GitHub/open-kbp-post' #'/home/user_name/open-kbp'  # directory where everything is stored
    # Define directory where given data is stored
    main_data_dir = '{}/provided-data'.format(primary_directory)
    training_data_dir = '{}/train-pats'.format(main_data_dir)
    validation_data_dir = '{}/validation-pats'.format(main_data_dir)
    testing_data_dir = '{}/test-pats'.format(main_data_dir)
    # Define hold out set
    test_time = False  # Only change this to True when the model has been fully tuned on the validation set
    # path where any data generated by this code (e.g., predictions, models) are stored
    results_dir = '{}/results'.format(primary_directory)
    # Name model to train and number of epochs to train it for
    prediction_name = 'baseline'
    number_of_training_epochs = 1  # This should probably be increased to 100-200 after your dry run

    # Prepare the data directory
    training_plan_paths = get_paths(training_data_dir, ext='')  # gets the path of each plan's directory

    # Train a model
    data_loader_train = DataLoader(training_plan_paths)
    dose_prediction_model_train = PredictionModel(data_loader_train, results_dir, model_name=prediction_name)
    dose_prediction_model_train.train_model(epochs=number_of_training_epochs, save_frequency=1, keep_model_history=1)

    # Define hold out set
    if test_time is False:
        hold_out_plan_paths = get_paths(validation_data_dir, ext='')  # list of paths used for held out validation
        stage_name = 'hold-out-validation'
    else:
        hold_out_plan_paths = get_paths(testing_data_dir, ext='')  # list of paths used for held out testing
        stage_name = 'hold-out-testing'

    # Predict dose for the held out set
    data_loader_hold_out = DataLoader(hold_out_plan_paths, mode_name='dose_prediction')
    dose_prediction_model_hold_out = PredictionModel(data_loader_hold_out, results_dir,
                                                     model_name=prediction_name, stage=stage_name)
    dose_prediction_model_hold_out.predict_dose(epoch=number_of_training_epochs)

    # Evaluate dose metrics
    data_loader_hold_out_eval = DataLoader(hold_out_plan_paths, mode_name='evaluation')  # Set data loader
    prediction_paths = get_paths(dose_prediction_model_hold_out.prediction_dir, ext='csv')
    hold_out_prediction_loader = DataLoader(prediction_paths, mode_name='predicted_dose')  # Set prediction loader
    dose_evaluator = EvaluateDose(data_loader_hold_out_eval, hold_out_prediction_loader)

    # print out scores if data was left for a hold out set
    if not data_loader_hold_out_eval.file_paths_list:
        print('No patient information was given to calculate metrics')
    else:
        dvh_score, dose_score = dose_evaluator.make_metrics()
        print('For this out-of-sample test:\n'
              '\tthe DVH score is {:.3f}\n '
              '\tthe dose score is {:.3f}'.format(dvh_score, dose_score))

    # Zip dose to submit
    submission_dir = make_directory_and_return_path('{}/submissions'.format(results_dir))
    shutil.make_archive('{}/{}'.format(submission_dir, prediction_name), 'zip', dose_prediction_model_hold_out.prediction_dir)