def main():
    argparser = argparse.ArgumentParser(usage="Used to determine the postprocessing for a trained model. Useful for "
                                              "when the best configuration (2d, 3d_fullres etc) as selected manually.")
    argparser.add_argument("-m", type=str, required=True, help="U-Net model (2d, 3d_lowres, 3d_fullres or "
                                                               "3d_cascade_fullres)")
    argparser.add_argument("-t", type=str, required=True, help="Task name or id")
    argparser.add_argument("-tr", type=str, required=False, default=None,
                           help="nnUNetTrainer class. Default: %s, unless 3d_cascade_fullres "
                                "(then it's %s)" % (default_trainer, default_cascade_trainer))
    argparser.add_argument("-pl", type=str, required=False, default=default_plans_identifier,
                           help="Plans name, Default=%s" % default_plans_identifier)
    argparser.add_argument("-val", type=str, required=False, default="validation_raw",
                           help="Validation folder name. Default: validation_raw")

    args = argparser.parse_args()
    model = args.m
    task = args.t
    trainer = args.tr
    plans = args.pl
    val = args.val

    if not task.startswith("Task"):
        task_id = int(task)
        task = convert_id_to_task_name(task_id)

    if trainer is None:
        if model == "3d_cascade_fullres":
            trainer = "nnUNetTrainerV2CascadeFullRes"
        else:
            trainer = "nnUNetTrainerV2"

    folder = get_output_folder_name(model, task, trainer, plans, None)

    consolidate_folds(folder, val)
