nnunet/inference/pretrained_models/collect_pretrained_models.py

#    Copyright 2020 Division of Medical Image Computing, German Cancer Research Center (DKFZ), Heidelberg, Germany
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.


import zipfile
from multiprocessing.pool import Pool

from batchgenerators.utilities.file_and_folder_operations import *
import shutil
from nnunet.paths import default_cascade_trainer, default_plans_identifier, default_trainer, network_training_output_dir
from nnunet.utilities.task_name_id_conversion import convert_id_to_task_name
from subprocess import call


def copy_fold(in_folder: str, out_folder: str):
    shutil.copy(join(in_folder, "debug.json"), join(out_folder, "debug.json"))
    shutil.copy(join(in_folder, "model_final_checkpoint.model"), join(out_folder, "model_final_checkpoint.model"))
    shutil.copy(join(in_folder, "model_final_checkpoint.model.pkl"),
                join(out_folder, "model_final_checkpoint.model.pkl"))
    shutil.copy(join(in_folder, "progress.png"), join(out_folder, "progress.png"))
    if isfile(join(in_folder, "network_architecture.pdf")):
        shutil.copy(join(in_folder, "network_architecture.pdf"), join(out_folder, "network_architecture.pdf"))


def copy_model(directory: str, output_directory: str):
    """

    :param directory: must have the 5 fold_X subfolders as well as a postprocessing.json and plans.pkl
    :param output_directory:
    :return:
    """
    expected_folders = ["fold_%d" % i for i in range(5)]
    assert all([isdir(join(directory, i)) for i in expected_folders]), "not all folds present"

    assert isfile(join(directory, "plans.pkl")), "plans.pkl missing"
    assert isfile(join(directory, "postprocessing.json")), "postprocessing.json missing"

    for e in expected_folders:
        maybe_mkdir_p(join(output_directory, e))
        copy_fold(join(directory, e), join(output_directory, e))

    shutil.copy(join(directory, "plans.pkl"), join(output_directory, "plans.pkl"))
    shutil.copy(join(directory, "postprocessing.json"), join(output_directory, "postprocessing.json"))


def copy_pretrained_models_for_task(task_name: str, output_directory: str,
                                    models: tuple = ("2d", "3d_lowres", "3d_fullres", "3d_cascade_fullres"),
                                    nnunet_trainer=default_trainer,
                                    nnunet_trainer_cascade=default_cascade_trainer,
                                    plans_identifier=default_plans_identifier):
    trainer_output_dir = nnunet_trainer + "__" + plans_identifier
    trainer_output_dir_cascade = nnunet_trainer_cascade + "__" + plans_identifier

    for m in models:
        to = trainer_output_dir_cascade if m == "3d_cascade_fullres" else trainer_output_dir
        expected_output_folder = join(network_training_output_dir, m, task_name, to)
        if not isdir(expected_output_folder):
            if m == "3d_lowres" or m == "3d_cascade_fullres":
                print("Task", task_name, "does not seem to have the cascade")
                continue
            else:
                raise RuntimeError("missing folder! %s" % expected_output_folder)
        output_here = join(output_directory, m, task_name, to)
        maybe_mkdir_p(output_here)
        copy_model(expected_output_folder, output_here)


def check_if_valid(ensemble: str, valid_models, valid_trainers, valid_plans):
    ensemble = ensemble[len("ensemble_"):]
    mb1, mb2 = ensemble.split("--")
    c1, tr1, p1 = mb1.split("__")
    c2, tr2, p2 = mb2.split("__")
    if c1 not in valid_models: return False
    if c2 not in valid_models: return False
    if tr1 not in valid_trainers: return False
    if tr2 not in valid_trainers: return False
    if p1 not in valid_plans: return False
    if p2 not in valid_plans: return False
    return True


def copy_ensembles(taskname, output_folder, valid_models=('2d', '3d_fullres', '3d_lowres', '3d_cascade_fullres'),
                   valid_trainers=(default_trainer, default_cascade_trainer),
                   valid_plans=(default_plans_identifier,)):
    ensemble_dir = join(network_training_output_dir, 'ensembles', taskname)
    if not isdir(ensemble_dir):
        print("No ensemble directory found for task", taskname)
        return
    subd = subdirs(ensemble_dir, join=False)
    valid = []
    for s in subd:
        v = check_if_valid(s, valid_models, valid_trainers, valid_plans)
        if v:
            valid.append(s)
    output_ensemble = join(output_folder, 'ensembles', taskname)
    maybe_mkdir_p(output_ensemble)
    for v in valid:
        this_output = join(output_ensemble, v)
        maybe_mkdir_p(this_output)
        shutil.copy(join(ensemble_dir, v, 'postprocessing.json'), this_output)


