2-Inference.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using MONAI to unlock clinically valuable insights from Digital Pathology"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 2 - Nuclei segmentation and classification with MONAI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"images/monai.png\" alt=\"MONAI\" style=\"width: 500px;\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the first section, you saw how we can speed up the loading and decoding of large images by using a high-performance image loader such as cuCIM. You also saw how we can use multi-threading or DASK to reduce the latency of loading a large image. So long as the image format and loader supports loading regions of interest then we can use different processes or threads to simultaneously get different parts of the image into memory. \n",
    "\n",
    "Often, the image loading can be the bottleneck that slows the whole processing pipeline down. Of course it is often not just the loading that we need to do. There may be a need to do some pre-processing on the image and, for digital pathology, we may need to threshold each region to ensure that we are not wasting time processing empty or background regions of the Whole Slide. There may also be some image transformation or augmentation to do. All of these operations can become the part that slows everything else down and results in under-utilised GPUs, if not dealt with efficiently.\n",
    "\n",
    "We are going to work through an example in which we use the HoVerNet [1] network and post-processing pipeline to detect, localise and classify nuclei and then, in the following notebook, analyse the output.\n",
    "\n",
    "HoVerNet is now one of the MONAI supported networks and can be found [here](https://github.com/Project-MONAI/tutorials/blob/main/pathology/hovernet/hovernet_torch.ipynb)\n",
    "\n",
    "![image](images/hovernet2.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[1] Simon Graham, Quoc Dang Vu, Shan E Ahmed Raza, Ayesha Azam, Yee Wah Tsang, Jin Tae Kwak, Nasir Rajpoot, Hover-Net: Simultaneous segmentation and classification of nuclei in multi-tissue histology images, Medical Image Analysis, 2019 https://doi.org/10.1016/j.media.2019.101563"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We start off by importing the libraries we need Python to be able to use"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright 2020 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.\n",
    "\n",
    "import sys\n",
    "\n",
    "# General Python libraries\n",
    "from functools import partial\n",
    "from matplotlib import pyplot as plt\n",
    "import logging\n",
    "import os\n",
    "import time\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "# torch-related classes\n",
    "import torch\n",
    "import numpy as np\n",
    "\n",
    "# Monai classes\n",
    "from monai.config import print_config\n",
    "from monai.data import (\n",
    "    DataLoader, \n",
    "    MaskedPatchWSIDataset, \n",
    "    decollate_batch, \n",
    "    PatchWSIDataset, \n",
    "    IterableDataset,\n",
    ")\n",
    "\n",
    "from monai.networks.nets import HoVerNet\n",
    "from monai.engines import IterationEvents, SupervisedEvaluator\n",
    "from monai.inferers import SimpleInferer\n",
    "\n",
    "# Pathology-specific transforms\n",
    "from monai.apps.pathology.transforms import (\n",
    "    GenerateWatershedMaskd,\n",
    "    GenerateInstanceBorderd,\n",
    "    GenerateDistanceMapd,\n",
    "    GenerateWatershedMarkersd,\n",
    "    GenerateInstanceContour,\n",
    "    GenerateInstanceCentroid,\n",
    "    GenerateInstanceType,\n",
    "    HoVerNetInstanceMapPostProcessingd, \n",
    "    HoVerNetNuclearTypePostProcessingd,\n",
    ")\n",
    "\n",
    "from monai.apps.pathology.transforms.post.dictionary import (\n",
    "    HoVerNetNuclearTypePostProcessingd, \n",
    "    Watershedd,\n",
    ")\n",
    "\n",
    "# Generic Transforms\n",
    "from monai.transforms import (\n",
    "    Compose,\n",
    "    ScaleIntensityRanged,\n",
    "    CastToTyped,\n",
    "    Lambdad,\n",
    "    LoadImage,\n",
    "    LoadImaged,\n",
    "    EnsureChannelFirst,\n",
    "    EnsureChannelFirstd,\n",
    "    ComputeHoVerMapsd,\n",
    "    BoundingRect,\n",
    "    ThresholdIntensity,\n",
    "    NormalizeIntensityd,\n",
    "    apply_transform,\n",
    ")\n",
    "\n",
    "# Event handlers\n",
    "from monai.handlers import EarlyStopHandler\n",
    "from monai.utils import convert_to_tensor, first, HoVerNetBranch\n",
    "\n",
    "print_config()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output of print_config provides us with a useful summary of the installed MONAI components and dependencies. If you encounter any issues when developing with MONAI, it can be useful to include this output in any issues you raise in the GitHub repo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As before, we we will be working with the same image from the TCGA archive. However, rather than loading the image directly, we are going to look at how MONAI abstracts the loading of images away. There are many different possibilities for loading data in MONAI because it is a very flexible API. \n",
    "Sometimes, it makes sense to just point a MONAI engine at a data source, such as a folder of images and let it run automatically, using a pre-configured pipeline. At other times you may have different requirements that necessitate some customisation at some level. We will take a look at a few of these options."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First off, let's look at some of the MONAI components that comprise a typical pipeline. \n",
    "We need to define a source of some images. We formulate this as a list of dictionaries to fit with the way that Monai operates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_list = [\n",
    "    {\"image\": \"data/images/im_test1.nii.gz\"},\n",
    "    {\"image\": \"data/images/im_test2.nii.gz\"},\n",
    "    {\"image\": \"data/images/im_test3.nii.gz\"},\n",
    "    {\"image\": \"data/images/im_test4.nii.gz\"},\n",
    "    {\"image\": \"data/images/im_test5.nii.gz\"},\n",
    "    {\"image\": \"data/images/im_test6.nii.gz\"},\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, we now have a list containing 6 dictionaries, each containing an \"image\" key and a filename as the value. We can use a Monai Transform to turn this list of images into a list of actual images. The Transform that we need to do this is LoadImage or LoadImaged. The first expects a list of file names whereas the second expects a list of dictionaries in which the key can be used for filtering.\n",
    "It is common to combine Transforms into a pipeline and the _Compose_ function allows us to do that. In this initial case, there is only one item in the Compose pipeline but we could easily chain more together."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trans = Compose([LoadImage(image_only=True)])\n",
    "img = trans(data_list[0][\"image\"])\n",
    "print(type(img), img.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, we have a 500 by 500 pixel image with 3 channels. Let's plot that out next"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Use Matplotlib to display the thumbnail view of the image\n",
    "plt.figure(figsize=(5,5))\n",
    "plt.imshow(np.array(img).astype(int))\n",
    "plt.title('tcga1.svs')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, one thing that needs to be accounted for is the channel ordering. Monai, like PyTorch, on which it is based, generally uses the convention of channels being the first dimension of an image. To make sure that this is the case we can use another transform - EnsureChannelFirst. For the case in which an image has only a single channel, there may not be a separate dimension (e.g. shape = [100, 100] rather than [1, 100, 100]) but this method will create one. Otherwise it will reorder the dimensions. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# channel_dim=-1 tells the Transform that the last dimension is the current channel dimension\n",
    "trans = Compose([LoadImage(image_only=True), EnsureChannelFirst(channel_dim=-1)])\n",
    "img = trans(data_list[0][\"image\"])\n",
    "print(type(img), img.shape)\n",
    "del(img)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, this has now provided the tensor in the correct format. Note that matplotlib uses the channel-last convention, so if you want to plot the image, you will need to reorder the channels again before displaying. The other thing to note is that the type of the tensor is reported as a Monai MetaTensor. You can find out more about this type in the Monai Documentation (https://docs.monai.io/en/stable/data.html#metatensor) but it is a useful feature of Monai that allows you to examine and add to the data associated with an image tensor e.g. the name of source image file, its dimensions or, in the case of Whole Slide Image processing, perhaps the coordinates of the current tile within the WSI."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To make this useful, we would most likely want to integrate a Transform into a data loading workflow. So far we used Array Tranforms, which operate on arrays of inputs. Now we will switch to Dictionary Tranforms, which are able to use the keys to filter the inputs. This can be useful, especially when training, because the filters can be used to deal with labels and input data separately. We are only doing inference in this case so we don't have labels but it can still be useful.\n",
    "\n",
    "Firstly, we will redefine the Transforms to use the Dictionary equivalent functions. They always include a 'd' suffix e.g. LoadImage becomes LoadImaged. This is also a feature of PyTorch, which also has the concept of Tranforms, which can be used with Monai Transforms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trans = Compose([LoadImaged(keys=\"image\"), EnsureChannelFirstd(keys=\"image\", channel_dim=-1)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This has defined the Transform. Next we need to use this Transform in some sort of data flow. We will use a MONAI Dataset to do this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an iterable from the image list\n",
    "data_iterator = iter(data_list)\n",
    "\n",
    "# Create a dataset with the data_iterator and the transforms that we defined\n",
    "dataset = IterableDataset(data=data_iterator, transform=trans)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can create a DataLoader to create some batches. This should yield 2 batches of 3 images, each of dimension (3, 500, 500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataloader = DataLoader(dataset=dataset, batch_size=3, num_workers=1)\n",
    "for d in dataloader:\n",
    "    print(d[\"image\"].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we wanted to create a Dataset rather than an IterableDataset then, rather than using the iter() function, we'd need to wrap the data_list in a class that provides the necessary methods e.g.:\n",
    "\n",
    "    class MyIterator:\n",
    "        def __init__(self, data):\n",
    "            self.data = data\n",
    "\n",
    "        def __iter__(self):\n",
    "            return iter(self.data)\n",
    "\n",
    "        def __getitem__(self, index):\n",
    "            return self.data[index]\n",
    "This would allow us to do something like:\n",
    "\n",
    "    data_iterator = MyIterator(data_list)\n",
    "    \n",
    "    dataset = Dataset(data=data_iterator, transform=trans)\n",
    "    for _ in dataset:\n",
    "        pass\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also add other Tranforms, if needed. See if you can add a Transform to Normalize the Intensity of the image (use the [documentation](https://docs.monai.io/en/stable/transforms.html#intensity) to find the right Tranform) ([solution](solutions/normalize.py))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_list = [\n",
    "    {\"image\": \"data/images/im_test1.nii.gz\"},\n",
    "]\n",
    "\n",
    "# TODO Add a new transform to the Compose input list\n",
    "trans = Compose([LoadImaged(keys=\"image\"), EnsureChannelFirstd(keys=\"image\", channel_dim=-1),])\n",
    "\n",
    "#Create the data iterator and dataset\n",
    "data_iterator = iter(data_list)\n",
    "dataset = IterableDataset(data=data_iterator, transform=trans)\n",
    "\n",
    "# Load the data\n",
    "dataloader = DataLoader(dataset=dataset, batch_size=3, num_workers=1)\n",
    "\n",
    "# retrieve the first batch\n",
    "d = first(dataloader)\n",
    "\n",
    "img = np.array(d[\"image\"][0])\n",
    "# TODO - Remember that MatplotLib expects channels last - use numpy to do this\n",
    "\n",
    "plt.figure(figsize=(5,5))\n",
    "plt.imshow(img)\n",
    "plt.title('Nomalized Image')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For WholeSlide Image Processing, we don't get the patches or tiles as the initial input. Instead it is necessary to slice the large images up into smaller patches or tiles using a Sliding Window approach. MONAI provides a number of tools to help with this process. To show this, we will start by defining a new datalist. This time we will supply the single WSI file, but provide some metadata that the MONAI Dataset can use to only load specified tiles (at a specified reolution, which defaults to full resolution)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate a data list that starts in the middle of the image and\n",
    "# specifies patches of 256 pixel at intervals of 164 (overlapping) within a 2k by 2k region\n",
    "width = 87647\n",
    "height = 52434\n",
    "\n",
    "x = 25000\n",
    "y = 47000\n",
    "\n",
    "data_list = []\n",
    "\n",
    "for i in range(x, x+2000, 164):\n",
    "    for j in range(y, y+2000, 164):\n",
    "        data_list.append({\"image\": \"data/tcga1.svs\",\"location\": [i, j], \"size\": [256,256]})\n",
    "\n",
    "print(data_list[0:10])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we can use Monai's PatchWSIDataset to handle the patch creation. Notice that this uses \"cuCIM\" by default as the image loader and this means that we don't actually need to create a Transform to do the Loading of Images. It also handles the ordering of the channels for us."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = PatchWSIDataset(\n",
    "    data_list,\n",
    "    patch_size=256,\n",
    "    patch_level=0,\n",
    "    include_label=False,\n",
    "    reader=\"cuCIM\",\n",
    "    additional_meta_keys=[\"location\", \"size\"],\n",
    ")\n",
    "print(\"dataset created\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can create the image batches with a DataLoader. The Monai DataLoaders are based on the equivalent PyTorch and are compatible with them. You can use the same features such as the number of concurrent workers to optimise the data loading."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the image batches\n",
    "dataloader = DataLoader(dataset=dataset, batch_size=3, num_workers=2)\n",
    "\n",
    "# retrieve the first batch\n",
    "d = first(dataloader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally we can show the first batch of 3 images. You should see that they are adjacent tiles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "imgs = np.array(d[\"image\"])\n",
    "imgs = np.moveaxis(imgs, 1, -1)\n",
    "\n",
    "print(imgs[0][46:-46,46:-46].shape)\n",
    "\n",
    "fig, ax = plt.subplots(1, 3, figsize=(10, 10))\n",
    "ax[0].imshow(imgs[0][46:-46,46:-46].astype(np.uint8))\n",
    "ax[1].imshow(imgs[1][46:-46,46:-46].astype(np.uint8))\n",
    "ax[2].imshow(imgs[2][46:-46,46:-46].astype(np.uint8))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, we have shown a few ways that you can load and process batches of images, but there is plenty more to play around with if you have specific needs that are not met by what we have seen already. \n",
    "\n",
    "We will now move on to doing some inference using HoVerNet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "HoVerNet can produces 3 branches of output which contain a nucleus probability map, a class probability map and a 2 channel Horizontal and Vertical distance from the centroid map. These outputs require some post-processing to convert the raw predictions into cleanly segmented and typed nuclei. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![image](images/Post_Processing_Workflow.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Monai contains all of the necessary capabilities to be able to do this post-processing and it is accomplished with - guess what - *Transforms*!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "log_dir = \"outputs\"\n",
    "\n",
    "torch.cuda.set_device(0)\n",
    "device = torch.device(\"cuda\")\n",
    "\n",
    "# Preprocessing transforms\n",
    "pre_transforms = Compose(\n",
    "    [\n",
    "        CastToTyped(keys=[\"image\"], dtype=torch.float32),\n",
    "        ScaleIntensityRanged(keys=[\"image\"], a_min=0.0, a_max=255.0, b_min=0.0, b_max=1.0, clip=True),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# Dataset of patches\n",
    "dataset = PatchWSIDataset(\n",
    "    data_list,\n",
    "    patch_size=256,\n",
    "    patch_level=0,\n",
    "    include_label=False,\n",
    "    transform=pre_transforms,\n",
    "    reader=\"cuCIM\",\n",
    "    additional_meta_keys=[\"location\"],\n",
    ")\n",
    "print(\"dataset created\")\n",
    "\n",
    "# Dataloader\n",
    "data_loader = DataLoader(dataset, \n",
    "                        num_workers=2, \n",
    "                        batch_size=8, \n",
    "                        pin_memory=True,)\n",
    "print(\"dataloader created\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To sanity-check the DataLoader, we can call the first() method, passing the DataLoader as the parameter. This will return the first batch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "first_sample = first(data_loader)\n",
    "\n",
    "print(\"image: \")\n",
    "print(\"    shape\", first_sample[\"image\"].shape)\n",
    "print(\"    type: \", type(first_sample[\"image\"]))\n",
    "print(\"    dtype: \", first_sample[\"image\"].dtype)\n",
    "print(\"    location: \", first_sample[\"image\"][0].meta[\"location\"])\n",
    "print(f\"number of batches: {len(data_loader)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This should have produced a batch of 8 images, each with 3 channels and height and width of 256.\n",
    "\n",
    "Now we can create the HoVerNet model, specifying the parameters (those provided are generally appropriate)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create model \n",
    "model = HoVerNet(\n",
    "    mode=\"fast\", \n",
    "    in_channels=3, \n",
    "    out_classes=5, \n",
    "    act=(\"relu\", {\"inplace\": True}), \n",
    "    norm=\"batch\", \n",
    "    dropout_prob=0.0,\n",
    ")\n",
    "\n",
    "# Specifies to use the first available GPU\n",
    "device = torch.device(\"cuda:0\")\n",
    "\n",
    "#Load the pre-trained weights\n",
    "model.load_state_dict(torch.load(\"data/fast.pt\")['model'])\n",
    "model = model.to(device)\n",
    "model.eval()\n",
    "\n",
    "inferer = SimpleInferer()\n",
    "\n",
    "dataloader = DataLoader(dataset=dataset, batch_size=4, num_workers=4)\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    for d in dataloader:\n",
    "        pred = inferer(inputs=d[\"image\"].to(device), network=model)\n",
    "        imgs = np.array(pred[\"nucleus_prediction\"][:4].cpu())\n",
    "        break\n",
    "        \n",
    "print(imgs[0,:,:,0].shape)\n",
    "fig, ax = plt.subplots(1, 4, figsize=(10, 10))\n",
    "ax[0].imshow(imgs[0,1,:,:])\n",
    "ax[1].imshow(imgs[1,1,:,:])\n",
    "ax[2].imshow(imgs[2,1,:,:])\n",
    "ax[3].imshow(imgs[3,1,:,:])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "So, we now have some inference results from HoVerNet, but this is without any post-processing. Pre-processing was already invoked when the data is loaded, to normalize the intensity. \n",
    "\n",
    "One important thing to notice is that the outpput size from HoVerNet is smaller (164 x 164) than the input size (256 x 256). This is because the series of convolutions reduce the input size. To compensate for this, we need to ingest tiles with overlapping borders (46 pixels on each side). The benefit of this approach is that we get much less pronounced tile border artefacts.\n",
    "\n",
    "Let's now define some post-processing to clean up the inference predictions. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Setup a post-processing pipeline to convert the predictions into the desired outputs\n",
    "Here we use several transforms to turn pixel-level predictions into maps, contours and images that can be saved to disk:\n",
    "1. `GenerateWatershedMaskd` Creates a binary mask within which to compute watershed\n",
    "2. `GenerateInstanceBorderd` Generates an instance border using a horizontal and vertical (hover) distance map\n",
    "3. `GenerateDistanceMapd` Within a segmentation region, computes the distance from the centre for horizontal and vertical axes\n",
    "4. `GenerateWatershedMarkersd` Generate markers to be used in `Watershed` algorithm\n",
    "6. `Watershedd` Uses the watershed algorithm to link pixels to specific object instances\n",
    "\n",
    "We could implement these directly or, for an easier life, we can just use the predefined HoVerNet PostProcessing pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Postprocessing transforms\n",
    "post_transforms = Compose(\n",
    "    [\n",
    "        HoVerNetInstanceMapPostProcessingd(sobel_kernel_size=21, marker_threshold=0.4, marker_radius=2),\n",
    "        HoVerNetNuclearTypePostProcessingd(),\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to create a dataloader, set the model to inference mode and then we can run the inference outputs through the post-processing pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataloader = DataLoader(dataset=dataset, batch_size=4, num_workers=4)\n",
    "model.eval()\n",
    "out=[]\n",
    "\n",
    "with torch.no_grad():\n",
    "    for d in dataloader:\n",
    "        pred = inferer(inputs=d[\"image\"].to(device), network=model)\n",
    "\n",
    "        nu = np.array(pred[\"nucleus_prediction\"].cpu())\n",
    "        hv = np.array(pred[\"horizontal_vertical\"].cpu())\n",
    "        tp = np.array(pred[\"type_prediction\"].cpu())\n",
    "        \n",
    "        for i in range(len(nu)):\n",
    "            inputs =  {\"nucleus_prediction\": nu[i], \"horizontal_vertical\": hv[i], \"type_prediction\": tp[i]}\n",
    "            out.append(post_transforms(inputs))  \n",
    "                       \n",
    "        break\n",
    "        \n",
    "# Plot out the cleaned up instance map\n",
    "fig, ax = plt.subplots(1, 4, figsize=(10, 10))\n",
    "ax[0].imshow(out[0][\"instance_map\"].squeeze())\n",
    "ax[1].imshow(out[1][\"instance_map\"].squeeze())\n",
    "ax[2].imshow(out[2][\"instance_map\"].squeeze())\n",
    "ax[3].imshow(out[3][\"instance_map\"].squeeze())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot out the type map (nucleus sub-types)\n",
    "fig, ax = plt.subplots(1, 4, figsize=(10, 10))\n",
    "ax[0].imshow(out[0][\"type_map\"].squeeze())\n",
    "ax[1].imshow(out[1][\"type_map\"].squeeze())\n",
    "ax[2].imshow(out[2][\"type_map\"].squeeze())\n",
    "ax[3].imshow(out[3][\"type_map\"].squeeze())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we look at the dictionary in each output row, we can see that the post-processing has created some additional information (the last three keys)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out[0].keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we look in the instance_info item, we can find the contours and centroids of each nuclei instance found."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out[0][\"instance_info\"].keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This reveals that for the first patch, there are 6 nuclei instance with 1-based numeric keys. Looking at the first of these reveals the metadata that the post-processing has computed for each nuclei instance. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out[0][\"instance_info\"][1].keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the next task, you are going to create a function that can write the centroid coordinates and type of each nucleus into an array which we can then use to map the various nuclei types found in the 2k x 2k region we just ran inference on. To do this you will need to examine the data that we looked at above and find the relevant items to write to disk. ([solution](solutions/get_centroids.py))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The offset of the current patch - use it to position the\n",
    "# centroids relative to the Region of Interest offset\n",
    "current_tile_offset = ()\n",
    "\n",
    "# Region of Interest Offset\n",
    "roi_y = 47000\n",
    "roi_x = 25000\n",
    "\n",
    "# create a function that will be added\n",
    "# to the post-processing Transforms\n",
    "def get_centroids(inst_info):\n",
    "    \n",
    "    centroids=[]\n",
    "    \n",
    "    # TODO - get controid info\n",
    "        \n",
    "    return centroids\n",
    "        \n",
    "# Postprocessing transforms\n",
    "post_transform_with_centroids = Compose(\n",
    "    [\n",
    "        HoVerNetInstanceMapPostProcessingd(sobel_kernel_size=21, marker_threshold=0.4, marker_radius=2),\n",
    "        HoVerNetNuclearTypePostProcessingd(),\n",
    "        # TODO - create Lambda function to call get_centroids\n",
    "    ]\n",
    ")\n",
    "\n",
    "dataloader = DataLoader(dataset=dataset, batch_size=4, num_workers=4)\n",
    "model.eval()\n",
    "out=[]\n",
    "\n",
    "with torch.no_grad():\n",
    "    \n",
    "    centroids = []\n",
    "    for d in dataloader:\n",
    "        # TODO get the offset for each image in the batch\n",
    "        offsets = ...\n",
    "        pred = inferer(inputs=d[\"image\"].to(device), network=model)\n",
    "        nu = np.array(pred[\"nucleus_prediction\"].cpu())\n",
    "        hv = np.array(pred[\"horizontal_vertical\"].cpu())\n",
    "        tp = np.array(pred[\"type_prediction\"].cpu())\n",
    "        \n",
    "        for i in range(len(nu)):\n",
    "            # TODO set the current offset\n",
    "            current_tile_offset = ...\n",
    "            raw_results =  {\"nucleus_prediction\": nu[i], \"horizontal_vertical\": hv[i], \"type_prediction\": tp[i]}\n",
    "            # TODO apply the postprocessing transform to the raw_results\n",
    "            # TODO Add the output to a list of centroids in dictionarry format \n",
    "\n",
    "print(\"Completed\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(centroids[0:10])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This next code cell should plot out the results. If all is working correctly, you should see a distribution of nuclei that corresponds to the input image region"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import ListedColormap\n",
    "from cucim import CuImage\n",
    "\n",
    "# load the image header\n",
    "wsi = CuImage(\"data/tcga1.svs\")\n",
    "\n",
    "# Get the resolution meta data\n",
    "sizes=wsi.metadata[\"cucim\"][\"resolutions\"]\n",
    "\n",
    "# Load the image data at this resolution\n",
    "wsi_thumb = wsi.read_region(location=(47000, 25000), size=(2000,2000), level=0)\n",
    "\n",
    "centres = np.zeros((len(centroids),3),dtype=int)\n",
    "label = [\"\"] * len(centroids)\n",
    "\n",
    "for i, centre in enumerate(centroids):\n",
    "    #invert row coordinate and swap x/y to match coordinates used in images\n",
    "    centres[i] = [centre[\"y\"], 2000-centre[\"x\"], centre[\"type\"]-1]  \n",
    "\n",
    "# plot\n",
    "cmap = ListedColormap([\"blue\", \"gold\", \"lawngreen\", \"red\"])\n",
    "fig, ax = plt.subplots(1,2,figsize = (10,5))\n",
    "\n",
    "ax[0].scatter(centres[:,0], centres[:,1], c=centres[:,2], cmap=cmap)\n",
    "ax[0].set(xlim=(0, 2000), xticks=np.arange(0, 2000, 500),\n",
    "       ylim=(0, 2000), yticks=np.arange(0, 2000, 500))\n",
    "\n",
    "ax[1].imshow(wsi_thumb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This sort of output, showing cell types and locations can provide useful clinical information (e.g. counts and densities of mitotic cells) and, as we will see in the next notebook, can also be used to provide raw data for other types of analysis."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are other abstractions that make using HoverNet even easier, such as using Ignite-based Evaluators, which include some nice features, such as progress bars and event-handling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def myfunc(engine, engine_state_batch):\n",
    "    print(\"Iteration Update Event Fired!\")\n",
    "    \n",
    "with tqdm(total=len(data_loader)) as pbar:\n",
    "\n",
    "    # Class used for event handling\n",
    "    class TestEvalIterEvents:\n",
    "        def attach(self, engine):\n",
    "            engine.add_event_handler(IterationEvents.FORWARD_COMPLETED, self._forward_completed)\n",
    "\n",
    "        def _forward_completed(self, engine):\n",
    "            pbar.update(engine.state.iteration)\n",
    "\n",
    "    # Define some Handlers\n",
    "    inference_handlers = [\n",
    "        TestEvalIterEvents(),\n",
    "    ]\n",
    "\n",
    "    model = model.to(device)\n",
    "\n",
    "    # Use an Ignite-based Evaluator\n",
    "    inference = SupervisedEvaluator(\n",
    "        device=device,\n",
    "        val_data_loader=data_loader,\n",
    "        network=model,\n",
    "        iteration_update=myfunc,\n",
    "        val_handlers=inference_handlers,\n",
    "        amp=True,\n",
    "    )\n",
    "    \n",
    "    # This event can be used to stop iteration during training or evaluation\n",
    "    EarlyStopHandler(\n",
    "        patience=20, score_function=lambda x: 1.0, epoch_level=False, trainer=inference\n",
    "    ).attach(inference)\n",
    "\n",
    "    inference.run()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should be able to see how the declared events have printed out messages.\n",
    "\n",
    "MONAI has a lot more to explore and the pathology capabilities are steadily growing, so you are encouraged to look at the documentation and tutorials to learn more.\n",
    "\n",
    "Although mentioned in the first notebook, one of the items that we did not cover in this notebook was the use of thresholding to eliminate regions of the image that contain no tissue. This is something that can be done using features and functions of MONAI, but has not been used in this tutorial. If you want to experiment further, you could adapt the thresholding done in the previous notebook so that it can be used with the inference examples from this notebook."
   ]
  }
 ],
 "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.9.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}