example_1.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multiple Sclerosis (MS) lesion segmentation of MRI images using a cascade of two 3D convolutional neural networks \n",
    "\n",
    "\n",
    "This script assumes that `Lasagne` and `nolearn` have been installed correctly and `CUDA / CUDNN` are configured. \n",
    "\n",
    "Import libraries: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using gpu device 0: GeForce GTX 1080 (CNMeM is disabled, cuDNN 5105)\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import os\n",
    "from collections import OrderedDict\n",
    "from base import *\n",
    "from build_model_nolearn import cascade_model\n",
    "from config import *\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model configuration:\n",
    "Configure the model options. Options are passed to the model using the dictionary `options`. The main options are:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "options = {}\n",
    "\n",
    "# --------------------------------------------------\n",
    "# Experiment parameters\n",
    "# --------------------------------------------------\n",
    "\n",
    "# image modalities used (T1, FLAIR, PD, T2, ...) \n",
    "options['modalities'] = ['T1', 'FLAIR']\n",
    "\n",
    "# Select an experiment name to store net weights and segmentation masks\n",
    "options['experiment'] = 'test_CNN'\n",
    "\n",
    "# In order to expose the classifier to more challeging samples, a threshold can be used to to select \n",
    "# candidate voxels for training. Note that images are internally normalized to 0 mean 1 standard deviation \n",
    "# before applying thresholding. So a value of t > 0.5 on FLAIR is reasonable in most cases to extract \n",
    "# all WM lesion candidates\n",
    "options['min_th'] = 0.5\n",
    "\n",
    "# randomize training features before fitting the model.  \n",
    "options['randomize_train'] = True\n",
    "\n",
    "# Select between pixel-wise or fully-convolutional training models. Although implemented, fully-convolutional\n",
    "# models have been not tested with this cascaded model \n",
    "options['fully_convolutional'] = False\n",
    "\n",
    "\n",
    "# --------------------------------------------------\n",
    "# model parameters\n",
    "# --------------------------------------------------\n",
    "\n",
    "# 3D patch size. So, far only implemented for 3D CNN models. \n",
    "options['patch_size'] = (11,11,11)\n",
    "\n",
    "# percentage of the training vector that is going to be used to validate the model during training\n",
    "options['train_split'] = 0.25\n",
    "\n",
    "# maximum number of epochs used to train the model\n",
    "options['max_epochs'] = 200\n",
    "\n",
    "# maximum number of epochs without improving validation before stopping training (early stopping) \n",
    "options['patience'] = 25\n",
    "\n",
    "# Number of samples used to test at once. This parameter should be around 50000 for machines\n",
    "# with less than 32GB of RAM\n",
    "options['batch_size'] = 50000\n",
    "\n",
    "# net print verbosity. Set to zero for this particular notebook, but setting this value to 11 is recommended\n",
    "options['net_verbose'] = 0\n",
    "\n",
    "# post-processing binary threshold. After segmentation, probabilistic masks are binarized using a defined threshold.\n",
    "options['t_bin'] = 0.8\n",
    "\n",
    "# The resulting binary mask is filtered by removing lesion regions with lesion size before a defined value\n",
    "options['l_min'] = 20\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Experiment configuration:\n",
    "\n",
    "Organize the experiment. Although not necessary, intermediate results, network weights and final lesion segmentation masks are stored inside a folder with name `options['experiment']`. This is extremely useful when a lot of experiments are computed on the same images to declutter the user space. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "exp_folder = os.path.join(os.getcwd(), options['experiment'])\n",
    "if not os.path.exists(exp_folder):\n",
    "    os.mkdir(exp_folder)\n",
    "    os.mkdir(os.path.join(exp_folder,'nets'))\n",
    "    os.mkdir(os.path.join(exp_folder,'.train'))\n",
    "\n",
    "# set the output name \n",
    "options['test_name'] = 'cnn_' + options['experiment'] + '.nii.gz'\n",
    " \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the training data:\n",
    "\n",
    "Training data is internally loaded by the method. So far, training and testing images are passed as dictionaries, where each training image is stored as follows: \n",
    "\n",
    "```\n",
    "traininig_X_data['image_identifier'] = {'modality_1': /path/to/image_modality_n.nii.gz/,\n",
    "                                         ....\n",
    "                                        'modality_n': /path/to/image_modality_n.nii.gz/}\n",
    "```\n",
    "\n",
    "And also for labels: \n",
    "\n",
    "```\n",
    "traininig_y_data['image_identifier_1'] = 'path/to/image_labels.nii.gz/'\n",
    "```\n",
    "\n",
    "**NOTE**: As stated in the paper, input images have been already skull-stripped and bias corrected (N3, etc...) by the user before running the classifer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "train_folder = '/mnt/DATA/w/CNN/images/train_images'\n",
    "train_x_data = {}\n",
    "train_y_data = {}\n",
    "\n",
    "# TRAIN X DATA\n",
    "train_x_data['im1'] = {'T1': os.path.join(train_folder,'im1', 'T1.nii.gz'), \n",
    "                       'FLAIR': os.path.join(train_folder,'im1', 'FLAIR.nii.gz')}\n",
    "train_x_data['im2'] = {'T1': os.path.join(train_folder,'im2', 'T1.nii.gz'), \n",
    "                       'FLAIR': os.path.join(train_folder,'im2', 'FLAIR.nii.gz')}\n",
    "train_x_data['im3'] = {'T1': os.path.join(train_folder,'im3', 'T1.nii.gz'), \n",
    "                       'FLAIR': os.path.join(train_folder,'im3', 'FLAIR.nii.gz')}\n",
    "\n",
    "# TRAIN LABELS \n",
    "train_y_data['im1'] = os.path.join(train_folder,'im1', 'lesion_bin.nii.gz')\n",
    "train_y_data['im2'] = os.path.join(train_folder,'im2', 'lesion_bin.nii.gz')\n",
    "train_y_data['im3'] = os.path.join(train_folder,'im3', 'lesion_bin.nii.gz')\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize the model:\n",
    "\n",
    "The model is initialized using the function `cascade_model`, which returns a list of two `NeuralNet` objects. Optimized weights are stored also inside the experiment folder for future use (testing different images without re-training the model.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "options['weight_paths'] = os.getcwd()\n",
    "model = cascade_model(options)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train the model:\n",
    "\n",
    "The function `train_cascaded_model` is used to train the model. The next image summarizes the training procedure. For further information about how this function optimizes the two CNN, please consult the original paper. (**NOTE**: For this example, `options['net_verbose`] has been set to `0` for simplicity)\n",
    "\n",
    "\n",
    "![](pipeline_training.png)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false,
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---> cnn1 loading training data\n",
      "---> cnn1 train_x  (32290, 2, 11, 11, 11) \n",
      "\n",
      "Early stopping.\n",
      "Best valid loss was 0.060829 at epoch 44.\n",
      "---> cnn2 loading training data\n",
      "---> cnn2 train_x  (32290, 2, 11, 11, 11) \n",
      "\n"
     ]
    }
   ],
   "source": [
    "model = train_cascaded_model(model, train_x_data, train_y_data,  options)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test the model:\n",
    "\n",
    "Once the model has been trained, it can e tested on other images. Please note that the same image modalities have to be used. Testing images are loaded equally to training_data, so a `dictionary` defines the modalities used:\n",
    "\n",
    "```\n",
    "test_X_data['image_identifier'] = {'modality_1': /path/to/image_modality_n.nii.gz/,\n",
    "                                         ....\n",
    "                                   'modality_n': /path/to/image_modality_n.nii.gz/}\n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " ---> testing the model\n"
     ]
    }
   ],
   "source": [
    "# TEST X DATA\n",
    "test_folder = '/mnt/DATA/w/CNN/images/test_images'\n",
    "test_x_data = {}\n",
    "test_x_data['im1'] = {'T1': os.path.join(test_folder,'im1', 'T1.nii.gz'), \n",
    "                       'FLAIR': os.path.join(test_folder,'im1', 'FLAIR.nii.gz')}\n",
    "\n",
    "\n",
    "# set the output_location of the final segmentation. In this particular example, \n",
    "# we are training and testing on the same images\n",
    "options['test_folder'] = test_folder\n",
    "options['test_scan'] = 'im1'\n",
    "out_seg = test_cascaded_model(model, test_x_data, options)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Compute different metrics on tested data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from metrics import *\n",
    "# load the GT annotation for the tested image \n",
    "GT = nib.load(os.path.join(test_folder,'im1', 'lesion_bin.nii.gz')).get_data()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}