def compress_everything(output_base, num_processes=8):
    p = Pool(num_processes)
    tasks = subfolders(output_base, join=False)
    tasknames = [i.split('/')[-1] for i in tasks]
    args = []
    for t, tn in zip(tasks, tasknames):
        args.append((join(output_base, tn + ".zip"), join(output_base, t)))
    p.starmap(compress_folder, args)
    p.close()
    p.join()


def compress_folder(zip_file, folder):
    """inspired by https://stackoverflow.com/questions/1855095/how-to-create-a-zip-archive-of-a-directory-in-python"""
    zipf = zipfile.ZipFile(zip_file, 'w', zipfile.ZIP_DEFLATED)
    for root, dirs, files in os.walk(folder):
        for file in files:
            zipf.write(join(root, file), os.path.relpath(join(root, file), folder))


def export_one_task(taskname, models, output_folder, nnunet_trainer=default_trainer,
                    nnunet_trainer_cascade=default_cascade_trainer,
                    plans_identifier=default_plans_identifier):
    copy_pretrained_models_for_task(taskname, output_folder, models, nnunet_trainer, nnunet_trainer_cascade,
                                    plans_identifier)
    copy_ensembles(taskname, output_folder, models, (nnunet_trainer, nnunet_trainer_cascade), (plans_identifier,))
    compress_folder(join(output_folder, taskname + '.zip'), join(output_folder, taskname))


def export_pretrained_model(task_name: str, output_file: str,
                            models: tuple = ("2d", "3d_lowres", "3d_fullres", "3d_cascade_fullres"),
                            nnunet_trainer=default_trainer,
                            nnunet_trainer_cascade=default_cascade_trainer,
                            plans_identifier=default_plans_identifier,
                            folds=(0, 1, 2, 3, 4), strict=True):
    zipf = zipfile.ZipFile(output_file, 'w', zipfile.ZIP_DEFLATED)
    trainer_output_dir = nnunet_trainer + "__" + plans_identifier
    trainer_output_dir_cascade = nnunet_trainer_cascade + "__" + plans_identifier

    for m in models:
        to = trainer_output_dir_cascade if m == "3d_cascade_fullres" else trainer_output_dir
        expected_output_folder = join(network_training_output_dir, m, task_name, to)
        if not isdir(expected_output_folder):
            if strict:
                raise RuntimeError("Task %s is missing the model %s" % (task_name, m))
            else:
                continue

        expected_folders = ["fold_%d" % i if i != 'all' else i for i in folds]
        assert all([isdir(join(expected_output_folder, i)) for i in expected_folders]), "not all requested folds " \
                                                                                        "present, " \
                                                                                        "Task %s model %s" % \
                                                                                        (task_name, m)

        assert isfile(join(expected_output_folder, "plans.pkl")), "plans.pkl missing, Task %s model %s" % (task_name, m)

        for e in expected_folders:
            zipf.write(join(expected_output_folder, e, "debug.json"),
                       os.path.relpath(join(expected_output_folder, e, "debug.json"),
                                       network_training_output_dir))
            zipf.write(join(expected_output_folder, e, "model_final_checkpoint.model"),
                       os.path.relpath(join(expected_output_folder, e, "model_final_checkpoint.model"),
                                       network_training_output_dir))
            zipf.write(join(expected_output_folder, e, "model_final_checkpoint.model.pkl"),
                       os.path.relpath(join(expected_output_folder, e, "model_final_checkpoint.model.pkl"),
                                       network_training_output_dir))
            zipf.write(join(expected_output_folder, e, "progress.png"),
                       os.path.relpath(join(expected_output_folder, e, "progress.png"), network_training_output_dir))
            if isfile(join(expected_output_folder, e, "network_architecture.pdf")):
                zipf.write(join(expected_output_folder, e, "network_architecture.pdf"),
                           os.path.relpath(join(expected_output_folder, e, "network_architecture.pdf"),
                                           network_training_output_dir))

        zipf.write(join(expected_output_folder, "plans.pkl"),
                   os.path.relpath(join(expected_output_folder, "plans.pkl"), network_training_output_dir))
        if not isfile(join(expected_output_folder, "postprocessing.json")):
            if strict:
                raise RuntimeError('postprocessing.json missing. Run nnUNet_determine_postprocessing or disable strict')
            else:
                print('WARNING: postprocessing.json missing')
        else:
            zipf.write(join(expected_output_folder, "postprocessing.json"),
                       os.path.relpath(join(expected_output_folder, "postprocessing.json"), network_training_output_dir))

    ensemble_dir = join(network_training_output_dir, 'ensembles', task_name)
    if not isdir(ensemble_dir):
        print("No ensemble directory found for task", task_name)
        return
    subd = subdirs(ensemble_dir, join=False)
    valid = []
    for s in subd:
        v = check_if_valid(s, models, (nnunet_trainer, nnunet_trainer_cascade), (plans_identifier))
        if v:
            valid.append(s)
    for v in valid:
        zipf.write(join(ensemble_dir, v, 'postprocessing.json'),
                   os.path.relpath(join(ensemble_dir, v, 'postprocessing.json'),
                                   network_training_output_dir))
    zipf.close()


