{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# nnU-Net for BraTS21\n",
    "\n",
    "# Table of contents\n",
    "- [Introduction](#introduction)\n",
    "- [Dataset](#dataset)\n",
    "- [Data pre-processing](#preprocessing)\n",
    "- [Data augmentations](#augmentations)\n",
    "- [Loss function](#loss)\n",
    "- [Model](#model)\n",
    "- [Training](#training)\n",
    "- [Inference](#inference)\n",
    "- [Post-processing](#postprocessing)\n",
    "\n",
    "# Introduction <a name=\"introduction\"></a>\n",
    "\n",
    "The goal of [BraTS 2021 challenge](https://www.med.upenn.edu/cbica/brats2021) was to create a model for segmenting the brain glioblastoma subregions in mpMRI scans. By using our nnU-Net implementation, NVIDIA data scientists have [won the BraTS21 validation phase](https://developer.nvidia.com/blog/nvidia-data-scientists-take-top-spots-in-miccai-2021-brain-tumor-segmentation-challenge). In this notebook, we will share with you the recipe we used for training our nnU-Net for BraTS21 challenge, so that you can reproduce our results. In particular, we will walk you through the following steps: data pre-processing, designing the loss function, building and training the model, running inference and finally post-processing the predictions.\n",
    "\n",
    "# Dataset <a name=\"dataset\"></a>\n",
    "\n",
    "The training dataset provided for the BraTS21 challenge consists of 1,251 brain mpMRI scans along with segmentation annotations of tumorous regions. The 3D volumes were skull-stripped and resampled to 1 mm isotropic resolution, with dimensions of (240, 240, 155) voxels. For each example, four modalities were given: Fluid Attenuated Inversion Recovery (FLAIR), native (T1), post-contrast T1-weighted (T1Gd), and T2-weighted (T2). See image below with each modality. Annotations consist of four classes: 1 for necrotic tumor core (NCR), 2 for peritumoral edematous tissue (ED), 4 for enhancing tumor (ET), and 0 for background (voxels that are not part of the tumor).\n",
    "\n",
    "To download the training and validation dataset, you need to have an account on https://www.synapse.org platform and be registered for BraTS21 challenge. We will assume that after downloading and unzipping, the dataset is organized as follows:\n",
    "\n",
    "```\n",
    "/data \n",
    " │\n",
    " ├───BraTS2021_train\n",
    " │      ├──BraTS2021_00000 \n",
    " │      │      └──BraTS2021_00000_flair.nii.gz\n",
    " │      │      └──BraTS2021_00000_t1.nii.gz\n",
    " │      │      └──BraTS2021_00000_t1ce.nii.gz\n",
    " │      │      └──BraTS2021_00000_t2.nii.gz\n",
    " │      │      └──BraTS2021_00000_seg.nii.gz\n",
    " │      ├──BraTS2021_00002\n",
    " │      │      └──BraTS2021_00002_flair.nii.gz\n",
    " │      ...    └──...\n",
    " │\n",
    " └────BraTS2021_val\n",
    "        ├──BraTS2021_00001 \n",
    "        │      └──BraTS2021_00001_flair.nii.gz\n",
    "        │      └──BraTS2021_00001_t1.nii.gz\n",
    "        │      └──BraTS2021_00001_t1ce.nii.gz\n",
    "        │      └──BraTS2021_00001_t2.nii.gz\n",
    "        ├──BraTS2021_00002\n",
    "        │      └──BraTS2021_00002_flair.nii.gz\n",
    "        ...    └──...\n",
    "```\n",
    "\n",
    "Let's visualize a BraTS2021_00000 example from the training dataset. Each plot presents a different modality (from left to right: FLAIR, T1, T1ce, T2), and an annotation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 100104 characters]\n",
      "text/plain": [
       "<Figure size 1080x1080 with 5 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import nibabel as nib\n",
    "from glob import glob\n",
    "\n",
    "imgs = [nib.load(f\"/data/BraTS2021_train/BraTS2021_00000/BraTS2021_00000_{m}.nii.gz\").get_fdata().astype(np.float32)[:, :, 75] for m in [\"flair\", \"t1\", \"t1ce\", \"t2\"]]\n",
    "lbl = nib.load(\"/data/BraTS2021_train/BraTS2021_00000/BraTS2021_00000_seg.nii.gz\").get_fdata().astype(np.uint8)[:, :, 75]\n",
    "fig, ax = plt.subplots(nrows=1, ncols=5, figsize=(15, 15))\n",
    "for i, img in enumerate(imgs):\n",
    "    ax[i].imshow(img, cmap='gray')\n",
    "    ax[i].axis('off')\n",
    "ax[-1].imshow(lbl, vmin=0, vmax=4)\n",
    "ax[-1].axis('off')\n",
    "plt.tight_layout()            \n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data pre-processing <a name=\"preprocessing\"></a>\n",
    "\n",
    "Each example of the BraTS21 dataset consists of four [NIfTI](https://nifti.nimh.nih.gov/) files with different MRI modalities (filenames with suffixes flair, t1, t1ce, t2). Additionally, examples in the training dataset have a NIfTI with annotation (filename with suffix seg). As a first step of data pre-processing, all four modalities are stacked such that each example has shape (4, 240, 240, 155) (input tensor is in the (C, H, W, D) layout, where C-channels, H-height, W-width and D-depth). Then redundant background voxels (with voxel value zero) on the borders of each volume are [cropped](https://docs.monai.io/en/latest/transforms.html#cropforeground), as they do not provide any useful information and can be ignored by the neural network. Subsequently, for each example, the mean and the standard deviation are computed within the non-zero region for each channel separately. All volumes are [normalized](https://docs.monai.io/en/latest/transforms.html#normalizeintensityd) by first subtracting the mean and then divided by the standard deviation. The background voxels are not normalized so that their value remained at zero. To distinguish between background voxels and normalized voxels which have values close to zero, we add an input channel with one-hot encoding for foreground voxels and stacked with the input data. As a result, each example has 5 channels.\n",
    "\n",
    "Let's start by preparing the raw training and validation datasets into stacked NIfTI files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Preparing BraTS21 dataset from: /data/BraTS2021_train\n",
      "Preparing time: 259.65\n",
      "Preparing BraTS21 dataset from: /data/BraTS2021_val\n",
      "Preparing time: 38.62\n",
      "Finished!\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "import os\n",
    "from glob import glob\n",
    "from subprocess import call\n",
    "import time\n",
    "\n",
    "import nibabel\n",
    "import numpy as np\n",
    "from joblib import Parallel, delayed\n",
    "\n",
    "\n",
    "def load_nifty(directory, example_id, suffix):\n",
    "    return nibabel.load(os.path.join(directory, example_id + \"_\" + suffix + \".nii.gz\"))\n",
    "\n",
    "\n",
    "def load_channels(d, example_id):\n",
    "    return [load_nifty(d, example_id, suffix) for suffix in [\"flair\", \"t1\", \"t1ce\", \"t2\"]]\n",
    "\n",
    "\n",
    "def get_data(nifty, dtype=\"int16\"):\n",
    "    if dtype == \"int16\":\n",
    "        data = np.abs(nifty.get_fdata().astype(np.int16))\n",
    "        data[data == -32768] = 0\n",
    "        return data\n",
    "    return nifty.get_fdata().astype(np.uint8)\n",
    "\n",
    "\n",
    "def prepare_nifty(d):\n",
    "    example_id = d.split(\"/\")[-1]\n",
    "    flair, t1, t1ce, t2 = load_channels(d, example_id)\n",
    "    affine, header = flair.affine, flair.header\n",
    "    vol = np.stack([get_data(flair), get_data(t1), get_data(t1ce), get_data(t2)], axis=-1)\n",
    "    vol = nibabel.nifti1.Nifti1Image(vol, affine, header=header)\n",
    "    nibabel.save(vol, os.path.join(d, example_id + \".nii.gz\"))\n",
    "\n",
    "    if os.path.exists(os.path.join(d, example_id + \"_seg.nii.gz\")):\n",
    "        seg = load_nifty(d, example_id, \"seg\")\n",
    "        affine, header = seg.affine, seg.header\n",
    "        vol = get_data(seg, \"unit8\")\n",
    "        vol[vol == 4] = 3\n",
    "        seg = nibabel.nifti1.Nifti1Image(vol, affine, header=header)\n",
    "        nibabel.save(seg, os.path.join(d, example_id + \"_seg.nii.gz\"))\n",
    "\n",
    "\n",
    "def prepare_dirs(data, train):\n",
    "    img_path, lbl_path = os.path.join(data, \"images\"), os.path.join(data, \"labels\")\n",
    "    call(f\"mkdir {img_path}\", shell=True)\n",
    "    if train:\n",
    "        call(f\"mkdir {lbl_path}\", shell=True)\n",
    "    dirs = glob(os.path.join(data, \"BraTS*\"))\n",
    "    for d in dirs:\n",
    "        if \"_\" in d.split(\"/\")[-1]:\n",
    "            files = glob(os.path.join(d, \"*.nii.gz\"))\n",
    "            for f in files:\n",
    "                if \"flair\" in f or \"t1\" in f or \"t1ce\" in f or \"t2\" in f:\n",
    "                    continue\n",
    "                if \"_seg\" in f:\n",
    "                    call(f\"mv {f} {lbl_path}\", shell=True)\n",
    "                else:\n",
    "                    call(f\"mv {f} {img_path}\", shell=True)\n",
    "        call(f\"rm -rf {d}\", shell=True)\n",
    "\n",
    "\n",
    "def prepare_dataset_json(data, train):\n",
    "    images, labels = glob(os.path.join(data, \"images\", \"*\")), glob(os.path.join(data, \"labels\", \"*\"))\n",
    "    images = sorted([img.replace(data + \"/\", \"\") for img in images])\n",
    "    labels = sorted([lbl.replace(data + \"/\", \"\") for lbl in labels])\n",
    "\n",
    "    modality = {\"0\": \"FLAIR\", \"1\": \"T1\", \"2\": \"T1CE\", \"3\": \"T2\"}\n",
    "    labels_dict = {\"0\": \"background\", \"1\": \"edema\", \"2\": \"non-enhancing tumor\", \"3\": \"enhancing tumour\"}\n",
    "    if train:\n",
    "        key = \"training\"\n",
    "        data_pairs = [{\"image\": img, \"label\": lbl} for (img, lbl) in zip(images, labels)]\n",
    "    else:\n",
    "        key = \"test\"\n",
    "        data_pairs = [{\"image\": img} for img in images]\n",
    "\n",
    "    dataset = {\n",
    "        \"labels\": labels_dict,\n",
    "        \"modality\": modality,\n",
    "        key: data_pairs,\n",
    "    }\n",
    "\n",
    "    with open(os.path.join(data, \"dataset.json\"), \"w\") as outfile:\n",
    "        json.dump(dataset, outfile)\n",
    "\n",
    "\n",
    "def run_parallel(func, args):\n",
    "    return Parallel(n_jobs=os.cpu_count())(delayed(func)(arg) for arg in args)\n",
    "\n",
    "\n",
    "def prepare_dataset(data, train):\n",
    "    print(f\"Preparing BraTS21 dataset from: {data}\")\n",
    "    start = time.time()\n",
    "    run_parallel(prepare_nifty, sorted(glob(os.path.join(data, \"BraTS*\"))))\n",
    "    prepare_dirs(data, train)\n",
    "    prepare_dataset_json(data, train)\n",
    "    end = time.time()\n",
    "    print(f\"Preparing time: {(end - start):.2f}\")\n",
    "\n",
    "prepare_dataset(\"/data/BraTS2021_train\", True)\n",
    "prepare_dataset(\"/data/BraTS2021_val\", False)\n",
    "print(\"Finished!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, lets preprocesses the datasets by cropping and normalizing the volumes. We will store the pre-processed volumes as NumPy arrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Preprocessing /data/BraTS2021_train\n",
      "Pre-processing time: 408.09\n",
      "Preprocessing /data/BraTS2021_val\n",
      "Pre-processing time: 68.28\n",
      "Finished!\n"
     ]
    }
   ],
   "source": [
    "!python3 ../preprocess.py --task 11 --ohe --exec_mode training\n",
    "!python3 ../preprocess.py --task 12 --ohe --exec_mode test\n",
    "print(\"Finished!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Augmentations <a name=\"augmentations\"></a>\n",
    "\n",
    "Data augmentation is a technique that alleviates the overfitting problem by artificially extending a dataset during the training phase. To make our method more robust, the following data augmentations are used during training phase:\n",
    "\n",
    "1. **Biased crop**: From the input volume, a patch of dimensions (5, 128, 128, 128) was randomly cropped. Additionally, with probability of 0.4 the patch selected via random biased crop is guaranteed that some foreground voxels (with positive class in the ground truth) are present in the cropped region.\n",
    "2. **Zoom**: With probability of 0.15, a zoom factor is sampled uniformly from (1.0, 1.4) and then input volume is zoomed by a sampled factor with cubic interpolation, while label with nearest neighbor interpolation.\n",
    "3. **Flips**: With probability of 0.5, for each x, y, z axis independently, volume was flipped along that axis.\n",
    "4. **Gaussian Noise**: With probability of 0.15, random Gaussian noise with mean zero and standard deviation sampled uniformly from (0, 0.33) is sampled for each voxel and added to the input volume. \n",
    "5. **Gaussian Blur**: With probability of 0.15, Gaussian blurring with standard deviation of the Gaussian Kernel sampled uniformly from (0.5, 1.5) is applied to the input volume.\n",
    "6. **Brightness**: With probability of 0.15, a random value is sampled uniformly from (0.7, 1.3) and then input volume voxels are multiplied by it.\n",
    "7. **Contrast**: With probability of 0.15, a random value is sampled uniformly from (0.65, 1.5) and then input volume voxels are multiplied by it and clipped to its original min and max values.\n",
    "\n",
    "The data loading pipeline is implemented with [NVIDIA Data Loading Library (DALI)](https://docs.nvidia.com/deeplearning/dali/user-guide/docs/index.html), which addresses the problem of the CPU bottleneck by offloading data augmentations to the GPU. We encourage you to check out the implementation details of our [DALI pipeline](https://github.com/NVIDIA/DeepLearningExamples/blob/master/PyTorch/Segmentation/nnUNet/data_loading/dali_loader.py)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Loss function <a name=\"loss\"></a>\n",
    "\n",
    "The BraTS leaderboard is computed based on three partially overlapping regions: whole tumor (1, 2, 4), tumor core (1, 4) and enhancing tumor (4), instead of classes present in the labels. Thus, it is beneficial to construct the loss function based on classes used for ranking calculation. Therefore, we optimize each region separately with a sum of binary Cross-Entropy with the Dice loss."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "from monai.losses import DiceLoss\n",
    "\n",
    "class Loss(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Loss, self).__init__()\n",
    "        self.dice = DiceLoss(sigmoid=True, batch=True)\n",
    "        self.ce = nn.BCEWithLogitsLoss()\n",
    "\n",
    "    def _loss(self, p, y):\n",
    "        return self.dice(p, y) + self.ce(p, y.float())\n",
    "\n",
    "    def forward(self, p, y):\n",
    "        y_wt, y_tc, y_et = y > 0, ((y == 1) + (y == 3)) > 0, y == 3\n",
    "        p_wt, p_tc, p_et = p[:, 0].unsqueeze(1), p[:, 1].unsqueeze(1), p[:, 2].unsqueeze(1)\n",
    "        l_wt, l_tc, l_et = self._loss(p_wt, y_wt), self._loss(p_tc, y_tc), self._loss(p_et, y_et)\n",
    "        return l_wt + l_tc + l_et"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model <a name=\"model\"></a>\n",
    "\n",
    "We have made some modifications to the U-Net architecture for the BraTS challenge with respect to the original nnU-Net template. In particular, the U-Net template in the nnU-Net has the encoder depth of 6, and the convolution channels at each encoder level are: 32, 64, 128, 256, 320, 320. Based on the experiments we run, increasing the depth of the encoder to 7, modifying the number of channels to: 64, 96, 128, 192, 256, 384, 512, and using deep supervision improves the final score.\n",
    "\n",
    "For deep supervision, we used two additional output heads at the decoder levels with feature map sizes (64, 64, 64) and (32, 32, 32). To match the shape of the additional predictions with the label shape of (128, 128, 128) we downsampled the label using the nearest neighbor interpolation to the (64, 64, 64) and (32, 32, 32) shapes, so that loss can be computed for additional outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 87272 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython.display import Image\n",
    "Image(filename=\"../images/unet-brats.jpg\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 1: *The final U-Net architecture used for BraTS21 challenge.*\n",
    "\n",
    "# Training <a name=\"training\"></a>\n",
    "\n",
    "Now, let's start training the model. For that, we will call the training script from our nnUNet repo with some additional command line arguments for BraTS challenge: \n",
    "\n",
    "- `--brats` - use loss function with partially overlapping regions (WT, TC, ET) and BraTS specific inference;\n",
    "- `--deep_supervision` - use deep supervision loss with two additional output heads;\n",
    "- `--more_chn` - create encoder with more channels than regular U-Net;\n",
    "- `--min_fmap 2` - create deeper encoder, with feature map size in the bottleneck 2x2x2;\n",
    "\n",
    "and the regular command line arguments:\n",
    "\n",
    "- `--scheduler` - use cosine decay learning rate scheduler with warm up 250 steps of warm up;\n",
    "- `--learning_rate 0.0003` - initial learning rate after warm up will be set to 0.0003;\n",
    "- `--epochs 30` - training will be done for 30 epochs;\n",
    "- `--fold 0` - training will be done for fold 0 (by default, 5-fold cross validation is used);\n",
    "- `--amp` - training with automatic mixed precision, for faster training and memory reduction;\n",
    "- `--gpus 1` - one GPU will be used during training;\n",
    "- `--task 11` - task number for BraTS21 training dataset. See file `data_preprocessing/configs.py` for more details;\n",
    "- `--save_ckpt` - save checkpoint with highest dice score acheived during training.\n",
    "\n",
    "We will run training on 1xA100 GPU. To train the model with [AMP](https://developer.nvidia.com/automatic-mixed-precision), you will need a GPU with at least 15G memory.\n",
    "\n",
    "Here, we will train the model on just 1-fold (fold with index 0) and 30 epochs. For the challenge submission, we have trained 5 models on each fold for 150 epochs, and averaged their predictions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Global seed set to 1\n",
      "Number of examples: Train 1000 - Val 251\n",
      "Filters: [64, 96, 128, 192, 256, 384, 512],\n",
      "Kernels: [[3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3]]\n",
      "Strides: [[1, 1, 1], [2, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2]]\n",
      "GPU available: True, used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "Using native 16bit precision.\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
      "\n",
      "  | Name  | Type      | Params\n",
      "------------------------------------\n",
      "0 | model | UNet      | 50.8 M\n",
      "1 | loss  | LossBraTS | 0     \n",
      "2 | dice  | Dice      | 0     \n",
      "------------------------------------\n",
      "50.8 M    Trainable params\n",
      "0         Non-trainable params\n",
      "50.8 M    Total params\n",
      "203.356   Total estimated model params size (MB)\n",
      "Epoch 0: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 09:47:10.369319 - Epoch: 0  Mean dice : 76.92  Highest : 76.92  L1 : 72.81  L2 : 72.71  L3 : 85.23  val_loss : 0.9542\n",
      "Epoch 1: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 09:50:06.658588 - Epoch: 1  Mean dice : 78.11  Highest : 78.11  L1 : 74.22  L2 : 76.79  L3 : 83.33  val_loss : 0.7793\n",
      "Epoch 2: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 09:53:03.036486 - Epoch: 2  Mean dice : 82.41  Highest : 82.41  L1 : 81.34  L2 : 77.48  L3 : 88.41  val_loss : 0.6169\n",
      "Epoch 3: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 09:55:59.582026 - Epoch: 3  Mean dice : 83.33  Highest : 83.33  L1 : 83.86  L2 : 78.26  L3 : 87.87  val_loss : 0.5845\n",
      "Epoch 4: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 09:58:55.668410 - Epoch: 4  Mean dice : 85.02  Highest : 85.02  L1 : 84.49  L2 : 80.92  L3 : 89.63  val_loss : 0.5207\n",
      "Epoch 5: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 10:01:51.775193 - Epoch: 5  Mean dice : 86.09  Highest : 86.09  L1 : 86.6  L2 : 81.56  L3 : 90.11  val_loss : 0.4856\n",
      "Epoch 6: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 10:04:47.909984 - Epoch: 6  Mean dice : 86.75  Highest : 86.75  L1 : 87.31  L2 : 81.84  L3 : 91.1  val_loss : 0.4545\n",
      "Epoch 7: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 10:07:44.297625 - Epoch: 7  Mean dice : 86.24  Highest : 86.75  L1 : 87.2  L2 : 80.88  L3 : 90.63  val_loss : 0.4771\n",
      "Epoch 8: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 10:10:39.488046 - Epoch: 8  Mean dice : 87.65  Highest : 87.65  L1 : 88.53  L2 : 83.03  L3 : 91.38  val_loss : 0.4257\n",
      "Epoch 9: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 10:13:35.567349 - Epoch: 9  Mean dice : 81.52  Highest : 87.65  L1 : 78.71  L2 : 78.49  L3 : 87.37  val_loss : 0.6421\n",
      "Epoch 10: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 10:16:30.845535 - Epoch: 10  Mean dice : 85.77  Highest : 87.65  L1 : 84.83  L2 : 82.32  L3 : 90.16  val_loss : 0.4978\n",
      "Epoch 11: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 10:19:26.081598 - Epoch: 11  Mean dice : 86.78  Highest : 87.65  L1 : 87.52  L2 : 82.28  L3 : 90.53  val_loss : 0.4516\n",
      "Epoch 12: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 10:22:21.393119 - Epoch: 12  Mean dice : 87.57  Highest : 87.65  L1 : 87.97  L2 : 83.67  L3 : 91.08  val_loss : 0.4225\n",
      "Epoch 13: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 10:25:16.783225 - Epoch: 13  Mean dice : 86.75  Highest : 87.65  L1 : 88.16  L2 : 82.88  L3 : 89.2  val_loss : 0.4548\n",
      "Epoch 14: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 10:28:12.184295 - Epoch: 14  Mean dice : 86.88  Highest : 87.65  L1 : 86.13  L2 : 82.73  L3 : 91.77  val_loss : 0.4468\n",
      "Epoch 15: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 10:31:07.374641 - Epoch: 15  Mean dice : 86.45  Highest : 87.65  L1 : 86.58  L2 : 83.46  L3 : 89.31  val_loss : 0.4728\n",
      "Epoch 16: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 10:34:03.013703 - Epoch: 16  Mean dice : 87.56  Highest : 87.65  L1 : 88.26  L2 : 82.44  L3 : 91.99  val_loss : 0.4162\n",
      "Epoch 17: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 10:36:58.202911 - Epoch: 17  Mean dice : 86.81  Highest : 87.65  L1 : 86.95  L2 : 83.3  L3 : 90.19  val_loss : 0.4453\n",
      "Epoch 18: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 10:39:53.265586 - Epoch: 18  Mean dice : 88.21  Highest : 88.21  L1 : 88.7  L2 : 84.11  L3 : 91.81  val_loss : 0.4033\n",
      "Epoch 19: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 10:42:49.002307 - Epoch: 19  Mean dice : 87.69  Highest : 88.21  L1 : 87.31  L2 : 83.86  L3 : 91.9  val_loss : 0.4166\n",
      "Epoch 20: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 10:45:44.182336 - Epoch: 20  Mean dice : 88.9  Highest : 88.9  L1 : 89.62  L2 : 84.57  L3 : 92.52  val_loss : 0.3775\n",
      "Epoch 21: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 10:48:40.095388 - Epoch: 21  Mean dice : 88.38  Highest : 88.9  L1 : 89.7  L2 : 83.1  L3 : 92.35  val_loss : 0.3971\n",
      "Epoch 22: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 10:51:35.371210 - Epoch: 22  Mean dice : 88.05  Highest : 88.9  L1 : 88.21  L2 : 84.05  L3 : 91.88  val_loss : 0.4091\n",
      "Epoch 23: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 10:54:30.622294 - Epoch: 23  Mean dice : 88.54  Highest : 88.9  L1 : 89.35  L2 : 83.88  L3 : 92.4  val_loss : 0.3882\n",
      "Epoch 24: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 10:57:25.706283 - Epoch: 24  Mean dice : 88.86  Highest : 88.9  L1 : 90.19  L2 : 84.43  L3 : 91.97  val_loss : 0.3735\n",
      "Epoch 25: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 11:00:20.871930 - Epoch: 25  Mean dice : 88.72  Highest : 88.9  L1 : 89.98  L2 : 84.97  L3 : 91.22  val_loss : 0.3824\n",
      "Epoch 26: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 11:03:15.888703 - Epoch: 26  Mean dice : 89.26  Highest : 89.26  L1 : 90.43  L2 : 84.83  L3 : 92.53  val_loss : 0.3639\n",
      "Epoch 27: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 11:06:12.074362 - Epoch: 27  Mean dice : 88.67  Highest : 89.26  L1 : 89.51  L2 : 84.29  L3 : 92.21  val_loss : 0.3861\n",
      "Epoch 28: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 11:09:07.148004 - Epoch: 28  Mean dice : 89.4  Highest : 89.4  L1 : 90.79  L2 : 85.35  L3 : 92.06  val_loss : 0.36\n",
      "Epoch 29: 100%|█████████████████████████████████████████████████████████████████████████████████████████▉| 751/751 [02:53<00:00]\n",
      "DLL 2021-10-07 11:12:03.077000 - Epoch: 29  Mean dice : 89.69  Highest : 89.69  L1 : 90.69  L2 : 85.5  L3 : 92.88  val_loss : 0.3489"
     ]
    }
   ],
   "source": [
    "!python ../main.py --brats --deep_supervision --depth 6 --filters 64 96 128 192 256 384 512 --min_fmap 2 --scheduler --learning_rate 0.0003 --epochs 30 --fold 0 --amp --gpus 1 --task 11 --save_ckpt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inference <a name=\"inference\"></a>\n",
    "\n",
    "During inference, the input volume can have arbitrary size, instead of the fixed patch size (128, 128, 128) as during the training phase. Thus, we use a [sliding window inference](https://docs.monai.io/en/latest/inferers.html) from [MONAI](https://monai.io/) library, where the window has the same size as the training patch, i.e., (128, 128, 128) and adjacent windows overlap by half the size of a patch. The predictions on the overlapping regions are then averaged with Gaussian importance weighting, such that the weights of the center voxels have higher importance.\n",
    "\n",
    "One of the known tricks to improve predictions robustness is to apply test time augmentations (TTA). During inference, we are creating eight versions of the input volume, such that each version corresponds to one of eight possible flips along the x, y, z axis combination. Then we run inference for each version of the input volume and transform the predictions back to the original input volume orientation by applying the same flips to predictions as was used for the input volume. Finally, the probabilities from all predictions were averaged.\n",
    "\n",
    "Let's run inference with TTA on the challenge validation dataset.\n",
    "\n",
    "Note: You will have to modify the `--ckpt_path` argument, such that the path to checkpoint is valid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Global seed set to 1\n",
      "Number of test examples: 219\n",
      "Filters: [64, 96, 128, 192, 256, 384, 512],\n",
      "Kernels: [[3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3]]\n",
      "Strides: [[1, 1, 1], [2, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2]]\n",
      "GPU available: True, used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "Using native 16bit precision.\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
      "Testing: 100%|███████████████████████████████████████████████████████████████████████████████████| 219/219 [10:39<00:00]\n"
     ]
    }
   ],
   "source": [
    "!python ../main.py --gpus 1 --amp --save_preds --exec_mode predict --brats --data /data/12_3d/test --ckpt_path /results/checkpoints/epoch=29-dice_mean=89.69.ckpt --tta"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Post-processing <a name=\"postprocessing\"></a>\n",
    "\n",
    "By optimizing the three overlapping regions (ET, TC, WT) we need to convert them back to the original classes (NCR, ED, ET). The strategy for transforming classes back to the original one is the following: if the WT probability for a given voxel is less than 0.45 then its class is set to 0 (background), otherwise if the probability for TC is less than 0.4 the voxel class is 2 (ED), and finally if probability for ET is less than 0.4 voxel has class 1 (NCR), or otherwise 4 (ET).\n",
    "\n",
    "Furthermore, we applied the following post-processing strategy: find ET connected components, for components smaller than 16 voxels with mean probability smaller than 0.9, replace their class to NCR (such that voxels are still considered part of the tumor core), next if there is overall less than 73 voxels with ET and their mean probability is smaller than 0.9 replace all ET voxels to NCR. With such post-processing we avoided the edge case where the model predicted a few voxels with enhancing tumor (ET) but there were not any in the ground truth. Such post-processing was beneficial to the final score as if there were no enhancing tumor voxels in the label, then the Dice score for zero false positive prediction was 1, and 0 otherwise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Preparing final predictions\n",
      "Finished!\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from glob import glob\n",
    "from subprocess import call\n",
    "\n",
    "import nibabel as nib\n",
    "import numpy as np\n",
    "from scipy.ndimage.measurements import label\n",
    "\n",
    "\n",
    "def to_lbl(pred):\n",
    "    enh = pred[2]\n",
    "    c1, c2, c3 = pred[0] > 0.5, pred[1] > 0.5, pred[2] > 0.5\n",
    "    pred = (c1 > 0).astype(np.uint8)\n",
    "    pred[(c2 == False) * (c1 == True)] = 2\n",
    "    pred[(c3 == True) * (c1 == True)] = 4\n",
    "\n",
    "    components, n = label(pred == 4)\n",
    "    for et_idx in range(1, n + 1):\n",
    "        _, counts = np.unique(pred[components == et_idx], return_counts=True)\n",
    "        if 1 < counts[0] and counts[0] < 8 and np.mean(enh[components == et_idx]) < 0.9:\n",
    "            pred[components == et_idx] = 1\n",
    "\n",
    "    et = pred == 4\n",
    "    if 0 < et.sum() and et.sum() < 73 and np.mean(enh[et]) < 0.9:\n",
    "        pred[et] = 1\n",
    "\n",
    "    pred = np.transpose(pred, (2, 1, 0)).astype(np.uint8)\n",
    "    return pred\n",
    "\n",
    "\n",
    "def prepare_preditions(e):\n",
    "    fname = e[0].split(\"/\")[-1].split(\".\")[0]\n",
    "    preds = [np.load(f) for f in e]\n",
    "    p = to_lbl(np.mean(preds, 0))\n",
    "\n",
    "    img = nib.load(f\"/data/BraTS2021_val/images/{fname}.nii.gz\")\n",
    "    nib.save(\n",
    "        nib.Nifti1Image(p, img.affine, header=img.header),\n",
    "        os.path.join(\"/results/final_preds\", fname + \".nii.gz\"),\n",
    "    )\n",
    "\n",
    "os.makedirs(\"/results/final_preds\")\n",
    "preds = sorted(glob(f\"/results/predictions*\"))\n",
    "examples = list(zip(*[sorted(glob(f\"{p}/*.npy\")) for p in preds]))\n",
    "print(\"Preparing final predictions\")\n",
    "for e in examples:\n",
    "    prepare_preditions(e)\n",
    "print(\"Finished!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualization\n",
    "\n",
    "Let's visualize the final prediction made on the challenge validation dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BraTS2021_00001\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 92336 characters]\n",
      "text/plain": [
       "<Figure size 864x864 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BraTS2021_00013\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 83724 characters]\n",
      "text/plain": [
       "<Figure size 864x864 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BraTS2021_00015\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 83680 characters]\n",
      "text/plain": [
       "<Figure size 864x864 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BraTS2021_00027\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 79900 characters]\n",
      "text/plain": [
       "<Figure size 864x864 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BraTS2021_00037\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 84604 characters]\n",
      "text/plain": [
       "<Figure size 864x864 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import nibabel as nib\n",
    "from glob import glob\n",
    "\n",
    "n, z = 5, 75\n",
    "data = sorted(glob(\"/results/final_preds/*.nii.gz\"))\n",
    "for i in range(n):\n",
    "    fname = data[i].split(\"/\")[-1].split(\".\")[0]\n",
    "    print(fname)\n",
    "    img = nib.load(f\"/data/BraTS2021_val/images/{fname}.nii.gz\").get_fdata().astype(np.float32)\n",
    "    pred = nib.load(data[i]).get_fdata().astype(np.uint8)[:, :, z]\n",
    "    imgs = [img[:, :, z, i] for i in [0, 3]] + [pred]\n",
    "    \n",
    "    fig, ax = plt.subplots(nrows=1, ncols=3, figsize=(12, 12))\n",
    "    for i in range(3):\n",
    "        if i < 2:\n",
    "            ax[i].imshow(imgs[i], cmap='gray')\n",
    "        else:\n",
    "            ax[i].imshow(imgs[i]);\n",
    "        ax[i].axis('off')  \n",
    "    plt.tight_layout()            \n",
    "    plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
