MAIA#
This tutorial has the purpose to showcase all the major applications and functionalities available in MAIA. In detail, we will cover all the essential Medical AI lifecycle stages, to provide a comprehensive overview of the platform.
The tutorial is based on two different datasets: - Decathlon Spleen: a dataset of 3D spleen CT scans from the Medical Segmentation Decathlon challenge. This NIFTI dataset is used to demonstrate the model preprocessing, training and evaluation functionalities in MAIA. - CT Lymph Nodes from TCIA: a dataset of 3D lymph node CT scans from The Cancer Imaging Archive. This DICOM dataset is used to demonstrate the data management functionalities in MAIA, including DICOM upload,visualization, annotation and AI model inference (including Active Learning with MONAI Label).
The tutorial will cover all the necessary steps to download the Decathlon Spleen dataset, preprocess it, train a ResEnc nnU-Net model as a MONAI Bundle, evaluate the model, and finally deploy it for inference on the CT Lymph Nodes dataset. The tutorial will also cover the necessary steps to upload the CT Lymph Nodes dataset, visualize it, annotate it, and perform AI model inference on it.
In this tutorial, we will use PyMAIA, a Python package developed to handle nnUNet-based experiments. PyMAIA provides a set of functionalities to prepare the data, run the training, and evaluate the model. PyMAIA also provides a set of functionalities to convert the MONAI Bundle to a nnUNet Bundle, which can be used to deploy the model for inference. To read more about PyMAIA, please refer to the documentation.
Prerequisites#
To run this notebook, you need to have the following prerequisites:
Python Environment: Ensure you have a Python environment set up. This notebook is designed to work with Python 3.8 or later.
Required Packages: Install the necessary Python packages. You can install them using the following command:
bash !{sys.executable} -m pip install odict plotly dtale "monai[nibabel, skimage, scipy, pillow, tensorboard, gdown, ignite, torchvision, itk, tqdm, pandas, mlflow, matplotlib, pydicom]" "pydicom==2.4.4" nnunetv2==2.5.1 dtale pymaia-learn fireDatasets: Download the required datasets:
3D Slicer: Ensure you have 3D Slicer installed for visualization purposes. You can download it from here.
Remote Desktop Access: Access to the remote desktop environment for interacting with 3D Slicer.
Make sure to follow the instructions in the notebook to set up the environment and download the datasets.
[ ]:
import pip
import sys
!{sys.executable} -m pip install odict plotly nbformat==4.2.0 dtale "monai[nibabel, skimage, scipy, pillow, tensorboard, gdown, ignite, torchvision, itk, tqdm, pandas, mlflow, matplotlib, pydicom]" "pydicom==2.4.4" nnunetv2==2.5.1 dtale pymaia-learn fire
[ ]:
from monai.apps import DecathlonDataset
import subprocess
import json
import os
import pandas as pd
import pathlib
import mlflow
from pathlib import Path
import numpy as np
import yaml
import dtale
import dtale.app as dtale_app
from monai.bundle import ConfigParser
from mlflow.models import ModelSignature
from mlflow.types.schema import Schema, TensorSpec
Download the Decathlon Spleen dataset#
[ ]:
DecathlonDataset("MAIA","Task09_Spleen","training",download=True)
Visualize Spleen CT and SEG in 3D Slicer#
You can navigate to Remote Desktop (password: maia) to interact with the Remote Desktop environment and 3D Slicer.
[ ]:
slicer_executable = "/home/maia-user/Documents/Slicer-5.7.0-2024-06-08-linux-amd64/Slicer"
ct_volume = "/home/maia-user/Tutorials/MAIA/Task09_Spleen/imagesTr/spleen_2.nii.gz"
seg_volume = "/home/maia-user/Tutorials/MAIA/Task09_Spleen/labelsTr/spleen_2.nii.gz"
subprocess.run([
slicer_executable,
"--python-code", f"slicer.util.loadVolume('{ct_volume}'); seg=slicer.util.loadSegmentation('{seg_volume}'); seg.CreateClosedSurfaceRepresentation()"
])
Prepare nnUNet ResEnc Training#
Following the PyMAIA workflow steps for nnUnet experiments, we will prepare the data folder and run the plan and preprocessing steps.
To start with the data preparation, we need to create a configuration file that contains the experiment name, seed, label suffix, modalities, label dictionary, number of folds, and file extension. The configuration file is saved as Task09_Spleen_config.json. We also need to create a dataset file that contains the training and test data paths. The dataset file is saved as Task09_Spleen_Dataset.json.
[ ]:
modality_conf = {
"CT": {
"suffix": ".nii.gz"
}
}
config_dict = {
"Experiment Name": "Task09_Spleen",
"Seed": 12345,
"label_suffix":".nii.gz",
"Modalities": {modality_conf[key]["suffix"]:key for key in modality_conf },
"label_dict":
{
"background": 0,
"Spleen": 1
},
"n_folds":5,
"FileExtension": ".nii.gz"
}
with open("MAIA/Task09_Spleen_config.json","w") as f:
json.dump(config_dict,f)
[ ]:
data_dir = "/home/maia-user/Tutorials/MAIA/Task09_Spleen"
data_list = []
image_suffix_list = [modality_conf[image_suffix]['suffix'] for image_suffix in modality_conf ]
subjects = []
for f in os.scandir(pathlib.Path(data_dir).joinpath("imagesTr")):
if f.is_file():
for image_suffix in image_suffix_list:
if f.name.endswith(image_suffix) and not f.name.startswith("."):
subjects.append(f.name[:-len(image_suffix)])
subjects = np.unique(subjects)
for subject in subjects:
data = {}
for modality in modality_conf:
data[modality] = str(pathlib.Path(data_dir).joinpath("imagesTr",subject+modality_conf[modality]['suffix']))
data["label"] = str(pathlib.Path(data_dir).joinpath("labelsTr",subject+".nii.gz"))
data_list.append(data)
data ={"train":data_list,"test":[]}
with open("MAIA/Task09_Spleen_Dataset.json","w") as f:
json.dump(data,f)
Visualize Sample Case from Dataset#
To verify that the data location and format are correct, we can visualize a sample case from the dataset in 3D Slicer:
[ ]:
sample_case = data["train"][0]
image_volume = sample_case["CT"]
label_volume = sample_case["label"]
slicer_executable = "/home/maia-user/Documents/Slicer-5.7.0-2024-06-08-linux-amd64/Slicer"
subprocess.run([
slicer_executable,
"--python-code", f"slicer.util.loadVolume('{image_volume}'); seg=slicer.util.loadSegmentation('{label_volume}'); seg.CreateClosedSurfaceRepresentation()"
])
Prepare nnUNet Experiment#
We will use the `nnunet_prepare_data_folder <https://pymaia.readthedocs.io/en/latest/apidocs/nnunet_prepare_data_folder.html>`__ command to prepare the data folder for nnUNet training. The command requires the dataset file and the configuration file as input. The command will create the necessary folders and files for the nnUNet training, including the dataset.json file, the plans folder, and the preprocessing folder.
[ ]:
%%bash
export ROOT_FOLDER=MAIA/Experiments
export PATH=/home/maia-user/.conda/envs/MAIA/bin:$PATH
nnunet_prepare_data_folder \
-i MAIA/Task09_Spleen_Dataset.json \
--task-name Task09_Spleen \
--task-ID 109 \
--test-split 0 \
--config-file MAIA/Task09_Spleen_config.json
nnUNet Plan and Preprocessing#
After completing the Data Preparation step, we will run the `nnunet_run_plan_and_preprocessing <https://pymaia.readthedocs.io/en/latest/apidocs/nnunet_run_plan_and_preprocessing.html>`__ command to generate the plans and preprocess the data. The command requires the configuration file as input. The command will generate the plans and preprocess the data according to the configuration file.
[ ]:
%%bash
export ROOT_FOLDER=MAIA/Experiments
export PATH=/home/maia-user/.conda/envs/MAIA/bin:$PATH
nnunet_run_plan_and_preprocessing \
--config-file MAIA/Experiments/Task09_Spleen/Task09_Spleen_results/Dataset109_Task09_Spleen.json \
--n-workers 2 \
-pl nnUNetPlannerResEncL
nnUNet Training#
We will train a nnUNet model using a MONAI Bundle specifically designed for nnUNet. To generate the MONAI Bundle configuration files, run the following cells:
[ ]:
def create_config(config_folder, output_file):
config_files = [f.path for f in os.scandir(config_folder) if f.path.endswith(".yaml")]
config = {}
for config_file in config_files:
with open(config_file, 'r') as file:
config.update(yaml.safe_load(file))
if output_file.endswith(".yaml"):
with open(output_file, 'w') as file:
yaml.dump(config, file)
if output_file.endswith(".json"):
with open(output_file, 'w') as file:
json.dump(config, file)
return config
[ ]:
train_config = create_config("nnUNetBundle/nnUNet/","nnUNetBundle/configs/train.json")
train_config = create_config("nnUNetBundle/nnUNet/","nnUNetBundle/configs/train.yaml")
[ ]:
evaluate_config = create_config("nnUNetBundle/nnUNet/evaluator/","nnUNetBundle/configs/evaluate.json")
evaluate_config = create_config("nnUNetBundle/nnUNet/evaluator/","nnUNetBundle/configs/evaluate.yaml")
Additionally, since the original nnUNet Scheduler implementation is not compatible with a MONAI Bundle training, we will create a custom PolyLRScheduler class that can be used in the nnUNet training, overriding the original implementation.
The incompatibility is derived from the missing get_last_lr method in the original implementation, which is used to log the learning rate in the MONAI Bundle training.
[ ]:
%%writefile /home/maia-user/.conda/envs/MAIA/lib/python3.10/site-packages/nnunetv2/training/lr_scheduler/polylr.py
from torch.optim.lr_scheduler import _LRScheduler
class PolyLRScheduler(_LRScheduler):
def __init__(self, optimizer, initial_lr: float, max_steps: int, exponent: float = 0.9, current_step: int = None):
self.optimizer = optimizer
self.initial_lr = initial_lr
self.max_steps = max_steps
self.exponent = exponent
self.ctr = 0
super().__init__(optimizer, current_step if current_step is not None else -1, False)
def step(self, current_step=None):
if current_step is None or current_step == -1:
current_step = self.ctr
self.ctr += 1
new_lr = self.initial_lr * (1 - current_step / self.max_steps) ** self.exponent
for param_group in self.optimizer.param_groups:
param_group['lr'] = new_lr
self._last_lr = [group['lr'] for group in self.optimizer.param_groups]
def get_last_lr(self):
return self._last_lr
As an alternative to the MONAI Bundle training, we can run the nnUNet training using the `nnunet_run_training <https://pymaia.readthedocs.io/en/latest/apidocs/nnunet_run_training.html>`__ command. The command requires the configuration file as input. The command will train the model according to the configuration file.
[ ]:
%%bash
export ROOT_FOLDER=./Experiments
/home/maia-user/.conda/envs/MAIA/bin/python -m monai.bundle run \/envs/MAIA/bin/nnunet_run_training --config-file MAIA/Experiments/Task09_Spleen/Task09_Spleen_results/Dataset109_Task09_Spleen.json --n-workers 2 \
-p nnUNetResEncUNetLPlans
#-tr nnUNetTrainerDemo
The nnUNet MONAI Bundle is designed to log the training metrics to MLFLow. Additionally, the training metrics are logged to Tensorboard, which can be accessed by running the following command:
[ ]:
%%bash
tensorboard --logdir nnUNet_Bundle/eval
To start the MONAI Bundle nnUNet training:
[ ]:
%%bash
export BUNDLE=$HOME/Tutorials/nnUNetBundle
export PYTHONPATH=$BUNDLE
export nnUNet_def_n_proc=2
export nnUNet_n_proc_DA=2
/home/maia-user/.conda/envs/MAIA/bin/python -m monai.bundle run \
--bundle-root $BUNDLE \
--pymaia-config-file MAIA/Experiments/Task09_Spleen/Task09_Spleen_results/Dataset109_Task09_Spleen.json \
--task-id 109 \
--token "" \
--nnunet_plans_identifier nnUNetResEncUNetLPlans \
--mlflow_experiment_name "Task09_Spleen" \
--tracking_uri $MLFLOW_TRACKING_URI \
--mlflow_run_name "nnUNetResEncUNetL" \
--nnunet_model_folder $HOME/Tutorials/MAIA/Experiments/Task09_Spleen/Task09_Spleen_results/Dataset109_Task09_Spleen/nnUNetTrainer__nnUNetResEncUNetLPlans__3d_fullres \
--config-file $BUNDLE/configs/train.yaml
Validation#
After the training on fold 0 is completed, we can perform the corresponding validation using the `nnunet_run_training <https://pymaia.readthedocs.io/en/latest/apidocs/nnunet_run_training.html>`__ command. The command requires the configuration file as input. The command will validate the model according to the configuration file.
Since the validation is performed according to the native nnUNet scripts, we need to locally convert the trained model from the MONAI Bundle format back to the original nnUNet format. To do this, we will use the convert_MONAI_to_nnUNet function:
[ ]:
def convert_MONAI_to_nnUNet(nnunet_root_folder, nnunet_config, bundle_config):
from PyMAIA.utils.file_utils import subfiles
from nnunetv2.training.logging.nnunet_logger import nnUNetLogger
from pathlib import Path
import json
import torch
from odict import odict
import os
import shutil
network_weights_prefix = "_orig_mod."
os.environ["ROOT_FOLDER"] = nnunet_root_folder
os.environ["RESULTS_FOLDER"] = str(
Path(os.environ["ROOT_FOLDER"]).joinpath(
nnunet_config["Experiment Name"], nnunet_config["Experiment Name"] + "_results"
)
)
nnunet_trainer = "nnUNetTrainer"
nnunet_plans = "nnUNetPlans"
if "nnunet_trainer" in nnunet_config:
nnunet_trainer = nnunet_config["nnunet_trainer"]
if "nnunet_plans" in nnunet_config:
nnunet_plans = nnunet_config["nnunet_plans"]
nnunet_model_folder = Path(os.environ["RESULTS_FOLDER"]).joinpath(
"Dataset" + nnunet_config["task_ID"] + "_" + nnunet_config[
"Experiment Name"],
f"{nnunet_trainer}__{nnunet_plans}__3d_fullres")
bundle_name = bundle_config["Bundle_Name"]
nnunet_checkpoint = torch.load(f"{bundle_name}/models/nnunet_checkpoint.pth")
latest_checkpoints = subfiles(Path(bundle_name).joinpath("models"),prefix="checkpoint_epoch",sort=True,join=False)
epochs = []
for latest_checkpoint in latest_checkpoints:
epochs.append(int(latest_checkpoint[len("checkpoint_epoch="):-len(".pt")]))
epochs.sort()
final_epoch = epochs[-1]
monai_last_checkpoint = torch.load(f"{bundle_name}/models/checkpoint_epoch={final_epoch}.pt")
best_checkpoints = subfiles(Path(bundle_name).joinpath("models"), prefix="checkpoint_key_metric", sort=True,
join=False)
key_metrics = []
for best_checkpoint in best_checkpoints:
key_metrics.append(str(best_checkpoint[len("checkpoint_key_metric="):-len(".pt")]))
key_metrics.sort()
best_key_metric = key_metrics[-1]
monai_best_checkpoint = torch.load(f"{bundle_name}/models/checkpoint_key_metric={best_key_metric}.pt")
nnunet_checkpoint['optimizer_state'] = monai_last_checkpoint['optimizer_state']
nnunet_checkpoint['network_weights'] = odict()
for key in monai_last_checkpoint['network_weights']:
nnunet_checkpoint['network_weights'][key] = monai_last_checkpoint['network_weights'][key]
nnunet_checkpoint['current_epoch'] = final_epoch
nnunet_checkpoint['logging'] = nnUNetLogger().get_checkpoint()
nnunet_checkpoint['_best_ema'] = 0
nnunet_checkpoint['grad_scaler_state'] = None
torch.save(nnunet_checkpoint, Path(nnunet_model_folder).joinpath("fold_0","checkpoint_final.pth"))
nnunet_checkpoint['network_weights'] = odict()
nnunet_checkpoint['optimizer_state'] = monai_best_checkpoint['optimizer_state']
for key in monai_best_checkpoint['network_weights']:
nnunet_checkpoint['network_weights'][key] = \
monai_best_checkpoint['network_weights'][key]
torch.save(nnunet_checkpoint, Path(nnunet_model_folder).joinpath("fold_0", "checkpoint_best.pth"))
shutil.move(f"{bundle_name}/models/checkpoint_epoch={final_epoch}.pt",f"{bundle_name}/models/model.pt")
shutil.move(f"{bundle_name}/models/checkpoint_key_metric={best_key_metric}.pt",f"{bundle_name}/models/best_model.pt")
#shutil.copy(Path(nnunet_model_folder).joinpath("fold_0", "checkpoint_final.pth"), f"{bundle_name}/models/model.pt")
[ ]:
nnunet_root_folder = "MAIA/Experiments"
nnunet_config = {
"Experiment Name": "Task09_Spleen",
"task_ID": "109",
"nnunet_plans":"nnUNetResEncUNetLPlans"
}
bundle_config = {
"Bundle_Name": "nnUNetBundle"
}
[ ]:
convert_MONAI_to_nnUNet(nnunet_root_folder,nnunet_config, bundle_config)
After converting the MONAI Bundle to the nnUNet format, we can run the validation:
[ ]:
%%bash
export ROOT_FOLDER=MAIA/Experiments
/home/maia-user/.conda/envs/MAIA/bin/nnunet_run_training --config-file MAIA/Experiments/Task09_Spleen/Task09_Spleen_results/Dataset109_Task09_Spleen.json --n-workers 2 \
-p nnUNetResEncUNetLPlans --run-fold 0 --run-validation-only yes
To visualize in the 3D Slicer the fold-0 validation predictions, together with the ground truth and the CT volume, we can use the following code:
[ ]:
slicer_executable = "/home/maia-user/Documents/Slicer-5.7.0-2024-06-08-linux-amd64/Slicer"
ct_volume = "/home/maia-user/Tutorials/MAIA/Task09_Spleen/imagesTr/spleen_10.nii.gz"
seg_volume = "/home/maia-user/Tutorials/MAIA/Task09_Spleen/labelsTr/spleen_10.nii.gz"
pred_volume = "/home/maia-user/Tutorials/MAIA/Experiments/Task09_Spleen/Task09_Spleen_results/Dataset109_Task09_Spleen/nnUNetTrainer__nnUNetResEncUNetLPlans__3d_fullres/fold_0/validation/spleen_10.nii.gz"
subprocess.run([
slicer_executable,
"--python-code", f"slicer.util.loadVolume('{ct_volume}'); pred=slicer.util.loadSegmentation('{pred_volume}'); seg=slicer.util.loadSegmentation('{seg_volume}'); seg.CreateClosedSurfaceRepresentation()"
])
Validation Metrics#
The validation metrics are stored in a JSON file named summary.json in the validation folder. We can load the file and visualize the metrics with DTale using the following code:
[53]:
summary_file= "MAIA/Experiments/Task09_Spleen/Task09_Spleen_results/Dataset109_Task09_Spleen/nnUNetTrainer__nnUNetResEncUNetLPlans__3d_fullres/fold_0/validation/summary.json"
[54]:
with open(summary_file) as f:
summary = json.load(f)
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[54], line 2
1 with open(summary_file) as f:
----> 2 summary = json.load(f)
NameError: name 'json' is not defined
[ ]:
df = []
label_to_name = {v: k for k, v in config_dict["label_dict"].items()}
for case in summary['metric_per_case']:
for label_id in case['metrics']:
for metric in case['metrics'][label_id]:
df.append({
"Case": Path(case['reference_file']).name[:-len(config_dict["label_suffix"])],
"Label": label_to_name[int(label_id)],
"Metric": metric,
"Value": case['metrics'][label_id][metric]
})
[ ]:
df = pd.DataFrame(df)
[ ]:
dtale_app.JUPYTER_SERVER_PROXY = True
d = dtale.show(df,host="0.0.0.0",)
[ ]:
from IPython.display import Markdown
from IPython.core.magic import register_cell_magic
import os
DTALE_URL = d._main_url
@register_cell_magic
def markdown(line, cell):
return Markdown(cell.format(**globals()))
[ ]:
%%markdown
[DTale]({DTALE_URL})
The DTale charts can be recreated and visualized in the notebook using the following code:
[ ]:
# DISCLAIMER: 'df' refers to the data you passed in when calling 'dtale.show'
import pandas as pd
if isinstance(df, (pd.DatetimeIndex, pd.MultiIndex)):
df = df.to_frame(index=False)
# remove any pre-existing indices for ease of use in the D-Tale code, but this is not required
df = df.reset_index().drop('index', axis=1, errors='ignore')
df.columns = [str(c) for c in df.columns] # update columns to strings in case they are numbers
df = df.query("""`Metric` == 'Dice'""")
chart_data = pd.concat([
df['Case'],
df['Value'],
], axis=1)
chart_data = chart_data.sort_values(['Case'])
chart_data = chart_data.rename(columns={'Case': 'x'})
chart_data = chart_data.dropna()
import plotly.graph_objs as go
charts = []
charts.append(go.Bar(
x=chart_data['x'],
y=chart_data['Value']
))
figure = go.Figure(data=charts, layout=go.Layout({
'barmode': 'group',
'legend': {'orientation': 'h', 'y': -0.3},
'title': {'text': 'Validation Fold 0, Dice score'},
'xaxis': {'title': {'text': 'Case'}},
'yaxis': {'title': {'text': 'Dice'}, 'type': 'linear'}
}))
# If you're having trouble viewing your chart in your notebook try passing your 'chart' into this snippet:
#
from plotly.offline import iplot, init_notebook_mode
#
init_notebook_mode(connected=True)
for chart in charts:
chart.pop('id', None) # for some reason iplot does not like 'id'
iplot(figure)
figure.write_html("Fold_0_Val_Dice.html")
[ ]:
df.groupby(["Metric"]).describe()
[ ]:
df.groupby(["Metric"]).describe()['Value']['mean'].values[0]
Finally, we can upload the validation metrics and the plots to MLFlow using the following code:
[ ]:
mlflow.set_tracking_uri(os.environ["MLFLOW_TRACKING_URI"])
mlflow.set_experiment("Task09_Spleen")
with mlflow.start_run(run_id="58d802747d7d493092d21287cc16af2d"):
mean_dice = df.groupby(["Metric"]).describe()['Value']['mean'].values[0]
mlflow.log_metric("Val_Dice_Fold_0", mean_dice)
mlflow.log_artifact("Fold_0_Val_Dice.html")
In the final step of the validation phase, we export the trained model, saving the nnUNet Bundle as a zip file (Task09_Spleen_nnUNet.zip). The zip file contains the model, the configuration files, and the environment files.
[ ]:
%%bash
export ROOT_FOLDER=MAIA/Experiments
export PATH=/home/maia-user/.conda/envs/MAIA/bin:$PATH
touch MAIA/Experiments/Task09_Spleen/Task09_Spleen_results/Dataset109_Task09_Spleen/nnUNetTrainer__nnUNetResEncUNetLPlans__3d_fullres/fold_0/progress.png
/home/maia-user/.conda/envs/MAIA/bin/nnunet_run_training --config-file MAIA/Experiments/Task09_Spleen/Task09_Spleen_results/Dataset109_Task09_Spleen.json --n-workers 2 \
-p nnUNetResEncUNetLPlans --run-fold -1 --output-model-file Task09_Spleen_nnUNet.zip --post-processing-folds 0
Package MONAI Bundle#
After completing the training and validation of the nnUNet model, we can package the model as a MONAI Bundle, to be used for inference (i.e. Active Learning with MONAI Label). We first copy the trained model into the bundle folder and remove the validation results:
[ ]:
%%bash
#OPTIONAL STEP, NOT REQUIRED ANYMORE
unzip Task09_Spleen_nnUNet.zip -d nnUNetBundle/models
rm -r nnUNetBundle/models/Dataset109_Spleen/nnUNetTrainer__nnUNetResEncUNetLPlans__3d_fullres/fold_0/validation
rm -r nnUNetBundle/models/Dataset109_Spleen/nnUNetTrainer__nnUNetResEncUNetLPlans__3d_fullres/crossval_results_folds_0
We then test it by running the model on a sample case from the dataset:
[ ]:
%%bash
rm -r MAIA/MONAI_Bundle
mkdir -p MAIA/MONAI_Bundle/input
mkdir -p MAIA/MONAI_Bundle/output
mkdir -p MAIA/MONAI_Bundle/input/spleen_1
cp MAIA/Task09_Spleen/imagesTs/spleen_1.nii.gz MAIA/MONAI_Bundle/input/spleen_1
[ ]:
%%bash
/home/maia-user/.conda/envs/MAIA/bin/python -m monai.bundle run \
--config-file nnUNetBundle/configs/inference.yaml \
--bundle-root nnUNetBundle \
--data-dir MAIA/MONAI_Bundle/input \
--output-dir MAIA/MONAI_Bundle/output \
--logging-file nnUNetBundle/configs/logging.conf
To visualize the prediction in 3D Slicer:
[ ]:
slicer_executable = "/home/maia-user/Documents/Slicer-5.7.0-2024-06-08-linux-amd64/Slicer"
ct_volume = "/home/maia-user/Tutorials/MAIA/MONAI_Bundle/input/spleen_1/spleen_1.nii.gz"
pred_volume = "/home/maia-user/Tutorials/MAIA/MONAI_Bundle/output/spleen_1/spleen_1_prediction.nii.gz"
subprocess.run([
slicer_executable,
"--python-code", f"slicer.util.loadVolume('{ct_volume}'); pred=slicer.util.loadSegmentation('{pred_volume}'); pred.CreateClosedSurfaceRepresentation()"
])
We finally export the python environment and the requirements for the MONAI Bundle:
[ ]:
%%bash
/opt/conda/bin/conda env export -n MAIA > nnUNetBundle/environment.yml
/home/maia-user/.conda/envs/MAIA/bin/python -m pip freeze > nnUNetBundle/requirements.txt
We can then create the metadata file for the MONAI Bundle:
[ ]:
import monai
import torch
import nnunetv2
import PyMAIA
import pydicom
import numpy
import ignite
#import lightning
from importlib.metadata import version
MONAI_VERSION = monai.__version__
PYTORCH_VERSION = torch.__version__
NNUNET_VERSION = version("nnunetv2")
PYMAIA_VERSION = PyMAIA.__version__
PYDICOM_VERSION = pydicom.__version__
NUMPY_VERSION = numpy.__version__
IGNITE_VERSION = ignite.__version__
#LIGHTNING_VERSION = lightning.__version__
patch_size = net.predictor.configuration_manager.patch_size
[ ]:
#nnUNetBundle/configs/metadata.json
metadata = {
"version": "0.1.0",
"monai_version": MONAI_VERSION,
"pytorch_version": PYTORCH_VERSION,
"numpy_version": NUMPY_VERSION,
"optional_packages_version": {
"nnunetv2": NNUNET_VERSION,
"pytorch-ignite": IGNITE_VERSION,
"pymaia_learn": PYMAIA_VERSION,
"pydicom": PYDICOM_VERSION,
#"lightning": lightning.__version__
},
"task": "Spleen Segmentation with ResEnc nnUNet [L]",
"description": "A nnUNet MONAI Bundle for Spleen Segmentation in CT, used for nnUNet-based Experiments within the MONAI Bundle framework.",
"authors": "Simone Bendazzoli",
"copyright": "Copyright 2024, Apache Software Foundation",
"data_source": "Decathlon Dataset, Task 09",
"data_type": "nifti",
"image_classes": "1 single channel volumes: CT",
"label_classes": "single channel data: 0 is background, 1 is spleen",
"pred_classes": "single channel data: 0 is background, 1 is spleen",
"network_data_format":
{
"inputs": {
"image": {
"type": "image",
"format": "magnitude",
"modality": "CT",
"num_channels": 1,
"spatial_shape": patch_size,
"dtype": "float32",
"value_range": [0, 1],
"is_patch_data": True,
"channel_def": {"0": "ct"}
}
},
"outputs":{
"pred": {
"type": "image",
"format": "segmentation",
"num_channels": 1,
"spatial_shape": patch_size,
"dtype": "float32",
"value_range": [0,1],
"is_patch_data": True,
"channel_def": {"0": "Background", "1": "Spleen"}
}
}
},
"references": [
"Isensee, F., Jaeger, P. F., Kohl, S. A., Petersen, J., & Maier-Hein, K. H. (2021). nnU-Net: a self-configuring method for deep learning-based biomedical image segmentation. Nature methods, 18(2), 203-211."
]
}
with open("nnUNetBundle/configs/metadata.json","w") as f:
json.dump(metadata,f,indent=4)
[51]:
%%writefile nnUNetBundle/requirements.txt
nnunetv2==2.5.1
odict
pymaia-learn
fire
Overwriting nnUNetBundle/requirements.txt
Optionally, before packaging the MONAI Bundle, we need to set the model_folder parameter in the inference.yaml file.
This step is only needed if the native nnUNet Inference is used.
model_folder: <BUNDLE_ROOT>/models/Dataset109_Spleen/nnUNetTrainer__nnUNetResEncUNetLPlans__3d_fullres
[52]:
%%bash
rm -r Task09_Spleen_Bundle
rm -r Task09_Spleen_Bundle.zip
cp -r nnUNetBundle Task09_Spleen_Bundle
zip -r Task09_Spleen_Bundle.zip Task09_Spleen_Bundle
adding: Task09_Spleen_Bundle/ (stored 0%)
adding: Task09_Spleen_Bundle/src/ (stored 0%)
adding: Task09_Spleen_Bundle/src/inferer.py (deflated 74%)
adding: Task09_Spleen_Bundle/src/__init__.py (stored 0%)
adding: Task09_Spleen_Bundle/src/__pycache__/ (stored 0%)
adding: Task09_Spleen_Bundle/src/__pycache__/__init__.cpython-310.pyc (deflated 22%)
adding: Task09_Spleen_Bundle/src/__pycache__/nnUNet_Trainer.cpython-310.pyc (deflated 44%)
adding: Task09_Spleen_Bundle/src/__pycache__/utils.cpython-310.pyc (deflated 40%)
adding: Task09_Spleen_Bundle/src/__pycache__/dataset.cpython-310.pyc (deflated 31%)
adding: Task09_Spleen_Bundle/src/__pycache__/inferer.cpython-310.pyc (deflated 47%)
adding: Task09_Spleen_Bundle/src/nnUNet_Trainer.py (deflated 76%)
adding: Task09_Spleen_Bundle/src/utils.py (deflated 68%)
adding: Task09_Spleen_Bundle/src/dataset.py (deflated 43%)
adding: Task09_Spleen_Bundle/configs/ (stored 0%)
adding: Task09_Spleen_Bundle/configs/evaluate.yaml (deflated 62%)
adding: Task09_Spleen_Bundle/configs/train.yaml (deflated 74%)
adding: Task09_Spleen_Bundle/configs/.ipynb_checkpoints/ (stored 0%)
adding: Task09_Spleen_Bundle/configs/.ipynb_checkpoints/inference-checkpoint.json (deflated 71%)
adding: Task09_Spleen_Bundle/configs/train.json (deflated 74%)
adding: Task09_Spleen_Bundle/configs/inference.yaml (deflated 63%)
adding: Task09_Spleen_Bundle/configs/metadata.json (deflated 68%)
adding: Task09_Spleen_Bundle/configs/evaluate.json (deflated 62%)
adding: Task09_Spleen_Bundle/environment.yml (deflated 63%)
adding: Task09_Spleen_Bundle/.ipynb_checkpoints/ (stored 0%)
adding: Task09_Spleen_Bundle/.ipynb_checkpoints/environment-checkpoint.yml (deflated 63%)
adding: Task09_Spleen_Bundle/.ipynb_checkpoints/requirements-checkpoint.txt (deflated 56%)
adding: Task09_Spleen_Bundle/nnUNet/ (stored 0%)
adding: Task09_Spleen_Bundle/nnUNet/evaluator/ (stored 0%)
adding: Task09_Spleen_Bundle/nnUNet/evaluator/evaluator.yaml (deflated 61%)
adding: Task09_Spleen_Bundle/nnUNet/run.yaml (deflated 46%)
adding: Task09_Spleen_Bundle/nnUNet/train.yaml (deflated 70%)
adding: Task09_Spleen_Bundle/nnUNet/imports.yaml (deflated 46%)
adding: Task09_Spleen_Bundle/nnUNet/validate.yaml (deflated 63%)
adding: Task09_Spleen_Bundle/nnUNet/params.yaml (deflated 38%)
adding: Task09_Spleen_Bundle/nnUNet/global.yaml (deflated 42%)
adding: Task09_Spleen_Bundle/nnUNet/train_handlers.yaml (deflated 57%)
adding: Task09_Spleen_Bundle/models/ (stored 0%)
adding: Task09_Spleen_Bundle/models/dataset.json (deflated 91%)
adding: Task09_Spleen_Bundle/models/.ipynb_checkpoints/ (stored 0%)
adding: Task09_Spleen_Bundle/models/best_model.pt (deflated 7%)
adding: Task09_Spleen_Bundle/models/plans.json (deflated 89%)
adding: Task09_Spleen_Bundle/models/nnunet_checkpoint.pth (deflated 66%)
adding: Task09_Spleen_Bundle/models/model.pt (deflated 7%)
adding: Task09_Spleen_Bundle/LICENSE (stored 0%)
adding: Task09_Spleen_Bundle/eval/ (stored 0%)
adding: Task09_Spleen_Bundle/eval/events.out.tfevents.1732567232.jupyter-demo-40maia-2ese.17102.0 (deflated 70%)
adding: Task09_Spleen_Bundle/eval/events.out.tfevents.1732567232.jupyter-demo-40maia-2ese.17102.1 (deflated 66%)
adding: Task09_Spleen_Bundle/requirements.txt (stored 0%)
adding: Task09_Spleen_Bundle/docs/ (stored 0%)
adding: Task09_Spleen_Bundle/docs/README.md (deflated 49%)
MLFlow Model Upload#
To store the model and be able to deploy it for inference in future use cases, we can upload it to MLFlow. We will use the MLFlow Python API to log the model:
[ ]:
import sys
import os
import yaml
from monai.bundle import ConfigParser
import torch
import numpy as np
import mlflow
from mlflow.models import ModelSignature
from mlflow.types.schema import Schema, TensorSpec
sys.path.append("Task09_Spleen_Bundle")
[ ]:
config_files = [f.path for f in os.scandir("Task09_Spleen_Bundle/configs") if f.path.endswith("inference.yaml")]
config = {}
for config_file in config_files:
with open(config_file, 'r') as file:
config.update(yaml.safe_load(file))
config["bundle_root"] = "Task09_Spleen_Bundle"
#config["model_folder"] = "Dataset109_Spleen/nnUNetTrainer__nnUNetResEncUNetLPlans__3d_fullres"
parser = ConfigParser(config,globals={"os": "os",
"pathlib":"pathlib",
"json":"json",
"ignite":"ignite"
})
parser.parse(True)
[ ]:
net = parser.get_parsed_content("network_def",instantiate=True)
[ ]:
net.network_weights.load_state_dict(torch.load("nnUNet_Bundle/models/model.pt")['network_weights'])
[ ]:
os.environ["MLFLOW_TRACKING_URI"] = "http://mlflow-v1-mkg:5000"
[ ]:
mlflow.set_experiment("Task09_Spleen")
mlflow.end_run()
input_schema = Schema(
[
TensorSpec(np.dtype(np.float32), (1, *net.predictor.configuration_manager.patch_size),name="ct")
]
)
output_schema = Schema([TensorSpec(np.dtype(np.float32), (1, *net.predictor.configuration_manager.patch_size),name="Spleen")])
signature = ModelSignature(inputs=input_schema, outputs=output_schema)
with mlflow.start_run(run_id='58d802747d7d493092d21287cc16af2d'):
mlflow.pytorch.log_model(
net,
"Task09_Spleen",
signature=signature,
conda_env = "/home/maia-user/Tutorials/nnUNet_Bundle/environment.yml",
registered_model_name = "Task09_Spleen",
extra_files = [
"/home/maia-user/Tutorials/Task09_Spleen_Bundle.zip",
"/home/maia-user/Tutorials/nnUNet_Bundle/environment.yml",
"/home/maia-user/Tutorials/nnUNet_Bundle/requirements.txt"
]
)
Deploy Model for Inference in MAIA#
To deploy the Spleen segmentation model as a service in MAIA, run the following command to install it as an Helm chart:
[ ]:
from kfp import kubernetes
from kfp import client
from kfp import dsl
from kfp import compiler
from pathlib import Path
[ ]:
Path("/home/maia-user/shared").joinpath("mlflow-models","Task09_Spleen","extra_files").mkdir(parents=True, exist_ok=True)
[ ]:
import shutil
shutil.copy("/home/maia-user/Tutorials/Task09_Spleen_Bundle.zip",Path("/home/maia-user/shared").joinpath("mlflow-models","Task09_Spleen","extra_files"))
[ ]:
@dsl.component(base_image='kthcloud/maia-workspace-admin:1.5')
def helm_install_monai_label_ohif(cluster_api: str,namespace: str, user_id: str, id_token: str):
import subprocess
from pathlib import Path
import os
import yaml
def generate_kubeconfig(id_token,user_id,namespace,cluster_api):
kube_config = {'apiVersion': 'v1', 'kind': 'Config', 'preferences': {},
'current-context': 'MAIA/{}'.format(user_id), 'contexts': [
{'name': 'MAIA/{}'.format(user_id),
'context': {'user': user_id, 'cluster': 'MAIA', 'namespace': namespace}}],
'clusters': [
{'name': 'MAIA', 'cluster': {'certificate-authority-data': "",
'server': cluster_api,
"insecure-skip-tls-verify": True}}],
"users": [{'name': user_id,
'user': {'token': id_token}}]}
return kube_config
kubeconfig_dict = generate_kubeconfig(id_token,user_id,namespace,cluster_api)
with open(Path(".").joinpath("kubeconfig"), "w") as f:
yaml.dump(kubeconfig_dict, f)
os.environ["KUBECONFIG"] = "kubeconfig"
subprocess.run(["helm","repo","add","maia","https://kthcloud.github.io/MAIA/"])
subprocess.run(["helm","repo","update"])
subprocess.run(["helm", "install",
"spleen-segmentation",
"-n", namespace,
"maia/monai-label-ohif-maia",
"--set", "hostname=monai-demo.maia.cloud.cbh.kth.se",
"--set", "pvc.pvc_type=nfs",
"--set", "bundle_model_name=Task09_Spleen_Bundle",
"--set", "mlflow_pvc_name=shared",
"--set", "mlflow_model_path=/workspace/mlflow/mlflow-models/Task09_Spleen",
])
[ ]:
@dsl.pipeline
def install_monai_label_ohif_pipeline(cluster_api: str,namespace: str, user_id: str, id_token: str):
task1 = helm_install_monai_label_ohif(cluster_api=cluster_api,namespace=namespace,user_id=user_id,id_token=id_token)
kubernetes.set_image_pull_secrets(task1, secret_names=["harbor-maia-docker-registry-secret"])
[ ]:
compiler.Compiler().compile(install_monai_label_ohif_pipeline, package_path='Install_MONAI_Label_OHIF_pipeline.yaml')
[ ]:
token = "eyJhbGciOiJSUzI1NiIsInR5cCIgOiAiSldUIiwia2lkIiA6ICJ2T3FLaTZ2MGNHYXhxVzFOaFZiOFh5ekZqbDJERFRkUEhHTWgtTXVEbXNNIn0.eyJleHAiOjE3MzU2NDg0MTksImlhdCI6MTczMzA1NjQxOSwiYXV0aF90aW1lIjoxNzMzMDU2NDE5LCJqdGkiOiI1YWVkNDJhYi1mZTAxLTRlNGYtYTFmOS1hMWExNWYzYmFiODQiLCJpc3MiOiJodHRwczovL2lhbS5jbG91ZC5jYmgua3RoLnNlL3JlYWxtcy9jbG91ZCIsImF1ZCI6Im1haWEiLCJzdWIiOiI4YTdjYTk4Ny05M2E0LTRlOTgtODY3Ny00NWI5NmRiYzAzYzUiLCJ0eXAiOiJJRCIsImF6cCI6Im1haWEiLCJzaWQiOiIyMDhlODI4MC04MzI1LTQ5MjYtYTQ4MC01NzA0ZWJlY2FiMDMiLCJhdF9oYXNoIjoiWTdIM2NwY0FKckwtUG13Q3FpdVlxdyIsImFjciI6IjEiLCJlbWFpbF92ZXJpZmllZCI6dHJ1ZSwibmFtZSI6IkRlbW8gTUFJQSIsImdyb3VwcyI6WyJNQUlBOk1PTkFJX0RlbW9fMiIsIk1BSUE6bW9uYWktZGVtbyIsIk1BSUE6dXNlcnMiXSwicHJlZmVycmVkX3VzZXJuYW1lIjoiZGVtb0BtYWlhLnNlIiwiZ2l2ZW5fbmFtZSI6IkRlbW8iLCJmYW1pbHlfbmFtZSI6Ik1BSUEiLCJlbWFpbCI6ImRlbW9AbWFpYS5zZSJ9.WJck1-YyVf12eQFjCW2kSFMjVuQJsou_5XPI9luGD5CgTj2MvguYm9vlpZivaDN_abaRzKz-5301Y2TJhQaLt-O2Nfs3rFo0ddcKxnZu1uAZzBxh1KhdtjKjbNfMIPaWLz7D8m664C4tOgHv6jk-sTiqBwE0R3aTSuFRzd7FgMxdn00e23bzIUnJYqSgwXKRkTCUivw2ZopcMjBj6ONFb3l_kHsBA-xo26CLtUtHIUlqndJgfMoa05RXkUavbLzS09vD6vZxVET-ZV3-NX7mMRbm3d7CHtE0cLATxv4tj8oO3KiirKFmtdnLqoFNC5tsgenqkcUhKal31EMOMYVnGNL_Ijbza_Jp40xCm3o-s2vyV3B0W-SQV53cj3jxztMQa1IC3PAQ_UxRiKxIUnc_YWzdNvt0cWoLz33OozOJvvkIJq6EwfbeL1kxHlflzZmo6qnl3MayW9-FMAtne4iE5r_d0IMaawMr0YyE8QxcKhOxckde5f0nrNZtjCfeykGAZwt99l-DPZqTDFoiokDWr7qVYvcU9QHQmI73Sfb8tj78KkMHgSOHs1Bsv2I5htBfOd9GyNhjmPN_b-9yiy-TcUs-Sx1oYR5ur8MMVcJi-ZNxv9egAhrgtQYnSA_XtJ7n_8SukrpeQTSMbGlF4RRYDBS1_Mte9bWW6KJZY8CwTBY"
[ ]:
kfp_client = client.Client(host="http://ml-pipeline-ui")
kfp_client.create_run_from_pipeline_package('Install_MONAI_Label_OHIF_pipeline.yaml',
arguments={
"cluster_api":"https://kubernetes.default:443",
"namespace":"monai-demo",
"user_id":"monai-demo@kth.se",
"id_token":token
})
After deploy the Helm Chart, you can get access to MONAI Label and the linked Orthanc server.
To perform Active Learning with MONAI Label from the OHIF Viewer: OHIF Viewer