def export_entry_point():
    import argparse
    parser = argparse.ArgumentParser(description="Use this script to export models to a zip file for sharing with "
                                                 "others. You can upload the zip file and then either share the url "
                                                 "for usage with nnUNet_download_pretrained_model_by_url, or share the "
                                                 "zip for usage with nnUNet_install_pretrained_model_from_zip")
    parser.add_argument('-t', type=str, help='task name or task id')
    parser.add_argument('-o', type=str, help='output file name. Should end with .zip')
    parser.add_argument('-m', nargs='+',
                        help='list of model configurations. Default: 2d 3d_lowres 3d_fullres 3d_cascade_fullres. Must '
                             'be adapted to fit the available models of a task',
                        default=("2d", "3d_lowres", "3d_fullres", "3d_cascade_fullres"), required=False)
    parser.add_argument('-tr', type=str, help='trainer class used for 2d 3d_lowres and 3d_fullres. '
                                              'Default: %s' % default_trainer, required=False, default=default_trainer)
    parser.add_argument('-trc', type=str, help='trainer class used for 3d_cascade_fullres. '
                                              'Default: %s' % default_cascade_trainer, required=False,
                        default=default_cascade_trainer)
    parser.add_argument('-pl', type=str, help='nnunet plans identifier. Default: %s' % default_plans_identifier,
                        required=False, default=default_plans_identifier)
    parser.add_argument('--disable_strict', action='store_true', help='set this if you want to allow skipping '
                                                                     'missing things', required=False)
    parser.add_argument('-f', nargs='+', help='Folds. Default: 0 1 2 3 4', required=False, default=[0, 1, 2, 3, 4])
    args = parser.parse_args()

    folds = args.f
    folds = [int(i) if i != 'all' else i for i in folds]

    taskname = args.t
    if taskname.startswith("Task"):
        pass
    else:
        try:
            taskid = int(taskname)
        except Exception as e:
            print('-t must be either a Task name (TaskXXX_YYY) or a task id (integer)')
            raise e
        taskname = convert_id_to_task_name(taskid)

    export_pretrained_model(taskname, args.o, args.m, args.tr, args.trc, args.pl, strict=not args.disable_strict,
                            folds=folds)


def export_for_paper():
    output_base = "/media/fabian/DeepLearningData/nnunet_trained_models"
    task_ids = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 17, 24, 27, 29, 35, 48, 55, 61, 38]
    for t in task_ids:
        if t == 61:
            models = ("3d_fullres",)
        else:
            models = ("2d", "3d_lowres", "3d_fullres", "3d_cascade_fullres")
        taskname = convert_id_to_task_name(t)
        print(taskname)
        output_folder = join(output_base, taskname)
        maybe_mkdir_p(output_folder)
        copy_pretrained_models_for_task(taskname, output_folder, models)
        copy_ensembles(taskname, output_folder)
    compress_everything(output_base, 8)