{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright (c) MONAI Consortium  \n",
    "Licensed under the Apache License, Version 2.0 (the \"License\");  \n",
    "you may not use this file except in compliance with the License.  \n",
    "You may obtain a copy of the License at  \n",
    "    http://www.apache.org/licenses/LICENSE-2.0  \n",
    "Unless required by applicable law or agreed to in writing, software  \n",
    "distributed under the License is distributed on an \"AS IS\" BASIS,  \n",
    "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  \n",
    "See the License for the specific language governing permissions and  \n",
    "limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MONet Bundle Training and Inference Tutorial\n",
    "\n",
    "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).\n",
    "\n",
    "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`.\n",
    "\n",
    "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)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "source": [
    "## Setup environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -c \"import monai\" || pip install -q \"monai-weekly[pillow, tqdm]\"\n",
    "!python -c \"import matplotlib\" || pip install -q matplotlib\n",
    "!python -c \"import nnunetv2\" || pip install -q nnunetv2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from monai.config import print_config\n",
    "from monai.apps import DecathlonDataset\n",
    "import os\n",
    "import tempfile\n",
    "from monai.bundle.config_parser import ConfigParser\n",
    "from monai.apps.nnunet import nnUNetV2Runner\n",
    "import random\n",
    "from monai.apps.nnunet.nnunet_bundle import convert_nnunet_to_monai_bundle\n",
    "import json\n",
    "from pathlib import Path\n",
    "\n",
    "print_config()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup data directory\n",
    "\n",
    "You can specify a directory with the `MONAI_DATA_DIRECTORY` environment variable.  \n",
    "This allows you to save results and reuse downloads.  \n",
    "If not specified a temporary directory will be used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"MONAI_DATA_DIRECTORY\"] = \"MONAI/Data\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "directory = os.environ.get(\"MONAI_DATA_DIRECTORY\")\n",
    "if directory is not None:\n",
    "    os.makedirs(directory, exist_ok=True)\n",
    "root_dir = tempfile.mkdtemp() if directory is None else directory\n",
    "print(root_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download Decathlon Spleen Dataset and Generate Data List\n",
    "\n",
    "To get the Decathlon Spleen dataset and generate the corresponding data list, you can follow the instructions in the [MSD Datalist Generator Notebook](https://github.com/Project-MONAI/tutorials/blob/main/auto3dseg/notebooks/msd_datalist_generator.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DecathlonDataset(root_dir, \"Task09_Spleen\", \"training\", download=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataroot = os.path.join(root_dir, \"Task09_Spleen/\")\n",
    "\n",
    "test_dir = os.path.join(dataroot, \"imagesTs/\")\n",
    "train_dir = os.path.join(dataroot, \"imagesTr/\")\n",
    "label_dir = os.path.join(dataroot, \"labelsTr/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datalist_json = {\"testing\": [], \"training\": []}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datalist_json[\"testing\"] = [\n",
    "    {\"image\": \"./imagesTs/\" + file} for file in os.listdir(test_dir) if (\".nii.gz\" in file) and (\"._\" not in file)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datalist_json[\"training\"] = [\n",
    "    {\"image\": \"./imagesTr/\" + file, \"label\": \"./labelsTr/\" + file, \"fold\": 0}\n",
    "    for file in os.listdir(train_dir)\n",
    "    if (\".nii.gz\" in file) and (\"._\" not in file)\n",
    "]  # Initialize as single fold"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "random.seed(42)\n",
    "random.shuffle(datalist_json[\"training\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_folds = 5\n",
    "fold_size = len(datalist_json[\"training\"]) // num_folds\n",
    "for i in range(num_folds):\n",
    "    for j in range(fold_size):\n",
    "        datalist_json[\"training\"][i * fold_size + j][\"fold\"] = i"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datalist_file = Path(root_dir).joinpath(\"Task09_Spleen\", \"Task09_Spleen_folds.json\")\n",
    "with open(datalist_file, \"w\", encoding=\"utf-8\") as f:\n",
    "    json.dump(datalist_json, f, ensure_ascii=False, indent=4)\n",
    "print(f\"Datalist is saved to {datalist_file}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## nnUNet Experiment with nnUNetV2Runner\n",
    "\n",
    "In the following sections, we will use the nnUNetV2Runner to train a model on the spleen dataset from the Medical Segmentation Decathlon."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We first create the Config file for the nnUNetV2Runner:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nnunet_root_dir = os.path.join(root_dir, \"nnUNet\")\n",
    "\n",
    "os.makedirs(nnunet_root_dir, exist_ok=True)\n",
    "\n",
    "data_src_cfg = os.path.join(nnunet_root_dir, \"data_src_cfg.yaml\")\n",
    "data_src = {\n",
    "    \"modality\": \"CT\",\n",
    "    \"dataset_name_or_id\": \"09\",\n",
    "    \"datalist\": str(datalist_file),\n",
    "    \"dataroot\": os.path.join(root_dir, \"Task09_Spleen\"),\n",
    "}\n",
    "\n",
    "ConfigParser.export_config_file(data_src, data_src_cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "runner = nnUNetV2Runner(\n",
    "    input_config=data_src_cfg, trainer_class_name=\"nnUNetTrainer_10epochs\", work_dir=nnunet_root_dir\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "runner.convert_dataset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "runner.plan_and_process(npfp=2, n_proc=[2, 2, 2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "runner.train_single_model(config=\"3d_fullres\", fold=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## MONet Bundle for Inference\n",
    "\n",
    "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."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We first download the MONet Bundle:\n",
    "\n",
    "```bash\n",
    "wget https://raw.githubusercontent.com/SimoneBendazzoli93/MONet-Bundle/main/MONetBundle.zip\n",
    "unzip MONetBundle.zip\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### nnUnet to MONAI Bundle Conversion\n",
    "\n",
    "Finally, we convert the nnUNet Trained Model to a Bundle-compatible format using the `convert_nnunet_to_monai_bundle` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nnunet_config = {\n",
    "    \"dataset_name_or_id\": \"009\",\n",
    "    \"nnunet_trainer\": \"nnUNetTrainer_10epochs\",\n",
    "}\n",
    "\n",
    "bundle_root = \"MONetBundle\"\n",
    "\n",
    "convert_nnunet_to_monai_bundle(nnunet_config, bundle_root, 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can then inspect the content of the `models` folder to verify that the model has been converted to the MONAI Bundle format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "\n",
    "which tree && tree MONetBundle/models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test the MONAI Bundle for Inference\n",
    "\n",
    "The MONAI Bundle for Inference is now ready to be used for inference on new data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "\n",
    "\n",
    "BUNDLE_ROOT=MONetBundle\n",
    "MONAI_DATA_DIRECTORY=MONAI/Data\n",
    "\n",
    "python -m monai.bundle run \\\n",
    "    --config-file $BUNDLE_ROOT/configs/inference.yaml \\\n",
    "    --bundle-root $BUNDLE_ROOT \\\n",
    "    --data_list_file  $MONAI_DATA_DIRECTORY/Task09_Spleen/Task09_Spleen_folds.json \\\n",
    "    --output-dir $BUNDLE_ROOT/pred_output \\\n",
    "    --data_dir $MONAI_DATA_DIRECTORY/Task09_Spleen \\\n",
    "    --logging-file$BUNDLE_ROOT/configs/logging.conf"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "MONAI",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}