ssn_demo.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "stochastic_segmentation_networks_demo.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fRcVNEp4EdPq",
        "colab_type": "text"
      },
      "source": [
        "# **Stochastic Segmentation Networks demo.**\n",
        "This is an intercative demo for our paper [Stochastic Segmentation Networks: Modelling Spatially Correlated Aleatoric Uncertainty](https://arxiv.org/abs/2006.06015).\n\n",
	"To run this demo:\n",
        "1.   Make sure to run the cell that downloads the images;\n",
        "2.   Please be patient with the image slider; it takes little time to read the necessary files whenever a new image is selected.\n",
        "\n",
        "The effect of the temperature sliders is defined by the sample direction and changes with each new sample.\n",
        "Increasing and decreasing the temperature does not necessarily correlate with increasing or decreasing the volume of a class.\n",
        "Some samples behave nicely, some not so much. In the future, we aim to obtain fine-grained sample control.\n",
	"\n",
        "The temperature slider controls the scaling of the entire covariance matrix, move it to 0 to get the mean.\n",
        "The individual class temperature sliders control the scaling of the part of the covariance matrix relating to only that class. (If the main temperature slider is set to 0 these have no effect)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Y86SDxmHYvFL",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title This cell downloads the images, it may take a few minutes since these are large 3D images. \n",
        "\n",
        "!pip install SimpleITK\n",
        "!pip install gdown\n",
        "!mkdir -p data/BraTS17_2013_10_1\n",
        "!gdown https://drive.google.com/uc?id=1PDQZROumSjbOBIg42yKzWlS0Z46DjOUd -O \"data/BraTS17_2013_10_1/t1ce.nii.gz\"\n",
        "!gdown https://drive.google.com/uc?id=1kFnIAa4lOjRnyKPJ8hqjibf5Q5OWJ4IW -O \"data/BraTS17_2013_10_1/seg.nii.gz\"\n",
        "!gdown https://drive.google.com/uc?id=1ZQ_nJyI7-fVXXpxABKp0opKhlzCKahQ2 -O \"data/BraTS17_2013_10_1/logit_mean.nii.gz\"\n",
        "!gdown https://drive.google.com/uc?id=1mE5UwDMOK0e7nHtt2wn6Dgj0b-0kft2b -O \"data/BraTS17_2013_10_1/cov_factor.nii.gz\"\n",
        "!gdown https://drive.google.com/uc?id=1GmmUfPyiFQVWktlwH5kwHLWSQkT-LtI3 -O \"data/BraTS17_2013_10_1/cov_diag.nii.gz\"\n",
        "!gdown https://drive.google.com/uc?id=1-fAt59jgkAGTyT7gDidIm_tj7fw2I7NN -O \"data/BraTS17_2013_10_1/brainmask.nii.gz\"\n",
        "!mkdir -p data/BraTS17_2013_12_1\n",
        "!gdown https://drive.google.com/uc?id=17BPLnfbZvl7jng4Kbnlk1oTNLWVQTkWG -O \"data/BraTS17_2013_12_1/t1ce.nii.gz\"\n",
        "!gdown https://drive.google.com/uc?id=1G5pMMUfuxOqLPeiVgNoJOmnazEmWau8p -O \"data/BraTS17_2013_12_1/seg.nii.gz\"\n",
        "!gdown https://drive.google.com/uc?id=17FEOGuNFQp5Ypm6GPIEC8n-cpDzYSoCl -O \"data/BraTS17_2013_12_1/logit_mean.nii.gz\"\n",
        "!gdown https://drive.google.com/uc?id=1NFkruCK4nUyveWhUJXJyjv8Z8TQzM0_Z -O \"data/BraTS17_2013_12_1/cov_factor.nii.gz\"\n",
        "!gdown https://drive.google.com/uc?id=1APnglw2vwkp2gbfhKfJysZu08rYViWqS -O \"data/BraTS17_2013_12_1/cov_diag.nii.gz\"\n",
        "!gdown https://drive.google.com/uc?id=1YWcb7g2kjJ1XS0jtrD43glJY3Bj3VVHL -O \"data/BraTS17_2013_12_1/brainmask.nii.gz\"\n",
        "!mkdir -p data/BraTS17_2013_20_1\n",
        "!gdown https://drive.google.com/uc?id=1j0jMay2EC1CHs5hRgFUBNOP2SmaUN0Qm -O \"data/BraTS17_2013_20_1/t1ce.nii.gz\"\n",
        "!gdown https://drive.google.com/uc?id=17M1Gzrkbf3TQ7n6MxPVBJN0g6D8b2Hhr -O \"data/BraTS17_2013_20_1/seg.nii.gz\"\n",
        "!gdown https://drive.google.com/uc?id=1VGJ-uWfRm69YqccHHzKvakkd5XYr9XoA -O \"data/BraTS17_2013_20_1/logit_mean.nii.gz\"\n",
        "!gdown https://drive.google.com/uc?id=1wdI8pGfbHY8lMmdpSkdmnnEkE4DhXXgd -O \"data/BraTS17_2013_20_1/cov_factor.nii.gz\"\n",
        "!gdown https://drive.google.com/uc?id=16exPfMcQ4Cly3zqgANwf21d2bIeiN4aI -O \"data/BraTS17_2013_20_1/cov_diag.nii.gz\"\n",
        "!gdown https://drive.google.com/uc?id=1Aol2mHxc5GZGaOPX3F0TqvnagyJEzj2V -O \"data/BraTS17_2013_20_1/brainmask.nii.gz\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xvzOImHb_E3f",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "# @title Brain Tumour Segmentation demo\n",
        "#Disclaimer: do not use this code for research, it has been optimised for visualisation in 2D\n",
        "import os\n",
        "import torch\n",
        "import torch.distributions as td\n",
        "from torch.distributions.lowrank_multivariate_normal import _standard_normal, _batch_mv\n",
        "import SimpleITK as sitk\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import warnings\n",
        "from ipywidgets import interactive, IntSlider, FloatSlider, HBox, VBox, Button, HTML, Layout\n",
        "from matplotlib import cm\n",
        "import matplotlib as mpl\n",
        "\n",
        "mpl.rcParams['figure.dpi'] = 300\n",
        "warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
        "torch.autograd.set_grad_enabled(False);\n",
        "COLOR_MAP = ((0, 0, 0), (1, 0, 0), (0, 1, 0), (0, 0, 1), (1, 1, 0), (0, 1, 1), (1, 0, 0))\n",
        "DEVICE = torch.device(0)\n",
        "\n",
        "\n",
        "def cast_to_tensor(array, device, dtype=torch.float32):\n",
        "    return torch.tensor(array.transpose((-1,) + tuple(range(array.ndim - 1))), dtype=dtype, device=device,\n",
        "                        requires_grad=False)\n",
        "\n",
        "\n",
        "class InteractiveSampler(object):\n",
        "    def __init__(self, logit_mean, cov_diag, cov_factor, device, mask, seed=None):\n",
        "        self.seed = seed\n",
        "        if seed is not None:\n",
        "            torch.manual_seed(seed)\n",
        "        self.device = device\n",
        "        self.logit_mean, self.cov_factor, self.cov_diag, self.shape, self.rank = self.build_distribution(logit_mean, cov_diag, cov_factor, device, mask)\n",
        "        self.num_classes = self.shape[0]\n",
        "        self.eps_W, self.eps_D = None, None\n",
        "        self.new_sample()\n",
        "\n",
        "    def new_sample(self):\n",
        "        shape = (1, ) + self.logit_mean.shape\n",
        "        w_shape = shape[:-1] + self.cov_factor.shape[-1:]\n",
        "        dtype=self.logit_mean.dtype\n",
        "        self.eps_W = _standard_normal(w_shape, dtype=dtype, device=self.device)\n",
        "        self.eps_D = _standard_normal(shape, dtype=dtype, device=self.device)\n",
        "\n",
        "    @staticmethod\n",
        "    def build_distribution(logit_mean, cov_diag, cov_factor, device, mask):\n",
        "        logit_mean = cast_to_tensor(logit_mean, device)\n",
        "        cov_diag = cast_to_tensor(cov_diag, device)\n",
        "        cov_factor = cast_to_tensor(cov_factor, device)\n",
        "        if mask is not None:\n",
        "            mask = torch.tensor(mask, device=device, requires_grad=False)\n",
        "            logit_mean[0, ~mask.type(torch.bool)] = 100\n",
        "            cov_factor = cov_factor * mask.unsqueeze(0)\n",
        "        shape = logit_mean.shape\n",
        "        num_classes = shape[0]\n",
        "        rank = int(cov_factor.shape[0] / num_classes)\n",
        "        logit_mean = logit_mean.reshape(-1)\n",
        "        cov_diag = cov_diag.reshape(-1)\n",
        "        cov_factor = cov_factor.reshape((rank, -1)).transpose(1, 0)\n",
        "        epsilon = 1e-3\n",
        "        return logit_mean, cov_factor, cov_diag, shape, rank\n",
        "\n",
        "    def get_manipulated_sample_slice_(self, slice_: int, temperature: float, class_weights: torch.Tensor):\n",
        "        _cov_factor = self.cov_factor.view(self.shape + (self.rank, ))[:, slice_].reshape((-1, self.rank))\n",
        "        factor_direction = _batch_mv(_cov_factor, self.eps_W).view((self.num_classes, ) + self.shape[2:])\n",
        "        \n",
        "        _cov_diag = self.cov_diag.view(self.shape)[:, slice_]\n",
        "        eps_D = self.eps_D.view(self.shape)[:, slice_]\n",
        "        diag_direction = _cov_diag.sqrt() * eps_D\n",
        "\n",
        "        spatial_size = int(np.prod(self.shape[2:]))\n",
        "        class_weights = torch.tensor(class_weights, device=self.device)\n",
        "\n",
        "        class_weights = class_weights.view(-1, 1, 1)\n",
        "        dist_loc = self.logit_mean.view(self.shape)[:, slice_]\n",
        "        logit_samples = dist_loc + temperature * class_weights * (factor_direction + diag_direction)\n",
        "        samples = torch.argmax(logit_samples, dim=0)\n",
        "        means = torch.argmax(dist_loc, dim=0)\n",
        "        return samples, means\n",
        "\n",
        "\n",
        "class Data(object):\n",
        "    def __init__(self):\n",
        "        self.ids = {0: 'BraTS17_2013_10_1', \n",
        "                    1: 'BraTS17_2013_12_1',\n",
        "                    2: 'BraTS17_2013_20_1'}\n",
        "        self.color_map = np.array(COLOR_MAP)\n",
        "        self.id_ = None\n",
        "        self.image_rgb = None\n",
        "        self.seg_rgb = None\n",
        "        self.sampler = None\n",
        "\n",
        "    def overlay_to_rgb(self, overlay):\n",
        "        overlay = overlay.astype(np.uint8)\n",
        "        new_overlay = np.zeros(shape=overlay.shape + (3,), dtype=np.uint8)\n",
        "        for i in range(3):\n",
        "            new_overlay[..., i] = self.color_map[overlay][..., i]\n",
        "        return new_overlay\n",
        "\n",
        "    @staticmethod\n",
        "    def image_to_rgb(image):\n",
        "        image = (image - image.min()) / (image.max() - image.min())\n",
        "        return np.stack((image,) * 3, axis=-1)\n",
        "\n",
        "    @staticmethod\n",
        "    def mix_image_and_overlay(image, overlay, opacity=.5):\n",
        "        overlay = overlay.reshape(image.shape)\n",
        "        new_image = np.copy(image)\n",
        "        ind = np.sum(overlay, axis=-1) > 0\n",
        "        new_image[ind] = opacity * overlay[ind] + (1 - opacity) * image[ind]\n",
        "        return new_image\n",
        "\n",
        "    def generate_data(self, id_):\n",
        "        path = f'data/{id_:s}'\n",
        "        image = sitk.GetArrayFromImage(sitk.ReadImage(os.path.join(path, 't1ce.nii.gz')))\n",
        "        image_rgb = self.image_to_rgb(image)\n",
        "        seg = sitk.GetArrayFromImage(sitk.ReadImage(os.path.join(path, 'seg.nii.gz'))).astype(np.uint8)\n",
        "        seg = self.overlay_to_rgb(seg)\n",
        "        seg_rgb = self.mix_image_and_overlay(image_rgb, seg)\n",
        "        self.image_rgb = image_rgb\n",
        "        self.seg_rgb = seg_rgb\n",
        "        mask = sitk.GetArrayFromImage(sitk.ReadImage(os.path.join(path, 'brainmask.nii.gz')))\n",
        "        logit_mean = sitk.GetArrayFromImage(sitk.ReadImage(os.path.join(path, 'logit_mean.nii.gz')))\n",
        "        cov_factor = sitk.GetArrayFromImage(sitk.ReadImage(os.path.join(path, 'cov_factor.nii.gz')))\n",
        "        cov_diag = sitk.GetArrayFromImage(sitk.ReadImage(os.path.join(path, 'cov_diag.nii.gz')))\n",
        "        self.sampler = InteractiveSampler(logit_mean, cov_diag, cov_factor, DEVICE, mask)\n",
        "\n",
        "    def get_data(self, slice_, image, temperature, class_weights):\n",
        "        id_ = self.ids[image]\n",
        "        if self.id_ is None or self.id_ != id_:\n",
        "            self.id_ = id_\n",
        "            self.generate_data(id_)\n",
        "\n",
        "        sample, mean = self.sampler.get_manipulated_sample_slice_(slice_, temperature, class_weights)\n",
        "\n",
        "        sample = sample.cpu().numpy().astype(np.uint8)\n",
        "        sample = self.overlay_to_rgb(sample)\n",
        "        sample_rgb = self.mix_image_and_overlay(self.image_rgb[slice_], sample)\n",
        "\n",
        "        mean = mean.cpu().numpy().astype(np.uint8)\n",
        "        mean = self.overlay_to_rgb(mean)\n",
        "        mean_rgb = self.mix_image_and_overlay(self.image_rgb[slice_], mean)\n",
        "\n",
        "        return self.image_rgb[slice_], self.seg_rgb[slice_], sample_rgb, mean_rgb\n",
        "\n",
        "\n",
        "data = Data()\n",
        "\n",
        "\n",
        "def plot(image, slice_, temperature, background, necrotic_core, oedema, enhancing_core):\n",
        "    class_weights = [background, necrotic_core, oedema, enhancing_core]\n",
        "\n",
        "    image_rgb, seg_rgb, sample_rgb, mean_rgb = data.get_data(slice_, image, temperature, class_weights)\n",
        "\n",
        "    fig, ax = plt.subplots(1, 4, figsize=(10, 2.5), gridspec_kw=dict(wspace=0, hspace=0))\n",
        "    ax[0].set_title('Image')\n",
        "    ax[0].imshow(image_rgb)\n",
        "    ax[1].set_title('Ground Truth')\n",
        "    ax[1].imshow(seg_rgb)\n",
        "    ax[2].set_title('Prediction (Mean)')\n",
        "    ax[2].imshow(mean_rgb)\n",
        "    ax[3].set_title('Sample')\n",
        "    ax[3].imshow(sample_rgb)\n",
        "\n",
        "    for axi in ax:\n",
        "        axi.axis('off')\n",
        "        axi.xaxis.set_major_locator(plt.NullLocator())\n",
        "        axi.yaxis.set_major_locator(plt.NullLocator())\n",
        "\n",
        "    plt.show()\n",
        "\n",
        "\n",
        "def interactive_plot():\n",
        "    style = {'description_width': '140px'}\n",
        "    layout = Layout(width='400px')\n",
        "    w = interactive(\n",
        "        plot,\n",
        "        image=IntSlider(min=0, max=2, description='Image #', style={'description_width': '100'}, layout=layout),\n",
        "        slice_=IntSlider(value=76, min=0, max=154, description='Slice #', style={'description_width': '100'}, layout=layout),\n",
        "        temperature=FloatSlider(value=1., min=-3., max=3., step=.5, continuous_update=False, description='Temperature', style=style, layout=layout),\n",
        "        background=FloatSlider(value=1., min=-3., max=3., step=.5, continuous_update=False, description='Background', style=style, layout=layout),\n",
        "        necrotic_core=FloatSlider(value=1., min=-3., max=3., step=.5, continuous_update=False, description='Necrotic Core (red)', style=style, layout=layout),\n",
        "        oedema=FloatSlider(value=1., min=-3., max=3., step=.5, continuous_update=False, description='Oedema (green)', style=style, layout=layout),\n",
        "        enhancing_core=FloatSlider(value=1., min=-3., max=3., step=.5, continuous_update=False, description='Enhancing Core (blue)', style=style, layout=layout),\n",
        "    )\n",
        "\n",
        "    new_sample = Button(description=\"New Sample!\")\n",
        "    def new_sample_on_click(b):\n",
        "        if data.sampler is not None:\n",
        "            data.sampler.new_sample()\n",
        "            w.update()\n",
        "    new_sample.on_click(new_sample_on_click)\n",
        "    message = HTML( value=\"Please be patient with the image slider until you see a new image.<br /> It takes a few seconds to load.\")\n",
        "    ui = VBox([\n",
        "        HBox([\n",
        "            VBox([\n",
        "                message,\n",
        "                w.children[0],\n",
        "                w.children[1],\n",
        "            ]),\n",
        "            VBox([\n",
        "                w.children[2],\n",
        "                w.children[3],\n",
        "                w.children[4],\n",
        "                w.children[5],\n",
        "                w.children[6],\n",
        "            ]),\n",
        "            new_sample\n",
        "        ]),\n",
        "        w.children[-1]\n",
        "    ])\n",
        "\n",
        "    display(ui)\n",
        "    w.update()\n",
        "\n",
        "\n",
        "interactive_plot()"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}