examples/Neuro_Demo.ipynb

{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  },
  "colab": {
   "name": "Neuro UNet Demo.ipynb",
   "provenance": [],
   "include_colab_link": true
  },
  "accelerator": "GPU"
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "view-in-github",
    "colab_type": "text"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/catalyst-team/neuro/blob/master/examples/Neuro_Demo.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FZOKa-kgP_L-"
   },
   "source": [
    "# Neuro UNet/Meshnet Tutorial\n",
    "\n",
    "This is a reimplementation of [An (almost) instant brain atlas segmentation for\n",
    "large-scale studies](https://arxiv.org/pdf/1711.00457.pdf) using the publicly available Mindboggle Dataset.  Given only 70 volumes for training (typically 700+) and very minimal augmentations we can achieve a 0.6688 Mean Dice Score for a brain with a majority vote classification.  With more brains and normalizing the T1 scans we could likely perform better\n",
    "\n",
    "\n",
    "Authors: [Kevin Wang](https://github.com/ssktotoro/), [Alex Fedorov](https://github.com/Entodi/), [Sergey Kolesnikov](https://github.com/Scitator)\n",
    "\n",
    "[![Catalyst logo](https://raw.githubusercontent.com/catalyst-team/catalyst-pics/master/pics/catalyst_logo.png)](https://github.com/catalyst-team/catalyst)\n",
    "\n",
    "### Colab setup\n",
    "\n",
    "First of all, do not forget to change the runtime type to GPU. <br/>\n",
    "To do so click `Runtime` -> `Change runtime type` -> Select `\\\"Python 3\\\"` and `\\\"GPU\\\"` -> click `Save`. <br/>\n",
    "After that you can click `Runtime` -> `Run all` and watch the tutorial."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z8hXadxdP_L-"
   },
   "source": [
    "## Requirements\n",
    "\n",
    "Download and install the latest versions of catalyst and other libraries required for this tutorial."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "-T0L6CqRQ-23"
   },
   "source": [
    "%%bash \n",
    "git clone https://github.com/catalyst-team/neuro.git\n",
    "pip install -r neuro/requirements/requirements.txt\n"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "k1BnwldEP_L-"
   },
   "source": [
    "import torch                                                                                                                                                                                              \n",
    "from tqdm import tqdm                                                                                                                                                                                     \n",
    "import numpy as np                                                                                                                                                                                        \n",
    "import nibabel as nib                                                                                                                                                                                     \n",
    "import collections                                                                                                                                                                                        \n",
    "from collections import OrderedDict                                                                                                                                                                       \n",
    "                                                                                                                                                                                                          \n",
    "import catalyst                                                                                                                                                                                           \n",
    "import pandas as pd  \n",
    "import os                                                                                                                                                                                     \n",
    "                                                                                                                                                                                                          \n",
    "from catalyst.contrib.utils.pandas import dataframe_to_list                                                                                                                                               \n",
    "from torch.utils.data import SequentialSampler                                                                                                                                                            \n",
    "from torch.utils.data import DataLoader                                                                                                                                                                   \n",
    "from catalyst.data import ReaderCompose                                                                                                                                                                   \n",
    "from torch.optim.lr_scheduler import CosineAnnealingLR, OneCycleLR                                                                                                                                        \n",
    "from catalyst.callbacks import CheckpointCallback                                                                                                                                                         \n",
    "from torch.nn import functional as F                                                                                                                                                                      \n",
    "from typing import List\n",
    "from catalyst import utils                                                                                                                                                                                   \n",
    "                                                                                                                                                                                                                                                                                                                                                                                    \n",
    "from catalyst import metrics                                                                                                                                                                              \n",
    "from catalyst.data import BatchPrefetchLoaderWrapper                                                                                                                                                      \n",
    "from catalyst.dl import Runner, LRFinder                                                                                                                                                                  \n",
    "                                                                                                                                                                                                          \n",
    "from catalyst.metrics.functional._segmentation import dice\n",
    "\n",
    "\n",
    "print(f\"torch: {torch.__version__}, catalyst: {catalyst.__version__}\")\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"  # \"\" - CPU, \"0\" - 1 GPU, \"0,1\" - MultiGPU\n",
    "\n",
    "SEED = 42\n",
    "utils.set_global_seed(SEED)\n",
    "utils.prepare_cudnn(deterministic=True)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DOGktFaFP_L_"
   },
   "source": [
    "# Dataset\n",
    "\n",
    "We'll be using the Mindboggle 101 dataset for a multiclass 3d segmentation task.\n",
    "The dataset can be downloaded off osf with the following command from osfclient after you register with osf.\n",
    "\n",
    "`osf -p 9ahyp clone .`\n",
    "\n",
    "Otherwise you can download it using a Catalyst utility `download-gdrive` which downloads a version from the Catalyst Google Drive\n",
    "\n",
    "`usage: download-gdrive {FILE_ID} {FILENAME}`"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "Oupu01n3RJb5"
   },
   "source": [
    "cd neuro"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "2VlogHK7RGkJ"
   },
   "source": [
    "%%bash\n",
    "mkdir Mindboggle_data \n",
    "mkdir -p data/Mindboggle_101/\n",
    "osf -p 9ahyp clone Mindboggle_data/\n",
    "cp -r Mindboggle_data/osfstorage/Mindboggle101_volumes/ data/Mindboggle_101/\n",
    "find data/Mindboggle_101 -name '*.tar.gz'| xargs -i tar zxvf {} -C data/Mindboggle_101\n",
    "find data/Mindboggle_101 -name '*.tar.gz'| xargs -i rm {}"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VCJ4_qnGP_MA"
   },
   "source": [
    "Run the prepare data script that limits the labels to the DKT cortical labels (31 labels).  We can use of course use more labels.\n",
    "\n",
    "`usage: python ../neuro/scripts/prepare_data.py ../data/Mindboggle_101 {N_labels)`"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "aBFRPKyuP_MA"
   },
   "source": [
    "%%bash \n",
    "\n",
    "python neuro/scripts/prepare_data.py data/Mindboggle_101/ 31"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kviNcrfIP_MA"
   },
   "source": [
    "### **Create the relevant Dataloaders for the specified train, validation, and inference BrainDatasets.**\n",
    "\n",
    "BrainDatasets comprise of T1 scans + the prepared limited labels.\n",
    "\n",
    "Training/ Validation batches: To avoid sampling outside the brain area we randomly sampled NxNxN subvolumes from a gaussian distribution with the mean set in the center of the volume and the diagonal covariance set at 50.  This helps avoid class imbalance that would lead the model to simpley prefer the background class.\n",
    "\n",
    "Inference batches: All non-overlapping NxNxN Subvolumes across the existing volume space with their corresponding labels + randomly sampled NxNxN subvolumes from a the same gaussian distribution accross volume space as above until the required number of subvolume is reached.\n",
    "\n",
    "More detail can be found in brain_dataset.py and generator_coords.py  "
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "IeFGpQ4_Ufdw"
   },
   "source": [
    "cd neuro/"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "zjVQEnItP_MA"
   },
   "source": [
    "from neuro.brain_dataset import BrainDataset\n",
    "from neuro.reader import NiftiFixedVolumeReader, NiftiReader\n",
    "from neuro.model import MeshNet, UNet"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "VSWPJkTgP_MA"
   },
   "source": [
    "def get_loaders(\n",
    "  random_state: int,\n",
    "  volume_shape: List[int],\n",
    "  subvolume_shape: List[int],\n",
    "  in_csv_train: str = None,\n",
    "  in_csv_valid: str = None,\n",
    "  in_csv_infer: str = None,\n",
    "  batch_size: int = 16,\n",
    "  num_workers: int = 10,\n",
    ") -> dict:\n",
    "\n",
    "  datasets = {}\n",
    "  open_fn = ReaderCompose(\n",
    "      [\n",
    "          NiftiFixedVolumeReader(input_key=\"images\", output_key=\"images\"),\n",
    "          NiftiReader(input_key=\"nii_labels\", output_key=\"targets\"),\n",
    "\n",
    "      ]\n",
    "  )\n",
    "\n",
    "  for mode, source in zip((\"train\", \"validation\", \"infer\"),\n",
    "                          (in_csv_train, in_csv_valid, in_csv_infer)):\n",
    "      if mode == \"infer\":\n",
    "          n_subvolumes = 512\n",
    "      else:\n",
    "          n_subvolumes = 128\n",
    "\n",
    "      if source is not None and len(source) > 0:\n",
    "          dataset = BrainDataset(\n",
    "              list_data=dataframe_to_list(pd.read_csv(source)),\n",
    "              list_shape=volume_shape,\n",
    "              list_sub_shape=subvolume_shape,\n",
    "              open_fn=open_fn,\n",
    "              n_subvolumes=n_subvolumes,\n",
    "              mode=mode,\n",
    "              input_key=\"images\",\n",
    "              output_key=\"targets\",\n",
    "          )\n",
    "\n",
    "      datasets[mode] = {\"dataset\": dataset}\n",
    "\n",
    "  def worker_init_fn(worker_id):\n",
    "      np.random.seed(np.random.get_state()[1][0] + worker_id)\n",
    "\n",
    "\n",
    "  train_loader = DataLoader(dataset=datasets['train']['dataset'], batch_size=batch_size,\n",
    "                            shuffle=True, worker_init_fn=worker_init_fn,\n",
    "                            num_workers=2, pin_memory=True)\n",
    "  valid_loader = DataLoader(dataset=datasets['validation']['dataset'],\n",
    "                            shuffle=True, worker_init_fn=worker_init_fn,\n",
    "                            batch_size=batch_size,\n",
    "                            num_workers=2, pin_memory=True,drop_last=True)\n",
    "  test_loader = DataLoader(dataset=datasets['infer']['dataset'],\n",
    "                           batch_size=batch_size, worker_init_fn=worker_init_fn,\n",
    "                           num_workers=2, pin_memory=True,drop_last=True)\n",
    "  train_loaders = collections.OrderedDict()\n",
    "  infer_loaders = collections.OrderedDict()\n",
    "  train_loaders[\"train\"] = BatchPrefetchLoaderWrapper(train_loader)\n",
    "  train_loaders[\"valid\"] = BatchPrefetchLoaderWrapper(valid_loader)\n",
    "  infer_loaders['infer'] = BatchPrefetchLoaderWrapper(test_loader)\n",
    "\n",
    "  return train_loaders, infer_loaders"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "XL43XWkYjCGx"
   },
   "source": [
    "cd ../"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "ORmV32oeP_MA"
   },
   "source": [
    "volume_shape = [256, 256, 256]\n",
    "subvolume_shape = [38, 38, 38]\n",
    "train_loaders, infer_loaders = get_loaders(\n",
    "    0, volume_shape, subvolume_shape,\n",
    "    \"./data/dataset_train.csv\",\n",
    "    \"./data/dataset_valid.csv\",\n",
    "    \"./data/dataset_infer.csv\"\n",
    ")"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "s_Qav0lxP_MA"
   },
   "source": [
    "# Model Training\n",
    "\n",
    "We'll train the model 30 epochs.\n",
    "\n",
    "* Scheduler: Adam with a One Cycle Learning Rate with a Max Learning Rate of .02\n",
    "* Batch Metric: DICE\n",
    "* Loss: CrossEntropyLoss\n",
    "* Logger: Tensorboard\n",
    "* CheckpointerCallback\n",
    "\n",
    "For training and validation we sample the volume with subvolumes specified in our Dataset "
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "TFnFWvyZXt5l"
   },
   "source": [
    "class CustomRunner(Runner):\n",
    "\n",
    "    def get_loaders(self, stage: str) -> \"OrderedDict[str, DataLoader]\":\n",
    "        \"\"\"Returns the loaders for a given stage.\"\"\"\n",
    "        self._loaders = self._loaders\n",
    "        return self._loaders\n",
    "\n",
    "    def predict_batch(self, batch):\n",
    "        # model inference step\n",
    "        batch = batch[0]\n",
    "        return self.model(batch['images'].float().to(self.device)), batch['coords']\n",
    "\n",
    "    def on_loader_start(self, runner):\n",
    "        super().on_loader_start(runner)\n",
    "        self.meters = {\n",
    "          key: metrics.AdditiveValueMetric(compute_on_call=False)\n",
    "          for key in [\"loss\", \"macro_dice\"]\n",
    "        }\n",
    "\n",
    "    def handle_batch(self, batch):\n",
    "        # model train/valid step\n",
    "        batch = batch[0]\n",
    "        x, y = batch['images'].float(), batch['targets']\n",
    "\n",
    "        if self.is_train_loader:\n",
    "          self.optimizer.zero_grad()\n",
    "\n",
    "        y_hat = self.model(x)\n",
    "        loss = F.cross_entropy(y_hat, y)\n",
    "\n",
    "        if self.is_train_loader:\n",
    "          loss.backward()\n",
    "          self.optimizer.step()\n",
    "          scheduler.step()\n",
    "\n",
    "        one_hot_targets = (\n",
    "          torch.nn.functional.one_hot(y, 31)\n",
    "          .permute(0, 4, 1, 2, 3)\n",
    "          .cuda()\n",
    "        )\n",
    "\n",
    "        logits_softmax = F.softmax(y_hat)\n",
    "        macro_dice = dice(logits_softmax, one_hot_targets, mode='macro')\n",
    "\n",
    "        self.batch_metrics.update({\"loss\": loss, 'macro_dice': macro_dice})\n",
    "        for key in [\"loss\", \"macro_dice\"]:\n",
    "          self.meters[key].update(self.batch_metrics[key].item(), self.batch_size)\n",
    "\n",
    "    def on_loader_end(self, runner):\n",
    "        for key in [\"loss\", \"macro_dice\"]:\n",
    "            self.loader_metrics[key] = self.meters[key].compute()[0]\n",
    "        super().on_loader_end(runner)\n",
    "                                           "
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "XQ9dzP86dAmx",
    "collapsed": true
   },
   "source": [
    "n_classes = 31\n",
    "n_epochs = 30\n",
    "meshnet = MeshNet(n_channels=1, n_classes=n_classes)\n",
    "\n",
    "logdir = \"logs/meshnet_mindboggle\"\n",
    "\n",
    "optimizer = torch.optim.Adam(meshnet.parameters(), lr=0.02)\n",
    "\n",
    "\n",
    "scheduler = OneCycleLR(optimizer, max_lr=.02,\n",
    "                 epochs=n_epochs, steps_per_epoch=len(train_loaders['train']))\n",
    "\n",
    "runner = CustomRunner()\n",
    "runner.train(\n",
    "    model=meshnet,\n",
    "    optimizer=optimizer,\n",
    "    scheduler=scheduler,\n",
    "    loaders=train_loaders,\n",
    "    num_epochs=n_epochs,\n",
    "    logdir=logdir,\n",
    "    callbacks=[CheckpointCallback(logdir=logdir)],\n",
    "    verbose=True\n",
    ")"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CqJgyfgIsLZ4"
   },
   "source": [
    "# Model Evaluation\n",
    "\n",
    "For every brain volume we implement a majority vote for every voxel and use  that to compute a Dice score.\n",
    "\n",
    "\n",
    "The initial volume is segmented into\n",
    "a regular grid of subvolumes partitioning the whole volume.\n",
    "These volumes ensure a prediction for each voxel. Then we\n",
    "sample overlapping volumes from the brain region until N\n",
    "subvolumes (512 in this case) are achieved for prediction.  \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "sijB1BVYisp4"
   },
   "source": [
    "def voxel_majority_predict_from_subvolumes(loader, n_classes, segmentations):\n",
    "    if segmentations is None:\n",
    "        for subject in range(loader.dataset.subjects):\n",
    "            segmentations[subject] = torch.zeros(\n",
    "                tuple(np.insert(loader.volume_shape, 0, n_classes)),\n",
    "                dtype=torch.uint8).cpu()\n",
    "\n",
    "    prediction_n = 0\n",
    "    for inference in tqdm(runner.predict_loader(loader=loader)):\n",
    "        coords = inference[1].cpu()\n",
    "        _, predicted = torch.max(F.log_softmax(inference[0].cpu(), dim=1), 1)\n",
    "        for j in range(predicted.shape[0]):\n",
    "            c_j = coords[j][0]\n",
    "            subj_id = prediction_n // loader.dataset.n_subvolumes\n",
    "            for c in range(n_classes):\n",
    "                segmentations[subj_id][c, c_j[0, 0]:c_j[0, 1],\n",
    "                                       c_j[1, 0]:c_j[1, 1],\n",
    "                                       c_j[2, 0]:c_j[2, 1]] += (predicted[j] == c)\n",
    "            prediction_n += 1\n",
    "\n",
    "    for i in segmentations.keys():\n",
    "        segmentations[i] = torch.max(segmentations[i], 0)[1]\n",
    "    return segmentations"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "al25kqCRV4En"
   },
   "source": [
    "segmentations = {}\n",
    "for subject in range(infer_loaders['infer'].dataset.subjects):\n",
    "    segmentations[subject] = torch.zeros(tuple(np.insert(volume_shape, 0, n_classes)), dtype=torch.uint8)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "0TdoZPOGVbaQ"
   },
   "source": [
    "segmentations = voxel_majority_predict_from_subvolumes(infer_loaders['infer'],\n",
    "                                                     n_classes, segmentations)\n",
    "subject_metrics = []\n",
    "for subject, subject_data in enumerate(tqdm(infer_loaders['infer'].dataset.data)):\n",
    "    seg_labels = nib.load(subject_data['nii_labels']).get_fdata()\n",
    "    segmentation_labels = torch.nn.functional.one_hot(\n",
    "        torch.from_numpy(seg_labels).to(torch.int64), n_classes)\n",
    "\n",
    "    inference_dice = dice(\n",
    "        torch.nn.functional.one_hot(segmentations[subject], n_classes).permute(0, 3, 1, 2),\n",
    "        segmentation_labels.permute(0, 3, 1, 2)\n",
    "    ).detach().numpy()\n",
    "    macro_inference_dice = dice(\n",
    "        torch.nn.functional.one_hot(segmentations[subject], n_classes).permute(0, 3, 1, 2),\n",
    "        segmentation_labels.permute(0, 3, 1, 2), mode='macro'\n",
    "    ).detach().numpy()\n",
    "    subject_metrics.append((inference_dice, macro_inference_dice))\n",
    "\n",
    "per_class_df = pd.DataFrame([metric[0] for metric in subject_metrics])\n",
    "macro_df = pd.DataFrame([metric[1] for metric in subject_metrics])\n",
    "print(per_class_df, macro_df)\n",
    "print(macro_df.mean())"
   ],
   "execution_count": null,
   "outputs": []
  }
 ]
}