Ejemplo n.º 2
0
def main():
    import argparse
    parser = argparse.ArgumentParser(
        usage=
        "This is intended to identify the best model based on the five fold "
        "cross-validation. Running this script requires all models to have been run "
        "already. This script will summarize the results of the five folds of all "
        "models in one json each for easy interpretability")

    parser.add_argument(
        "-m",
        '--models',
        nargs="+",
        required=False,
        default=['2d', '3d_lowres', '3d_fullres', '3d_cascade_fullres'])
    parser.add_argument("-t", '--task_ids', nargs="+", required=True)

    parser.add_argument("-tr",
                        type=str,
                        required=False,
                        default=default_trainer,
                        help="nnUNetTrainer class. Default: %s" %
                        default_trainer)
    parser.add_argument(
        "-ctr",
        type=str,
        required=False,
        default=default_cascade_trainer,
        help="nnUNetTrainer class for cascade model. Default: %s" %
        default_cascade_trainer)
    parser.add_argument("-pl",
                        type=str,
                        required=False,
                        default=default_plans_identifier,
                        help="plans name, Default: %s" %
                        default_plans_identifier)
    parser.add_argument('-f',
                        '--folds',
                        nargs='+',
                        default=(0, 1, 2, 3, 4),
                        help="use this if you have non-standard folds")
    parser.add_argument(
        "--strict",
        required=False,
        default=False,
        action="store_true",
        help=
        "set this flag if you want this script to crash of one of the models is missing"
    )

    args = parser.parse_args()
    tasks = [int(i) for i in args.task_ids]

    models = args.models
    tr = args.tr
    trc = args.ctr
    strict = args.strict
    pl = args.pl
    folds = tuple(int(i) for i in args.folds)

    validation_folder = "validation_raw"

    # this script now acts independently from the summary jsons. That was unnecessary
    id_task_mapping = {}
    # for each task, run ensembling using all combinations of two models
    for t in tasks:
        # first collect pure model performance (postprocessed)
        results = {}
        all_results = {}
        valid_models = []
        for m in models:
            try:
                if m == "3d_cascade_fullres":
                    trainer = trc
                else:
                    trainer = tr

                if t not in id_task_mapping.keys():
                    task_name = find_task_name(get_output_folder_name(m), t)
                    id_task_mapping[t] = task_name

                output_folder = get_output_folder_name(m, id_task_mapping[t],
                                                       trainer, pl)
                assert isdir(
                    output_folder
                ), "Output folder for model %s is missing, expected: %s" % (
                    m, output_folder)

                # we need a postprocessing_json for inference, so that must be present
                postprocessing_json = join(output_folder,
                                           "postprocessing.json")
                # we need cv_niftis_postprocessed to know the single model performance
                cv_niftis_folder = join(output_folder, "cv_niftis_raw")
                if not isfile(postprocessing_json) or not isdir(
                        cv_niftis_folder):
                    print(
                        "running missing postprocessing for %s and model %s" %
                        (id_task_mapping[t], m))
                    consolidate_folds(output_folder, folds=folds)
                assert isfile(
                    postprocessing_json
                ), "Postprocessing json missing, expected: %s" % postprocessing_json
                assert isdir(
                    cv_niftis_folder
                ), "Folder with niftis from CV missing, expected: %s" % cv_niftis_folder

                # obtain mean foreground dice
                summary_file = join(cv_niftis_folder, "summary.json")
                results[m] = get_mean_foreground_dice(summary_file)
                foreground_mean(summary_file)
                all_results[m] = load_json(summary_file)['results']['mean']
                valid_models.append(m)

            except Exception as e:
                if strict:
                    raise e
                else:
                    print("WARNING!")
                    print(e)

        # now run ensembling and add ensembling to results
        print("\nFound the following valid models:\n", valid_models)
        if len(valid_models) > 1:
            # for m1, m2 in combinations(valid_models, 2):
            #
            #     trainer_m1 = trc if m1 == "3d_cascade_fullres" else tr
            #     trainer_m2 = trc if m2 == "3d_cascade_fullres" else tr
            #
            #     ensemble_name = "ensemble_" + m1 + "__" + trainer_m1 + "__" + pl + "--" + m2 + "__" + trainer_m2 + "__" + pl
            #     output_folder_base = join(network_training_output_dir, "ensembles", id_task_mapping[t], ensemble_name)
            #     maybe_mkdir_p(output_folder_base)
            #
            #     network1_folder = get_output_folder_name(m1, id_task_mapping[t], trainer_m1, pl)
            #     network2_folder = get_output_folder_name(m2, id_task_mapping[t], trainer_m2, pl)
            #
            #     print("ensembling", network1_folder, network2_folder)
            #     ensemble2(network1_folder, network2_folder, output_folder_base, id_task_mapping[t], validation_folder, folds)
            #     # ensembling will automatically do postprocessingget_foreground_mean
            #
            #     # now get result of ensemble
            #     results[ensemble_name] = get_mean_foreground_dice(join(output_folder_base, "ensembled_raw", "summary.json"))
            #     summary_file = join(output_folder_base, "ensembled_raw", "summary.json")
            #     foreground_mean(summary_file)
            #     all_results[ensemble_name] = load_json(summary_file)['results']['mean']

            # for m1, m2, m3 in combinations(valid_models, 3):
            #
            #     trainer_m1 = trc if m1 == "3d_cascade_fullres" else tr
            #     trainer_m2 = trc if m2 == "3d_cascade_fullres" else tr
            #     trainer_m3 = trc if m3 == "3d_cascade_fullres" else tr
            #
            #     ensemble_name = "ensemble_" + m1 + "__" + trainer_m1 + "__" + pl + "--" + m2 + "__" + trainer_m2 + "__" + pl + "--" + m3 + "__" + trainer_m3 + "__" + pl
            #     output_folder_base = join(network_training_output_dir, "ensembles", id_task_mapping[t], ensemble_name)
            #     maybe_mkdir_p(output_folder_base)
            #
            #     network1_folder = get_output_folder_name(m1, id_task_mapping[t], trainer_m1, pl)
            #     network2_folder = get_output_folder_name(m2, id_task_mapping[t], trainer_m2, pl)
            #     network3_folder = get_output_folder_name(m3, id_task_mapping[t], trainer_m3, pl)
            #
            #     print("ensembling", network1_folder, network2_folder, network3_folder)
            #     ensemble3(network1_folder, network2_folder, network3_folder, output_folder_base, id_task_mapping[t], validation_folder, folds)
            #     # ensembling will automatically do postprocessingget_foreground_mean
            #
            #     # now get result of ensemble
            #     results[ensemble_name] = get_mean_foreground_dice(join(output_folder_base, "ensembled_raw", "summary.json"))
            #     summary_file = join(output_folder_base, "ensembled_raw", "summary.json")
            #     foreground_mean(summary_file)
            #     all_results[ensemble_name] = load_json(summary_file)['results']['mean']

            for m1, m2, m3, m4 in combinations(valid_models, 4):

                trainer_m1 = trc if m1 == "3d_cascade_fullres" else tr
                trainer_m2 = trc if m2 == "3d_cascade_fullres" else tr
                trainer_m3 = trc if m3 == "3d_cascade_fullres" else tr
                trainer_m4 = trc if m4 == "3d_cascade_fullres" else tr

                ensemble_name = "ensemble_" + m1 + "__" + trainer_m1 + "__" + pl + "--" + m2 + "__" + trainer_m2 + "__" + pl + "--" + m3 + "__" + trainer_m3 + "__" + pl + m4 + "__" + trainer_m4 + "__" + pl
                output_folder_base = join(network_training_output_dir,
                                          "ensembles", id_task_mapping[t],
                                          ensemble_name)
                maybe_mkdir_p(output_folder_base)

                network1_folder = get_output_folder_name(
                    m1, id_task_mapping[t], trainer_m1, pl)
                network2_folder = get_output_folder_name(
                    m2, id_task_mapping[t], trainer_m2, pl)
                network3_folder = get_output_folder_name(
                    m3, id_task_mapping[t], trainer_m3, pl)
                network4_folder = get_output_folder_name(
                    m4, id_task_mapping[t], trainer_m4, pl)

                print("ensembling", network1_folder, network2_folder,
                      network3_folder, network4_folder)
                ensemble4(network1_folder, network2_folder, network3_folder,
                          network4_folder, output_folder_base,
                          id_task_mapping[t], validation_folder, folds)
                # ensembling will automatically do postprocessingget_foreground_mean

                # now get result of ensemble
                results[ensemble_name] = get_mean_foreground_dice(
                    join(output_folder_base, "ensembled_raw", "summary.json"))
                summary_file = join(output_folder_base, "ensembled_raw",
                                    "summary.json")
                foreground_mean(summary_file)
                all_results[ensemble_name] = load_json(
                    summary_file)['results']['mean']

        # now print all mean foreground dice and highlight the best
        foreground_dices = list(results.values())
        best = np.max(foreground_dices)
        for k, v in results.items():
            print(k, v)

        predict_str = ""
        best_model = None
        for k, v in results.items():
            if v == best:
                print("%s submit model %s" % (id_task_mapping[t], k), v)
                best_model = k
                print(
                    "\nHere is how you should predict test cases. Run in sequential order and replace all input and output folder names with your personalized ones\n"
                )
                if k.startswith("ensemble"):
                    tmp = k[len("ensemble_"):]
                    model1, model2 = tmp.split("--")
                    m1, t1, pl1 = model1.split("__")
                    m2, t2, pl2 = model2.split("__")
                    predict_str += "nnUNet_predict -i FOLDER_WITH_TEST_CASES -o OUTPUT_FOLDER_MODEL1 -tr " + tr + " -ctr " + trc + " -m " + m1 + " -p " + pl + " -t " + \
                                   id_task_mapping[t] + "\n"
                    predict_str += "nnUNet_predict -i FOLDER_WITH_TEST_CASES -o OUTPUT_FOLDER_MODEL2 -tr " + tr + " -ctr " + trc + " -m " + m2 + " -p " + pl + " -t " + \
                                   id_task_mapping[t] + "\n"

                    predict_str += "nnUNet_ensemble -f OUTPUT_FOLDER_MODEL1 OUTPUT_FOLDER_MODEL2 -o OUTPUT_FOLDER -pp " + join(
                        network_training_output_dir, "ensembles",
                        id_task_mapping[t], k, "postprocessing.json") + "\n"
                else:
                    predict_str += "nnUNet_predict -i FOLDER_WITH_TEST_CASES -o OUTPUT_FOLDER_MODEL1 -tr " + tr + " -ctr " + trc + " -m " + k + " -p " + pl + " -t " + \
                                   id_task_mapping[t] + "\n"
                print(predict_str)

        summary_folder = join(network_training_output_dir, "ensembles",
                              id_task_mapping[t])
        maybe_mkdir_p(summary_folder)
        with open(join(summary_folder, "prediction_commands.txt"), 'w') as f:
            f.write(predict_str)

        num_classes = len(
            [i for i in all_results[best_model].keys() if i != 'mean'])
        with open(join(summary_folder, "summary.csv"), 'w') as f:
            f.write("model")
            for c in range(1, num_classes):
                f.write(",class%d" % c)
            f.write(",average")
            f.write("\n")
            for m in all_results.keys():
                f.write(m)
                for c in range(1, num_classes):
                    f.write(",%01.4f" % all_results[m][str(c)]["Dice"])
                f.write(",%01.4f" % all_results[m]['mean']["Dice"])
                f.write("\n")
