Copyright (c) MONAI Consortium

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.

MONet Bundle Training and Inference Tutorial

This notebook demonstrates how to create a MONAI Bundle for a trained nnUNet and use it for inference. This is needed when some other application from the MONAI EcoSystem require a MONAI Bundle (MONAI Label, MonaiAlgo for Federated Learning, etc).

This notebook cover the steps to convert a trained nnUNet model to a consumable MONAI Bundle. The nnUNet training is here perfomed using the nnUNetV2Runner.

Optionally, the notebook also demonstrates how to use the same MONet Bundle for training a new model. This might be needed in some applications where the nnUNet training needs to be performed through a MONAI Bundle (i.e., Active Learning in MONAI Label, MonaiAlgo for Federated Learning, etc).

Setup environment

!python -c "import monai" || pip install -q "monai-weekly[pillow, tqdm]"
!python -c "import matplotlib" || pip install -q matplotlib
!python -c "import nnunetv2" || pip install -q nnunetv2

Setup imports

from monai.config import print_config
from monai.apps import DecathlonDataset
import os
import tempfile
from monai.bundle.config_parser import ConfigParser
from monai.apps.nnunet import nnUNetV2Runner
import random
from monai.apps.nnunet.nnunet_bundle import convert_nnunet_to_monai_bundle
import json
from pathlib import Path

print_config()

Setup data directory

You can specify a directory with the MONAI_DATA_DIRECTORY environment variable.

This allows you to save results and reuse downloads.

If not specified a temporary directory will be used.

os.environ["MONAI_DATA_DIRECTORY"] = "MONAI/Data"

directory = os.environ.get("MONAI_DATA_DIRECTORY")
if directory is not None:
    os.makedirs(directory, exist_ok=True)
root_dir = tempfile.mkdtemp() if directory is None else directory
print(root_dir)

Download Decathlon Spleen Dataset and Generate Data List

To get the Decathlon Spleen dataset and generate the corresponding data list, you can follow the instructions in the MSD Datalist Generator Notebook

DecathlonDataset(root_dir, "Task09_Spleen", "training", download=True)

dataroot = os.path.join(root_dir, "Task09_Spleen/")

test_dir = os.path.join(dataroot, "imagesTs/")
train_dir = os.path.join(dataroot, "imagesTr/")
label_dir = os.path.join(dataroot, "labelsTr/")

datalist_json = {"testing": [], "training": []}

datalist_json["testing"] = [
    {"image": "./imagesTs/" + file} for file in os.listdir(test_dir) if (".nii.gz" in file) and ("._" not in file)
]

datalist_json["training"] = [
    {"image": "./imagesTr/" + file, "label": "./labelsTr/" + file, "fold": 0}
    for file in os.listdir(train_dir)
    if (".nii.gz" in file) and ("._" not in file)
]  # Initialize as single fold

random.seed(42)
random.shuffle(datalist_json["training"])

num_folds = 5
fold_size = len(datalist_json["training"]) // num_folds
for i in range(num_folds):
    for j in range(fold_size):
        datalist_json["training"][i * fold_size + j]["fold"] = i

datalist_file = Path(root_dir).joinpath("Task09_Spleen", "Task09_Spleen_folds.json")
with open(datalist_file, "w", encoding="utf-8") as f:
    json.dump(datalist_json, f, ensure_ascii=False, indent=4)
print(f"Datalist is saved to {datalist_file}")

nnUNet Experiment with nnUNetV2Runner

In the following sections, we will use the nnUNetV2Runner to train a model on the spleen dataset from the Medical Segmentation Decathlon.

We first create the Config file for the nnUNetV2Runner:

nnunet_root_dir = os.path.join(root_dir, "nnUNet")

os.makedirs(nnunet_root_dir, exist_ok=True)

data_src_cfg = os.path.join(nnunet_root_dir, "data_src_cfg.yaml")
data_src = {
    "modality": "CT",
    "dataset_name_or_id": "09",
    "datalist": str(datalist_file),
    "dataroot": os.path.join(root_dir, "Task09_Spleen"),
}

ConfigParser.export_config_file(data_src, data_src_cfg)

runner = nnUNetV2Runner(
    input_config=data_src_cfg, trainer_class_name="nnUNetTrainer_10epochs", work_dir=nnunet_root_dir
)

runner.convert_dataset()

runner.plan_and_process(npfp=2, n_proc=[2, 2, 2])

runner.train_single_model(config="3d_fullres", fold=0)

MONet Bundle for Inference

This section is the relevant part of the MONet Bundle for Inference, showing how to use the trained model to perform inference on new data through the use of a MONAI Bundle, wrapping the native nnUNet model and its pre- and post-processing steps.

We first download the MONet Bundle:

wget https://raw.githubusercontent.com/SimoneBendazzoli93/MONet-Bundle/main/MONetBundle.zip
unzip MONetBundle.zip

nnUnet to MONAI Bundle Conversion

Finally, we convert the nnUNet Trained Model to a Bundle-compatible format using the convert_nnunet_to_monai_bundle function:

nnunet_config = {
    "dataset_name_or_id": "009",
    "nnunet_trainer": "nnUNetTrainer_10epochs",
}

bundle_root = "MONetBundle"

convert_nnunet_to_monai_bundle(nnunet_config, bundle_root, 0)

You can then inspect the content of the models folder to verify that the model has been converted to the MONAI Bundle format.

%%bash

which tree && tree MONetBundle/models

Test the MONAI Bundle for Inference

The MONAI Bundle for Inference is now ready to be used for inference on new data

%%bash


BUNDLE_ROOT=MONetBundle
MONAI_DATA_DIRECTORY=MONAI/Data

python -m monai.bundle run \
    --config-file $BUNDLE_ROOT/configs/inference.yaml \
    --bundle-root $BUNDLE_ROOT \
    --data_list_file  $MONAI_DATA_DIRECTORY/Task09_Spleen/Task09_Spleen_folds.json \
    --output-dir $BUNDLE_ROOT/pred_output \
    --data_dir $MONAI_DATA_DIRECTORY/Task09_Spleen \
    --logging-file$BUNDLE_ROOT/configs/logging.conf