01_getting_started_with_MONAI.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"./assets_bootcamp/monai.png\" style=\"width: 700px;\"/>\n",
    "\n",
    "### Acknowledgment\n",
    "\n",
    "This notebook was directly imported from the official [MONAI Bootcamp repository](https://github.com/Project-MONAI/monai-bootcamp/tree/main/MONAICore).\n",
    "\n",
    "### Welcome Note\n",
    "\n",
    "Welcome to the MONAI bootcamp! This notebook will introduce you to the MONAI design and architecture, and then hands-on with its transforms."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using Google Colab\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/nvahmadi/NVIDIA_CAMP_Workshop/blob/main/01_getting_started_with_MONAI.ipynb)\n",
    "\n",
    "This notebook has the pip command for installing MONAI and will be added to any subsequent notebook.\n",
    "\n",
    "**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.    \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": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!python -c \"import monai\" || pip install -qU \"monai[ignite, nibabel, torchvision, tqdm]==1.0.0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Getting Started with MONAI\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "MONAI is a PyTorch-based, open-source framework for deep learning in healthcare imaging, part of PyTorch Ecosystem.\n",
    "\n",
    "*Its ambitions are:*\n",
    "\n",
    "- developing a community of academic, industrial and clinical researchers collaborating on a common foundation;\n",
    "- creating state-of-the-art, end-to-end training workflows for healthcare imaging;\n",
    "- providing researchers with the optimized and standardized way to create and evaluate deep learning models."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "MONAI aims at supporting deep learning in medical image analysis at multiple granularities. This figure shows a typical example of the end-to-end workflow in medical deep learning area:\n",
    "\n",
    "<img src=\"./assets_bootcamp/end_to_end.png\" style=\"width: 700px;\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What's the Need?\n",
    "\n",
    "* Biomedical applications have specific requirements\n",
    "* Image modalities (MR, CT, US, etc.) require specific data processing\n",
    "* Data formats (DICOM, NIfTI, etc.) are specific to medical applications and require special support\n",
    "* Certain network architectures are designed for, or are highly suitable for, biomedical applications\n",
    "* Data transforms specific to biomedical applications, and to image modalities, are very useful when pre-processing data, augmenting data during training, and post-processing\n",
    "* Reproducible science requires reproducible experiments which in turn rely on software accessible to other scientists, even if just as a common baseline\n",
    "* A community-driven library to provide solutions to these requirements reduces duplication/re-implementation\n",
    "* Baseline implementations of common networks, and implementations of networks specific to certain papers, provides a basis for comparison between other's networks and results\n",
    "\n",
    "## How Does MONAI Address This Need?\n",
    "\n",
    "MONAI provides a framework of deep learning facilities and infrastructure to meet these needs in a flexible Pytorch-compatible way:\n",
    "* Data loading and handling library for biomedical file types\n",
    "* Large set of data transforms to process, regularize, and augment image date before, during, and after training\n",
    "* Library of general-purpose network, metric, and loss function definitions implementing common architectures\n",
    "* Set of ready-made components for training and inference to utilize computing infrastructure efficiently\n",
    "\n",
    "MONAI contributes to reproducibilty by making the training and distribution of experiments easier:\n",
    "* Common underlying framework betweeen different users' implementations aids in comparison and interoperability\n",
    "* Users can pick and choose components and existing networks then contribute new code to make their results available to others\n",
    "* Reference implementations of networks from papers provides baselines for replication studies and comparisons\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## MONAI architecture\n",
    "The design principle of MONAI is to provide flexible and light APIs for users with varying expertise. All the core components are independent modules, which can be easily integrated into any existing PyTorch programs. Users can leverage the workflows in MONAI to quickly set up a robust training or evaluation program for research experiments.  Rich examples and demos are provided to demonstrate the key features. Researchers contribute implementations based on the state-of-the-art for the latest research challenges, including COVID-19 image analysis, Model Parallel, etc.\n",
    "\n",
    "<img src=\"./assets_bootcamp/arch_modules_v0.4.png\" style=\"width: 700px;\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## MONAI Design Philosophy\n",
    "\n",
    "Key principles:\n",
    "* MONAI looks like PyTorch, uses/extends PyTorch types and adheres to it's general design philosophy\n",
    "* MONAI is additive on top of PyTorch, providing extensions or wrappers\n",
    "* MONAI is opt-in and incremental, no need to rewrite entire models to integrate existing code\n",
    "* MONAI is collaborative, providing adapters and loosely coupled components to ease integration with third party code\n",
    "* MONAI is PyTorch ecosystem friendly, and part of the official ecosystem!\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## MONAI Transforms\n",
    "\n",
    "To help you understand more about MONAI transforms, this guide will help you answer five key questions:\n",
    "\n",
    "1. **What transforms are available to help create a data pipeline for training?**\n",
    "2. **What are array transforms?**\n",
    "3. **What is required to write a custom transform?**\n",
    "4. **What are dictionary transforms?**\n",
    "5. **How do I create a basic MONAI dataset with transforms?**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's get started by importing our dependecies.  We're going to load everything that we'll need for the remainder of the notebook here.  You'll see a lot of import statements, but we'll make sure to go over each of them throughout the rest of the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tempfile\n",
    "import nibabel as nib\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from typing import Optional, Any, Mapping, Hashable\n",
    "\n",
    "import torch\n",
    "\n",
    "import monai\n",
    "from monai.config import print_config\n",
    "from monai.utils import first\n",
    "from monai.config import KeysCollection\n",
    "from monai.data import Dataset, ArrayDataset, create_test_image_3d, DataLoader\n",
    "from monai.transforms import (\n",
    "    Transform,\n",
    "    MapTransform,\n",
    "    Randomizable,\n",
    "    EnsureChannelFirst,\n",
    "    EnsureChannelFirstd,\n",
    "    Compose,\n",
    "    LoadImage,\n",
    "    LoadImaged,\n",
    "    Lambda,\n",
    "    Lambdad,\n",
    "    RandSpatialCrop,\n",
    "    RandSpatialCropd,\n",
    "    Orientation, \n",
    "    Rotate\n",
    ")\n",
    "print_config()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **1. What transforms are available to help create a data pipeline for training?**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Medical image data I/O, processing and augmentation\n",
    "\n",
    "Medical images require highly specialized methods for I/O, preprocessing, and augmentation. Medical images are often in specialized formats with rich meta-information, and the data volumes are often high-dimensional. These require carefully designed manipulation procedures. The medical imaging focus of MONAI is enabled by powerful and flexible image transformations that facilitate user-friendly, reproducible, optimized medical data pre-processing pipelines.\n",
    "\n",
    "<img src=\"./assets_bootcamp/medical_transforms.png\" style=\"width: 700px;\"/>\n",
    "\n",
    "#### Transforms support both Dictionary and Array format data\n",
    "The widely used computer vision packages (such as torchvision) focus on spatially 2D array image processing. MONAI provides more domain-specific transformations for both spatially 2D and 3D and retains the flexible transformation “compose” feature.\n",
    "\n",
    "As medical image preprocessing often requires additional fine-grained system parameters, MONAI provides transforms for input data encapsulated in python dictionaries. Users can specify the keys corresponding to the expected data fields and system parameters to compose complex transformations.\n",
    "\n",
    "There is a rich set of transforms in six categories: Crop & Pad, Intensity, IO, Post-processing, Spatial, and Utilities. For more details, please visit all the transforms in MONAI.\n",
    "\n",
    "#### Medical specific transforms\n",
    "MONAI aims at providing a comprehensive medical image specific transformations. These currently include, for example:\n",
    "\n",
    "    - LoadImage: Load medical specific formats file from provided path\n",
    "    - Spacing: Resample input image into the specified pixdim\n",
    "    - Orientation: Change the image’s orientation into the specified axcodes\n",
    "    - RandGaussianNoise: Perturb image intensities by adding statistical noises\n",
    "    - NormalizeIntensity: Intensity Normalization based on mean and standard deviation\n",
    "    - Affine: Transform image based on the affine parameters\n",
    "    - Rand2DElastic: Random elastic deformation and affine in 2D\n",
    "    - Rand3DElastic: Random elastic deformation and affine in 3D"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll create a temporary directory and populate it with a few example Nifti file-format images containing a random assortment of spheres.  We're also creating a matching segmentation pair that will be used later in the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fn_keys = (\"img\", \"seg\")  # filename keys for image and seg files\n",
    "\n",
    "root_dir = tempfile.mkdtemp()\n",
    "filenames = []\n",
    "\n",
    "for i in range(5):\n",
    "    im, seg = create_test_image_3d(128, 128, 128, num_objs=16, rad_max=25)\n",
    "\n",
    "    im_filename = f\"{root_dir}/im{i}.nii.gz\"\n",
    "    seg_filename = f\"{root_dir}/seg{i}.nii.gz\"\n",
    "    filenames.append({\"img\": im_filename, \"seg\": seg_filename})\n",
    "\n",
    "    n = nib.Nifti1Image(im, np.eye(4))\n",
    "    nib.save(n, im_filename)\n",
    "\n",
    "    n = nib.Nifti1Image(seg, np.eye(4))\n",
    "    nib.save(n, seg_filename)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. **What are array transforms?**\n",
    "\n",
    "Transforms in MONAI are callable objects accepting inputs from initial data in a dataset or previous transforms. We can create and call these directly without any infrastructure or system setup as components in MONAI are designed to be as decoupled as possible. For example we can load one of our Nifti files directly by creating the transform and calling it."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Transforms are composed with `Compose` to create a sequence of operations. `Compose` itself being a transform we can also call it directly. The type of img here is `MetaTensor` which is our extension of Pytorch's `Tensor` class adding metadata dictionary support:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trans = Compose([LoadImage(image_only=True), EnsureChannelFirst()])\n",
    "img = trans(filenames[0][\"img\"])\n",
    "print(type(img), img.shape, img.get_device())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **3. How do I create a custom transform?**\n",
    "We can define our own custom transform operation in a number of ways. If a simple callable is used as an operator, `Lambda` can be used to wrap it as a transform. We define in this example a transform to sum the image in the 1st (width) dimension to produce a 2D image:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sum_width(img):\n",
    "    return img.sum(1)\n",
    "\n",
    "\n",
    "trans = Compose([LoadImage(image_only=True), EnsureChannelFirst(), Lambda(sum_width)])\n",
    "img = trans(filenames[0][\"img\"])\n",
    "plt.imshow(img[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Creating a subclass of Transform is the second method, and this has the advantage of being able to define attributes with the instantiated objects. Let's define a class to sum in a chosen dimension, and use it to sum in the 2nd (height) dimension:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SumDimension(Transform):\n",
    "    def __init__(self, dim=1):\n",
    "        self.dim = dim\n",
    "\n",
    "    def __call__(self, inputs):\n",
    "        return inputs.sum(self.dim)\n",
    "\n",
    "\n",
    "trans = Compose([LoadImage(image_only=True), EnsureChannelFirst(), SumDimension(2)])\n",
    "img = trans(filenames[0][\"img\"])\n",
    "plt.imshow(img[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All of these example transforms so far have been deterministic, to define transforms which perform some stochastic operation on input data we want to also inherit from `Randomizable`. This class is used to randomize variables but also distinguish from deterministic transforms. We'll see why this is important later in caching data loaders. \n",
    "\n",
    "In this class we have a `numpy.random.RandomState` object to provide stochastic values. This can be replaced using `Randomizable.set_random_state()` to control the randomization process. The `randomize()` method is responsible for determining if the random operation is to be performed based on the `prob` probability member, then creates the random noise array if so. This functionality is in this method so that it can be called by `Compose` or other external controllers. \n",
    "\n",
    "For now lets define a simple transform to add noise. \n",
    "\n",
    "**Run this cell a few times to see the random transform being applied 50% of the time.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RandAdditiveNoise(Randomizable, Transform):\n",
    "    def __init__(self, prob: float = 0.5, max_add: float = 1.0) -> None:\n",
    "        self.prob = np.clip(prob, 0.0, 1.0)\n",
    "        self.max_add = max_add\n",
    "        self._noise = 0\n",
    "\n",
    "    def randomize(self, data: torch.tensor) -> None:\n",
    "        self._noise = 0\n",
    "\n",
    "        if self.R.random() < self.prob:\n",
    "            noise_array = self.R.rand(*data.shape[1:])[None]\n",
    "            self._noise = noise_array * self.max_add\n",
    "\n",
    "    def add_noise(self, img: torch.tensor) -> torch.tensor:\n",
    "        return img + torch.as_tensor(self._noise).to(img)\n",
    "\n",
    "    def __call__(self, img: torch.tensor) -> torch.tensor:\n",
    "        self.randomize(img)\n",
    "        return self.add_noise(img)\n",
    "\n",
    "\n",
    "trans = Compose([LoadImage(image_only=True), EnsureChannelFirst(), RandAdditiveNoise()])\n",
    "img = trans(filenames[0][\"img\"])\n",
    "plt.imshow(img[0, 64])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **4. What are dictionary transforms?**\n",
    "\n",
    "So far we have seen transforms which are applied to individual Numpy arrays, however for most training schemes a pipeline with multiple values is needed. To address this MONAI includes transforms for operating on dictionaries of arrays, one for each equivalent array transform. These can be applied to named values in an input dictionary while leaving unnamed values untouched, for example adding noise to an image while leaving the associated label image untouched.\n",
    "\n",
    "Earlier in the notebook we imported the dictionary equivalent transforms which have a `d` appended to their names, we'll use those transforms in this section.  The `keys` argument in `LoadNiftid` is used to state which keys contain paths to Nifti files, all other values in the input dictionary will be retained. With this set we can look at the keys returned when calling the transform:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trans = LoadImaged(keys=fn_keys)\n",
    "data = trans(filenames[0])\n",
    "print(list(data.keys()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`Lambdad` applies the given callable to each array named by `keys` separately. We can use this to define transforms operating on different named values in the dictionary at different points in the sequence:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sum_width(img):\n",
    "    return img.sum(1)\n",
    "\n",
    "\n",
    "def max_width(img):\n",
    "    return img.max(1)\n",
    "\n",
    "\n",
    "trans = Compose(\n",
    "    [LoadImaged(fn_keys), EnsureChannelFirstd(fn_keys), Lambdad((\"img\",), sum_width), Lambdad((\"seg\",), max_width)]\n",
    ")\n",
    "\n",
    "imgd = trans(filenames[0])\n",
    "img = imgd[\"img\"]\n",
    "seg = imgd[\"seg\"]\n",
    "\n",
    "plt.imshow(np.hstack((img[0] * 5 / img.max(), np.squeeze(seg[0]))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above applies one operation to one member of the dictionary and different operation to another. A reasonable re-implementation of this in one transform would require retaining the names of which members to apply which transform to and applying the operations in one method. Adapting array-based transforms to operate over dictionaries is relatively straight-forward.\n",
    "\n",
    "**Run this cell a few times to see the random transform being applied 50% of the time.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RandAdditiveNoised(Randomizable, MapTransform):\n",
    "    def __init__(self, keys: KeysCollection, prob: float = 0.5, max_add: float = 1.0) -> None:\n",
    "        super(Randomizable, self).__init__(keys)\n",
    "        self.transform = RandAdditiveNoise(prob, max_add)\n",
    "\n",
    "    def set_random_state(\n",
    "        self, seed: Optional[int] = None, state: Optional[np.random.RandomState] = None\n",
    "    ) -> \"RandAdditiveNoised\":\n",
    "        self.transform.set_random_state(seed, state)\n",
    "        super().set_random_state(seed, state)\n",
    "        return self\n",
    "\n",
    "    def randomize(self, data: Optional[Any] = None) -> None:\n",
    "        self.transform.randomize(data)\n",
    "\n",
    "    def __call__(self, data: Mapping[Hashable, torch.tensor]) -> Mapping[Hashable, torch.tensor]:\n",
    "        self.randomize(data[first(self.keys)])\n",
    "\n",
    "        d = dict(data)\n",
    "        for key in self.keys:\n",
    "            d[key] = self.transform.add_noise(d[key])\n",
    "        return d\n",
    "\n",
    "\n",
    "trans = Compose([LoadImaged(fn_keys), EnsureChannelFirstd(fn_keys), RandAdditiveNoised((\"img\",))])\n",
    "img = trans(filenames[0])\n",
    "\n",
    "# We're adding random noise to the image, not the segmentation\n",
    "plt.imshow(np.hstack([img[\"img\"][0, 50], img[\"seg\"][0, 50]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Method overrides are provided to delegate behaviour to an internal `RandAdditiveNoise` instance:\n",
    "* `set_random_state` sets the state of the delegate and the current object\n",
    "* `randomize` delegates the randomization to the `RandAdditiveNoise` instance\n",
    "* `__call__` causes the delegate to randomize then applies the transform to each named member of the dictionary. The delegate transform is randomized only once, this ensures the same random field is added to each named member of the dictionary, a slightly different implementation adding a per-key random field would be needed if this were the desired behaviour. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **5. How do I create a basic MONAI dataset with transforms?**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we've taken a look at transform, let's take a look at datasets. With a data source and transforms defined we can now create a dataset object. The base class for MONAI is `Dataset`, created here to load the image Nifti files only.\n",
    "\n",
    "`Dataset` inherits from the Pytorch class of that name and adds only the ability to apply the given transform to selected items. If you're familiar with the class from Pytorch this will work the same way. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images = [fn[\"img\"] for fn in filenames]\n",
    "\n",
    "transform = Compose([LoadImage(image_only=True), EnsureChannelFirst()])\n",
    "ds = Dataset(images, transform)\n",
    "img_tensor = ds[0]\n",
    "print(img_tensor.shape, img_tensor.get_device())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "MONAI provides the `ArrayDataset` for supervised training applications specifically. It can accept data arrays for images separate from those for segmentations or labels with their own separate transforms. Here we will again separate out the image and segmentation filenames to demonstrate this usage:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images = [fn[\"img\"] for fn in filenames]\n",
    "segs = [fn[\"seg\"] for fn in filenames]\n",
    "\n",
    "img_transform = Compose(\n",
    "    [\n",
    "        LoadImage(image_only=True),\n",
    "        EnsureChannelFirst(),\n",
    "        RandSpatialCrop((128, 128, 128), random_size=False),\n",
    "        RandAdditiveNoise(),\n",
    "    ]\n",
    ")\n",
    "seg_transform = Compose(\n",
    "    [LoadImage(image_only=True), EnsureChannelFirst(), RandSpatialCrop((128, 128, 128), random_size=False)]\n",
    ")\n",
    "\n",
    "ds = ArrayDataset(images, img_transform, segs, seg_transform)\n",
    "im, seg = ds[0]\n",
    "plt.imshow(np.hstack([im.numpy()[0, 48], seg.numpy()[0, 48]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Instead of returning a single image, accessing a member of the dataset produces a pair containing the image and segmentation after being pass through their respective transforms. One important aspect of this class is that the random state of each transform (`Compose` in this case) is set to that of the dataset before being applied. This ensures the same random operations are applied to each output, which is why the `RandSpatialCrop` operation chooses the same crop window for the image as well as the segmentation. By having separate transforms one can apply operations to images and not to segmentations (or vice versa), being careful that these unshared operations come after the shared ones.\n",
    "\n",
    "Alternatively, `Dataset` can be used with dictionary-based transforms to construct a result mapping. For training applications beyond simple input/ground-truth pairs like the above this would be more suitable:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trans = Compose(\n",
    "    [\n",
    "        LoadImaged(fn_keys),\n",
    "        EnsureChannelFirstd(fn_keys),\n",
    "        RandAdditiveNoised((\"img\",)),\n",
    "        RandSpatialCropd(fn_keys, (128, 128, 128), random_size=False),\n",
    "    ]\n",
    ")\n",
    "\n",
    "ds = Dataset(filenames, trans)\n",
    "item = ds[0]\n",
    "im, seg = item[\"img\"], item[\"seg\"]\n",
    "plt.imshow(np.hstack([im.numpy()[0, 48], seg.numpy()[0, 48]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the dataset defined, we can now create the dataloader to create data batches. This inherits directly from Pytorch's `DataLoader` class with a few changes to the default constructor arguments. MONAI functionality should be compatible with the PyTorch DataLoader, but it was subclasses to include additional functionality that we consider key and which cannot be realized with the standard DataLoader class.\n",
    "\n",
    "The `DataLoader` will use five worker processes to load the actual data. MONAI provides a number of `Dataset` subclasses to improve the efficiency of this process. These and other features will be covered in subsequent labs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loader = DataLoader(ds, batch_size=5, num_workers=5)\n",
    "batch = first(loader)\n",
    "print(list(batch.keys()), batch[\"img\"].shape)\n",
    "\n",
    "f, ax = plt.subplots(2, 1, figsize=(8, 4))\n",
    "ax[0].imshow(np.hstack(batch[\"img\"][:, 0, 64]))\n",
    "ax[1].imshow(np.hstack(batch[\"seg\"][:, 0, 64]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **Summary**\n",
    "\n",
    "We've covered MONAI Transforms.  Some key highlights are:\n",
    "\n",
    "- There is a long list of medical specific transforms available in MONAI\n",
    "- There are array and dictionary versions of transforms.\n",
    "- You can create a simple callable lambda function or create a class based on transform to create your own custom tranform\n",
    "- You can create a MONAI dataset and directly pass a compose tranform chain to it"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **Assignment 1**\n",
    "\n",
    "Now you get to practice creating a transform of your own and integrating it into a compose array.\n",
    "\n",
    "Create a transform which selects a patch from an input image of a given size and randomly chooses to transpose this data along randomly chosen axes. For example, this could produce an image where a patch in the center has been flipped in the horizontal direction while the rest of the image is left unchanged.\n",
    "\n",
    "Assumptions:\n",
    "* Implement this as an array randomized transform\n",
    "* Use a default probability ratio of 40% chance to crop\n",
    "* Use one of the test 3D images from above to validate\n",
    "* Default crop size is 150 but allow it to be variable\n",
    "* Don't worry about edge cases\n",
    "\n",
    "Below is the initial class definition you'll need to implement.  We also include a compose chain and image visualization call that should work once your function is fully implemented."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RandFlipRect(Randomizable, Transform):\n",
    "    \"\"\"\n",
    "    Define a transform to randomly transpose a patch of the input image along randomly chosen spatial dimensions. This \n",
    "    uses arguments `prob` for the probability the transform is applied, and `size` giving the size of the cube to flip.\n",
    "    \"\"\"\n",
    "    def __init__(self, prob: float = 0.4, size: int = 100) -> None:\n",
    "        self.prob = np.clip(prob, 0.0, 1.0)\n",
    "        self.size = size\n",
    "        \n",
    "    def __call__(self, data: torch.tensor) -> torch.tensor:\n",
    "        pass  # replace with your code\n",
    "    \n",
    "trans = Compose([LoadImage(image_only=True), EnsureChannelFirst(), RandFlipRect()])\n",
    "img = trans(filenames[0][\"img\"])\n",
    "plt.imshow(img[0, img.shape[1]//2])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **Assignment 2**\n",
    "\n",
    "In the cell below define a dictionary-based transform pipeline which will take input images and segmentations and ensure:\n",
    "* both image and segmentation arrays have a channel dimension\n",
    "* each image/segmentation pair is 200x200x200 pixels in shape\n",
    "* the value range of each image is between 0 and 1\n",
    "* the dtype of the image is float32 and that of the segmentation is int32\n",
    "* both are produced in the end as tensors\n",
    "\n",
    "All this can be done with built-in transforms so you shouldn't need a lambda transform. Once your pipeline works, use it to construct and dataset and then a data loader. Iterate over every value in the data loader and check that the image and segmentation members of the produced dictionaries meet these requirements."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "trans = Compose(\n",
    "    [\n",
    "        # ???\n",
    "    ]\n",
    ")\n",
    "\n",
    "data = trans(filenames[0])\n",
    "img = data[\"img\"]\n",
    "seg = data[\"seg\"]\n",
    "\n",
    "print(img.shape)  # should be (1, 200, 200, 200)\n",
    "print(img.dtype, seg.dtype)  # should be float32  int32\n",
    "print(img.min(), img.max())  # should be 0.0 1.0\n",
    "print(torch.unique(seg))  # should be [0 1 2 3 4 5]\n",
    "\n",
    "plt.imshow(img[0, img.shape[1] // 2])\n",
    "\n",
    "# construct dataset and dataloder here and iterate over batches, batch_size of 1 should be used\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **Next Steps**\n",
    "\n",
    "In the next notebook, we cover more advanced uses of dataset and dataset caching.\n",
    "\n",
    "You can find more information about everything covered here on our [MONAI Documentation Page](https://docs.monai.io/).  \n",
    "\n",
    "If you're looking for more examples and tutorials, we have a repo dedicated just to that!  You can find it on our [GitHub Organization Page](https://github.com/Project-MONAI/tutorials). "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}