Ejemplo n.º 3
0
def main():
    import argparse
    parser = argparse.ArgumentParser(
        usage=
        "This is intended to identify the best model based on the five fold "
        "cross-validation. Running this script requires all models to have been run "
        "already. This script will summarize the results of the five folds of all "
        "models in one json each for easy interpretability")

    parser.add_argument(
        "-m",
        '--models',
        nargs="+",
        required=False,
        default=['2d', '3d_lowres', '3d_fullres', '3d_cascade_fullres'])
    parser.add_argument("-t", '--task_ids', nargs="+", required=True)

    parser.add_argument("-tr",
                        type=str,
                        required=False,
                        default=default_trainer,
                        help="nnUNetTrainer class. Default: %s" %
                        default_trainer)
    parser.add_argument(
        "-ctr",
        type=str,
        required=False,
        default=default_cascade_trainer,
        help="nnUNetTrainer class for cascade model. Default: %s" %
        default_cascade_trainer)
    parser.add_argument("-pl",
                        type=str,
                        required=False,
                        default=default_plans_identifier,
                        help="plans name, Default: %s" %
                        default_plans_identifier)
    parser.add_argument('-f',
                        '--folds',
                        nargs='+',
                        default=(0, 1, 2, 3, 4),
                        help="Use this if you have non-standard "
                        "folds. Experienced users only.")
    parser.add_argument(
        '--disable_ensembling',
        required=False,
        default=False,
        action='store_true',
        help=
        'Set this flag to disable the use of ensembling. This will find the best single '
        'configuration for each task.')
    parser.add_argument(
        "--disable_postprocessing",
        required=False,
        default=False,
        action="store_true",
        help="Set this flag if you want to disable the use of postprocessing")

    args = parser.parse_args()
    tasks = [int(i) for i in args.task_ids]

    models = args.models
    tr = args.tr
    trc = args.ctr
    pl = args.pl
    disable_ensembling = args.disable_ensembling
    disable_postprocessing = args.disable_postprocessing
    folds = tuple(int(i) for i in args.folds)

    validation_folder = "validation_raw"

    # this script now acts independently from the summary jsons. That was unnecessary
    id_task_mapping = {}

    for t in tasks:
        # first collect pure model performance
        results = {}
        all_results = {}
        valid_models = []
        for m in models:
            if m == "3d_cascade_fullres":
                trainer = trc
            else:
                trainer = tr

            if t not in id_task_mapping.keys():
                task_name = find_task_name(get_output_folder_name(m), t)
                id_task_mapping[t] = task_name

            output_folder = get_output_folder_name(m, id_task_mapping[t],
                                                   trainer, pl)
            if not isdir(output_folder):
                raise RuntimeError(
                    "Output folder for model %s is missing, expected: %s" %
                    (m, output_folder))

            if disable_postprocessing:
                # we need to collect the predicted niftis from the 5-fold cv and evaluate them against the ground truth
                cv_niftis_folder = join(output_folder, 'cv_niftis_raw')

                if not isfile(join(cv_niftis_folder, 'summary.json')):
                    print(t, m, ': collecting niftis from 5-fold cv')
                    if isdir(cv_niftis_folder):
                        shutil.rmtree(cv_niftis_folder)

                    collect_cv_niftis(output_folder, cv_niftis_folder,
                                      validation_folder, folds)

                    niftis_gt = subfiles(join(output_folder, "gt_niftis"),
                                         suffix='.nii.gz',
                                         join=False)
                    niftis_cv = subfiles(cv_niftis_folder,
                                         suffix='.nii.gz',
                                         join=False)
                    if not all([i in niftis_gt for i in niftis_cv]):
                        raise AssertionError("It does not seem like you trained all the folds! Train " \
                                             "all folds first! There are %d gt niftis in %s but only " \
                                             "%d predicted niftis in %s" % (len(niftis_gt), niftis_gt,
                                                                            len(niftis_cv), niftis_cv))

                    # load a summary file so that we can know what class labels to expect
                    summary_fold0 = load_json(
                        join(output_folder, "fold_%d" % folds[0],
                             validation_folder,
                             "summary.json"))['results']['mean']
                    # read classes from summary.json
                    classes = tuple((int(i) for i in summary_fold0.keys()))

                    # evaluate the cv niftis
                    print(t, m, ': evaluating 5-fold cv results')
                    evaluate_folder(join(output_folder, "gt_niftis"),
                                    cv_niftis_folder, classes)

            else:
                postprocessing_json = join(output_folder,
                                           "postprocessing.json")
                cv_niftis_folder = join(output_folder, "cv_niftis_raw")

                # we need cv_niftis_postprocessed to know the single model performance. And we need the
                # postprocessing_json. If either of those is missing, rerun consolidate_folds
                if not isfile(postprocessing_json) or not isdir(
                        cv_niftis_folder):
                    print(
                        "running missing postprocessing for %s and model %s" %
                        (id_task_mapping[t], m))
                    consolidate_folds(output_folder, folds=folds)

                assert isfile(
                    postprocessing_json
                ), "Postprocessing json missing, expected: %s" % postprocessing_json
                assert isdir(
                    cv_niftis_folder
                ), "Folder with niftis from CV missing, expected: %s" % cv_niftis_folder

            # obtain mean foreground dice
            summary_file = join(cv_niftis_folder, "summary.json")
            results[m] = get_mean_foreground_dice(summary_file)
            foreground_mean(summary_file)
            all_results[m] = load_json(summary_file)['results']['mean']
            valid_models.append(m)

        if not disable_ensembling:
            # now run ensembling and add ensembling to results
            print(
                "\nI will now ensemble combinations of the following models:\n",
                valid_models)
            if len(valid_models) > 1:
                for m1, m2 in combinations(valid_models, 2):

                    trainer_m1 = trc if m1 == "3d_cascade_fullres" else tr
                    trainer_m2 = trc if m2 == "3d_cascade_fullres" else tr

                    ensemble_name = "ensemble_" + m1 + "__" + trainer_m1 + "__" + pl + "--" + m2 + "__" + trainer_m2 + "__" + pl
                    output_folder_base = join(network_training_output_dir,
                                              "ensembles", id_task_mapping[t],
                                              ensemble_name)
                    os.makedirs(output_folder_base, exist_ok=True)

                    network1_folder = get_output_folder_name(
                        m1, id_task_mapping[t], trainer_m1, pl)
                    network2_folder = get_output_folder_name(
                        m2, id_task_mapping[t], trainer_m2, pl)

                    print("ensembling", network1_folder, network2_folder)
                    ensemble(network1_folder,
                             network2_folder,
                             output_folder_base,
                             id_task_mapping[t],
                             validation_folder,
                             folds,
                             allow_ensembling=not disable_postprocessing)
                    # ensembling will automatically do postprocessingget_foreground_mean

                    # now get result of ensemble
                    results[ensemble_name] = get_mean_foreground_dice(
                        join(output_folder_base, "ensembled_raw",
                             "summary.json"))
                    summary_file = join(output_folder_base, "ensembled_raw",
                                        "summary.json")
                    foreground_mean(summary_file)
                    all_results[ensemble_name] = load_json(
                        summary_file)['results']['mean']

        # now print all mean foreground dice and highlight the best
        foreground_dices = list(results.values())
        best = np.max(foreground_dices)
        for k, v in results.items():
            print(k, v)

        predict_str = ""
        best_model = None
        for k, v in results.items():
            if v == best:
                print("%s submit model %s" % (id_task_mapping[t], k), v)
                best_model = k
                print(
                    "\nHere is how you should predict test cases. Run in sequential order and replace all input and output folder names with your personalized ones\n"
                )
                if k.startswith("ensemble"):
                    tmp = k[len("ensemble_"):]
                    model1, model2 = tmp.split("--")
                    m1, t1, pl1 = model1.split("__")
                    m2, t2, pl2 = model2.split("__")
                    predict_str += "nnUNet_predict -i FOLDER_WITH_TEST_CASES -o OUTPUT_FOLDER_MODEL1 -tr " + tr + " -ctr " + trc + " -m " + m1 + " -p " + pl + " -t " + \
                                   id_task_mapping[t] + "\n"
                    predict_str += "nnUNet_predict -i FOLDER_WITH_TEST_CASES -o OUTPUT_FOLDER_MODEL2 -tr " + tr + " -ctr " + trc + " -m " + m2 + " -p " + pl + " -t " + \
                                   id_task_mapping[t] + "\n"

                    if not disable_postprocessing:
                        predict_str += "nnUNet_ensemble -f OUTPUT_FOLDER_MODEL1 OUTPUT_FOLDER_MODEL2 -o OUTPUT_FOLDER -pp " + join(
                            network_training_output_dir, "ensembles",
                            id_task_mapping[t], k,
                            "postprocessing.json") + "\n"
                    else:
                        predict_str += "nnUNet_ensemble -f OUTPUT_FOLDER_MODEL1 OUTPUT_FOLDER_MODEL2 -o OUTPUT_FOLDER\n"
                else:
                    predict_str += "nnUNet_predict -i FOLDER_WITH_TEST_CASES -o OUTPUT_FOLDER_MODEL1 -tr " + tr + " -ctr " + trc + " -m " + k + " -p " + pl + " -t " + \
                                   id_task_mapping[t] + "\n"
                print(predict_str)

        summary_folder = join(network_training_output_dir, "ensembles",
                              id_task_mapping[t])
        os.makedirs(summary_folder, exist_ok=True)
        with open(join(summary_folder, "prediction_commands.txt"), 'w') as f:
            f.write(predict_str)

        num_classes = len([
            i for i in all_results[best_model].keys()
            if i != 'mean' and i != '0'
        ])
        with open(join(summary_folder, "summary.csv"), 'w') as f:
            f.write("model")
            for c in range(1, num_classes + 1):
                f.write(",class%d" % c)
            f.write(",average")
            f.write("\n")
            for m in all_results.keys():
                f.write(m)
                for c in range(1, num_classes + 1):
                    f.write(",%01.4f" % all_results[m][str(c)]["Dice"])
                f.write(",%01.4f" % all_results[m]['mean']["Dice"])
                f.write("\n")