4. MONAI Post-Transforms and Inference Methods.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using Google Colab\n",
    "\n",
    "This notebook has the pip command for installing MONAI and will be added to any subsequent notebook.\n",
    "\n",
    "**Required Packages for Colab Execution**\n",
    "\n",
    "Execute the following cell to install MONAI the first time a colab notebook is run:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -qU \"monai[ignite, nibabel, torchvision, tqdm]==0.5.0\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Enabling GPU Support**\n",
    "\n",
    "To use GPU resources through Colab, change the runtime to GPU:\n",
    "\n",
    "1. From the **\"Runtime\"** menu select **\"Change Runtime Type\"**\n",
    "2. Choose **\"GPU\"** from the drop-down menu\n",
    "3. Click **\"SAVE\"**\n",
    "\n",
    "This will reset the notebook and probably ask you if you are a robot (these instructions assume you are not). Running\n",
    "\n",
    "**!nvidia-smi**\n",
    "\n",
    "in a cell will verify this has worked and show you what kind of hardware you have access to.    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MONAI Sliding Inference and Post-processing\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To run model inferences and evaluate the model quality, MONAI provides reference implementations for the relevant widely-used approaches. Currently, several popular evaluation metrics and inference patterns are included.\n",
    "\n",
    "MONAI also includes post-processing tranform functions to help handle the model outputs like removing segmentation noise or extracting contour of segmentation results.\n",
    "\n",
    "## MONAI Sliding Inference and Post-processing\n",
    "\n",
    "To help you understand more about MONAI Datasets and Caching options, this guide will help you answer five key questions:\n",
    "\n",
    "1. **What is sliding inference?**\n",
    "2. **What are post-processing transforms?**\n",
    "3. **How do I use sliding inference on real data?**\n",
    "4. **How do I use post-processing on my sliding inference results?**\n",
    "5. **How can I use tensorboard to visualize the results?**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "Let's get started by importing our dependencies.  We'll also use the Jupyter load extension command to load tensorboard so we can vizualize our results later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext tensorboard\n",
    "\n",
    "import monai\n",
    "monai.config.print_config()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **1. What is sliding inference?**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### A toy model for inference\n",
    "For model inferences on large volumes, the sliding window approach is a popular choice to achieve high performance while having flexible memory requirements. It also supports overlap and blending_mode configurations to handle the overlapped windows for better performances.\n",
    "\n",
    "A typical process is:\n",
    "\n",
    "- Select continuous windows on the original image.\n",
    "- Iteratively run batched window inferences until all windows are analyzed.\n",
    "- Aggregate the inference outputs to a single segmentation map.\n",
    "- Save the results to file or compute some evaluation metrics.\n",
    "\n",
    "<img src=\"sliding_window.png\" style=\"width: 700px;\"/>\n",
    "\n",
    "The [sliding_window_inference](https://docs.monai.io/en/latest/inferers.html?highlight=sliding#sliding-window-inference) requires a callable function which takes a batch of image windows as the input.\n",
    "\n",
    "Here we construct a toy model. It has a single model parameter, `self.pred`. The inference outcome is just `input + self.pred`.\n",
    "\n",
    "Every time the model is called, it also increases `self.pred` by one. This is to demonstrate that the model can be \"stateful\", and also so that we can conveniently visualize the inference outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ToyModel:\n",
    "    # A simple model generates the output by adding an integer `pred` to input.\n",
    "    # each call of this instance increases the integer by 1.\n",
    "    pred = 0\n",
    "    def __call__(self, input):\n",
    "        self.pred = self.pred + 1\n",
    "        return input + self.pred"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Run the inference using sliding window\n",
    "\n",
    "We're going to create a 200x200-pixel image and pass it to `sliding_window_inference` with a 40x40 window size then display the image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "from monai.inferers import sliding_window_inference\n",
    "\n",
    "input_tensor = torch.zeros(1, 1, 200, 200)\n",
    "output_tensor = sliding_window_inference(\n",
    "    inputs=input_tensor, \n",
    "    predictor=ToyModel(), \n",
    "    roi_size=(40, 40), \n",
    "    sw_batch_size=1, \n",
    "    overlap=0.5, \n",
    "    mode=\"constant\")\n",
    "plt.imshow(output_tensor[0, 0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Gaussian weighted windows\n",
    "For a given input image window, the convolutional neural networks often predict the central regions more accurately than the border regions, usually due to the stacked convolutions' receptive field.\n",
    "\n",
    "Therefore, it is worth considering a \"Gaussian weighted\" prediction to emphasize the central region predictions when we stitch the windows into a complete inference output.\n",
    "\n",
    "By simply changing the inference mode to \"gaussian\", the sliding window module will use this \"weighted stitching\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_tensor = torch.zeros(1, 1, 200, 200)\n",
    "output_tensor_1 = sliding_window_inference(\n",
    "    inputs=input_tensor, \n",
    "    predictor=ToyModel(), \n",
    "    roi_size=(40, 40), \n",
    "    sw_batch_size=1, \n",
    "    overlap=0.5, \n",
    "    mode=\"gaussian\")\n",
    "plt.imshow(output_tensor_1[0, 0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compared with the previous inferences, the overlapping windows are stitched together with fewer border artifacts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.subplots(1, 2)\n",
    "plt.subplot(1, 2, 1); plt.imshow(output_tensor[0, 0])\n",
    "plt.subplot(1, 2, 2); plt.imshow(output_tensor_1[0, 0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **2. What are post-processing transforms?**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "MONAI also provides post-processing transforms for handling the model outputs. Currently, the transforms include:\n",
    "\n",
    "- Adding activation layer (Sigmoid, Softmax, etc.).\n",
    "- Converting to discrete values (Argmax, One-Hot, Threshold value, etc), as below figure (b).\n",
    "- Splitting multi-channel data into multiple single channels.\n",
    "- Removing segmentation noise based on Connected Component Analysis, as below figure (c).\n",
    "- Extracting contour of segmentation result, which can be used to map to original image and evaluate the model, as below figure (d) and (e).\n",
    "\n",
    "After applying the post-processing transforms, it’s easier to compute metrics, save model output into files or visualize data in the TensorBoard.\n",
    "\n",
    "<img src=\"post_transforms.png\" style=\"width: 700px;\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **3. How do I use sliding inference on real data?**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "This section will set up and load a [SegResNet](https://docs.monai.io/en/latest/networks.html?highlight=segresnet#segresnet) model, run sliding window inference, and post-process the model output volumes:\n",
    "- Argmax to get a discrete prediction map\n",
    "- Remove small isolated predicted regions\n",
    "- Convert the segmentation regions into contours\n",
    "\n",
    "We'll start by importing all of our dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "\n",
    "from monai.apps import download_and_extract\n",
    "from monai.utils import set_determinism\n",
    "from monai.data import CacheDataset, DataLoader\n",
    "from monai.networks.nets import SegResNet\n",
    "from monai.transforms import (\n",
    "    AddChanneld,\n",
    "    AsDiscrete,\n",
    "    Compose,\n",
    "    CropForegroundd,\n",
    "    KeepLargestConnectedComponent,\n",
    "    LabelToContour,\n",
    "    LoadImaged,\n",
    "    Orientationd,\n",
    "    ScaleIntensityRanged,\n",
    "    Spacingd,\n",
    "    ToTensord,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Download image and labels\n",
    "\n",
    "We'll download the .tar file for Task09_Spleen from the link below.  We'll then extract the data, so it's ready to use in the next section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resource = \"https://msd-for-monai.s3-us-west-2.amazonaws.com/Task09_Spleen.tar\"\n",
    "md5 = \"410d4a301da4e5b2f6f86ec3ddba524e\"\n",
    "\n",
    "root_dir = '.'\n",
    "compressed_file = os.path.join(root_dir, \"Task09_Spleen.tar\")\n",
    "data_dir = os.path.join(root_dir, \"Task09_Spleen\")\n",
    "download_and_extract(resource, compressed_file, root_dir, md5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up the validation data, preprocessing transforms, and data loader\n",
    "\n",
    "We'll put the data and labels into a data dictionary, create a sequence of transforms with `Compose`, using a `CacheDataset`, and then load the data using `DataLoader`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images = sorted(glob.glob(os.path.join(data_dir, \"imagesTr\", \"*.nii.gz\")))\n",
    "labels = sorted(glob.glob(os.path.join(data_dir, \"labelsTr\", \"*.nii.gz\")))\n",
    "data_dicts = [\n",
    "    {\"image\": image_name, \"label\": label_name}\n",
    "    for image_name, label_name in zip(images, labels)\n",
    "]\n",
    "val_files = data_dicts[-9:]\n",
    "\n",
    "val_transforms = Compose(\n",
    "    [\n",
    "        LoadImaged(keys=[\"image\", \"label\"]),\n",
    "        AddChanneld(keys=[\"image\", \"label\"]),\n",
    "        Spacingd(keys=[\"image\", \"label\"], pixdim=(1.5, 1.5, 2.0), mode=(\"bilinear\", \"nearest\")),\n",
    "        Orientationd(keys=[\"image\", \"label\"], axcodes=\"RAS\"),\n",
    "        ScaleIntensityRanged(\n",
    "            keys=[\"image\"], a_min=-57, a_max=164, b_min=0.0, b_max=1.0, clip=True,\n",
    "        ),\n",
    "        CropForegroundd(keys=[\"image\", \"label\"], source_key=\"image\"),\n",
    "        ToTensord(keys=[\"image\", \"label\"]),\n",
    "    ]\n",
    ")\n",
    "val_ds = CacheDataset(data=val_files, transform=val_transforms, cache_rate=1.0, num_workers=4)\n",
    "val_loader = DataLoader(val_ds, batch_size=1, num_workers=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Set up the model\n",
    "\n",
    "We want to utilize the GPU, so we'll check to see if it's available and set it as the primary device; otherwise we'll use the CPU.\n",
    "\n",
    "We'll then instantiate the model with the selected device."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch \n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu:0\")\n",
    "model = SegResNet(\n",
    "    spatial_dims=3,\n",
    "    in_channels=1,\n",
    "    out_channels=2,\n",
    ").to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_path = os.path.join(root_dir, \"segresnet_model_epoch30.pth\")\n",
    "model.load_state_dict(torch.load(model_path))\n",
    "model.eval()\n",
    "print(f\"model from {model_path}.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Run the sliding window inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_data = next(iter(val_loader))\n",
    "val_data = val_data[\"image\"].to(device)\n",
    "\n",
    "roi_size = (88, 88, 88)\n",
    "sw_batch_size = 1\n",
    "with torch.no_grad():\n",
    "  val_output = sliding_window_inference(\n",
    "      val_data, roi_size, sw_batch_size=sw_batch_size, predictor=model, mode=\"gaussian\", overlap=0.2)\n",
    "print(val_output.shape, val_output.device)\n",
    "\n",
    "slice_idx = 80\n",
    "plt.title(f\"image -- slice {slice_idx}\")\n",
    "plt.imshow(val_output.detach().cpu()[0, 1, :, :, 80], cmap=\"gray\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **4. How do I use post-processing on my sliding inference results?**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Post-processing: argmax over the output probabilities into a discrete map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "argmax = AsDiscrete(argmax=True)(val_output)\n",
    "print(argmax.shape)\n",
    "\n",
    "slice_idx = 80\n",
    "plt.subplots(1, 2)\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.title(f\"image -- slice {slice_idx}\")\n",
    "plt.imshow(val_data.detach().cpu()[0, 0, :, :, 80], cmap=\"gray\")\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.title(f\"argmax -- slice {slice_idx}\")\n",
    "plt.imshow(argmax.detach().cpu()[0, 0, :, :, 80])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Post-processing: connected component analysis to select the largest segmentation region"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "largest = KeepLargestConnectedComponent(applied_labels=[1])(argmax)\n",
    "print(largest.shape)\n",
    "\n",
    "slice_idx = 80\n",
    "plt.subplots(1, 2)\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.title(f\"image -- slice {slice_idx}\")\n",
    "plt.imshow(val_data.detach().cpu()[0, 0, :, :, 80], cmap=\"gray\")\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.title(f\"largest component -- slice {slice_idx}\")\n",
    "plt.imshow(largest.detach().cpu()[0, 0, :, :, 80])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Post-processing: convert the region into a contour map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "contour = LabelToContour()(largest)\n",
    "print(contour.shape)\n",
    "\n",
    "slice_idx = 80\n",
    "plt.subplots(1, 2)\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.title(f\"image -- slice {slice_idx}\")\n",
    "plt.imshow(val_data.detach().cpu()[0, 0, :, :, 80], cmap=\"gray\")\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.title(f\"contour -- slice {slice_idx}\")\n",
    "plt.imshow(contour.detach().cpu()[0, 0, :, :, 80], cmap=\"Greens\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Visualise the contour over the original input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_image = contour + val_data\n",
    "\n",
    "slice_idx = 80\n",
    "plt.subplots(1, 2)\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.title(f\"image -- slice {slice_idx}\")\n",
    "plt.imshow(val_data.detach().cpu()[0, 0, :, :, 80], cmap=\"gray\")\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.title(f\"contour -- slice {slice_idx}\")\n",
    "plt.imshow(map_image.detach().cpu()[0, 0, :, :, 80], cmap=\"gray\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For more details about the post-postprocessing transformations, please visit:\n",
    "https://docs.monai.io/en/latest/transforms.html#post-processing-dict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. **How can I use tensorboard to visualize the results?**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualising the results in animation with tensorboard (please start tensorboard manually in a separate terminal):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from monai.visualize import plot_2d_or_3d_image\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "\n",
    "with SummaryWriter(log_dir=\"./\") as writer:\n",
    "    plot_2d_or_3d_image(map_image, step=0, writer=writer, tag=\"segmentation\")\n",
    "    plot_2d_or_3d_image(val_output, step=0, max_channels=2, writer=writer, tag=\"Probability\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%tensorboard --logdir=\"./\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **Summary**\n",
    "\n",
    "We've covered MONAI Sliding Inference and Post-Processing. Here are some key highlights:\n",
    "\n",
    "- Sliding inference runs inference on a moving window of images.  It can be helpful with large volume images and can help improve performance.\n",
    "- Post Processing transforms can help make handling the model output easier"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **Next Steps**\n",
    "\n",
    "Start exploring MONAI on your own!  There are lots of great tutorials that can help guide you along the way. You can find it on our [GitHub Organization Page](https://github.com/Project-MONAI/tutorials).  We also have all of our videos from our first ever MONAI Bootcamp available on our [Youtube Channel](https://www.youtube.com/c/ProjectMONAI)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "monai",
   "language": "python",
   "name": "monai"
  },
  "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.9.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}