guide/training_with_augmentations.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training a segmentation model with data augmentation\n",
    "\n",
    "In this notebook, we will use Nobrainer to augment training data which will be useful to learn more robust model for brain extraction task. Nobrainer provides several methods of augmenting volumetric data including spatial and intensity transforms.Augmentation is useful when the training data is insufficient to train large convolutional neural networks. Augmentation also improves robustness of the network and improves generalization to different datasets. \n",
    "\n",
    "\n",
    "In the following cells, we will:\n",
    "\n",
    "1. Get sample T1-weighted MR scans as features and FreeSurfer segmentations as labels.\n",
    "    - We will binarize the FreeSurfer to get a precise brainmask.\n",
    "2. Convert the data to TFRecords format.\n",
    "3. Create two Datasets of the features and labels.\n",
    "    - One dataset will be for training and the other will be for evaluation.\n",
    "    - Set a sequence of augmentations for training \n",
    "4. Instantiate a 3D convolutional neural network.\n",
    "5. Choose a loss function and metrics to use.\n",
    "6. Train on part of the data.\n",
    "7. Evaluate on the rest of the data.\n",
    "\n",
    "## Google Colaboratory\n",
    "\n",
    "If you are using Colab, please switch your runtime to GPU. To do this, select `Runtime > Change runtime type` in the top menu. Then select GPU under `Hardware accelerator`. A GPU greatly speeds up training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nobrainer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Get sample features and labels\n",
    "\n",
    "We use 9 pairs of volumes for training and 1 pair of volumes for evaulation. Many more volumes would be required to train a model for any useful purpose."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "csv_of_filepaths = nobrainer.utils.get_data()\n",
    "filepaths = nobrainer.io.read_csv(csv_of_filepaths)\n",
    "\n",
    "train_paths = filepaths[:9]\n",
    "evaluate_paths = filepaths[9:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Convert medical images to TFRecords\n",
    "\n",
    "Remember how many full volumes are in the TFRecords files. This will be necessary to know how many steps are in on training epoch. The default training method needs to know this number, because Datasets don't always know how many items they contain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Verify that all volumes have the same shape and that labels are integer-ish.\n",
    "\n",
    "invalid = nobrainer.io.verify_features_labels(train_paths)\n",
    "assert not invalid\n",
    "\n",
    "invalid = nobrainer.io.verify_features_labels(evaluate_paths)\n",
    "assert not invalid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Convert training and evaluation data to TFRecords.\n",
    "\n",
    "nobrainer.tfrecord.write(\n",
    "    features_labels=train_paths,\n",
    "    filename_template='data/data-train_shard-{shard:03d}.tfrec',\n",
    "    examples_per_shard=3)\n",
    "\n",
    "nobrainer.tfrecord.write(\n",
    "    features_labels=evaluate_paths,\n",
    "    filename_template='data/data-evaluate_shard-{shard:03d}.tfrec',\n",
    "    examples_per_shard=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_classes = 1\n",
    "batch_size = 2\n",
    "volume_shape = (256, 256, 256)\n",
    "block_shape = (32, 32, 32)\n",
    "n_epochs = None\n",
    "shuffle_buffer_size = 10\n",
    "num_parallel_calls = 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Take a look at different augmentation options in Nobrainer spatial and intensity transforms. To set training with multiple augmentations, the 'augment' will be set as a list where their order will determine the sequence of execution. For example augment option below will first add Gaussian noise and will then perform the random flip. \n",
    "Parameters of any given augmentation techniques can be set as shown below ( eg. 'noise_mean':0.1') otherwise default parameter settings will be applied. \n",
    "\n",
    "For training without augmentation, set 'augment = None'. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nobrainer.intensity_transforms import *\n",
    "from nobrainer.spatial_transforms import *\n",
    "\n",
    "augment = [(addGaussianNoise, {'noise_mean':0.1,'noise_std':0.5}), (randomflip_leftright)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_train = nobrainer.dataset.get_dataset(\n",
    "    file_pattern='data/data-train_shard-*.tfrec',\n",
    "    n_classes=n_classes,\n",
    "    batch_size=batch_size,\n",
    "    volume_shape=volume_shape,\n",
    "    block_shape=block_shape,\n",
    "    n_epochs=n_epochs,\n",
    "    augment=augment,\n",
    "    shuffle_buffer_size=shuffle_buffer_size,\n",
    "    num_parallel_calls=num_parallel_calls,\n",
    ")\n",
    "\n",
    "dataset_evaluate = nobrainer.dataset.get_dataset(\n",
    "    file_pattern='data/data-evaluate_shard-*.tfrec',\n",
    "    n_classes=n_classes,\n",
    "    batch_size=batch_size,\n",
    "    volume_shape=volume_shape,\n",
    "    block_shape=block_shape,\n",
    "    n_epochs=1,\n",
    "    augment=False,\n",
    "    shuffle_buffer_size=None,\n",
    "    num_parallel_calls=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_evaluate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Instantiate a neural network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = nobrainer.models.unet(\n",
    "    n_classes=n_classes, \n",
    "    input_shape=(*block_shape, 1),\n",
    "    batchnorm=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Choose a loss function and metrics\n",
    "\n",
    "We have many choices of loss functions for binary segmentation. One can choose from binary crossentropy, Dice, Jaccard, Tversky, and many other loss functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = tf.keras.optimizers.Adam(learning_rate=1e-04)\n",
    "\n",
    "model.compile(\n",
    "    optimizer=optimizer,\n",
    "    loss=nobrainer.losses.dice,\n",
    "    metrics=[nobrainer.metrics.dice, nobrainer.metrics.jaccard],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train and evaluate model\n",
    "\n",
    "$$\n",
    "steps = \\frac{nBlocks}{volume} * \\frac{nVolumes}{batchSize}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "steps_per_epoch = nobrainer.dataset.get_steps_per_epoch(\n",
    "    n_volumes=len(train_paths),\n",
    "    volume_shape=volume_shape,\n",
    "    block_shape=block_shape,\n",
    "    batch_size=batch_size)\n",
    "\n",
    "steps_per_epoch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "validation_steps = nobrainer.dataset.get_steps_per_epoch(\n",
    "    n_volumes=len(evaluate_paths),\n",
    "    volume_shape=volume_shape,\n",
    "    block_shape=block_shape,\n",
    "    batch_size=batch_size)\n",
    "\n",
    "validation_steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "model.fit(\n",
    "    dataset_train,\n",
    "    epochs=5,\n",
    "    steps_per_epoch=steps_per_epoch, \n",
    "    validation_data=dataset_evaluate, \n",
    "    validation_steps=validation_steps)"
   ]
  }
 ],
 "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}