weakly/tutorial_training.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import sys\n",
    "import random\n",
    "import time\n",
    "\n",
    "import labelreg.helpers as helper\n",
    "import labelreg.networks as network\n",
    "import labelreg.utils as util\n",
    "import labelreg.losses as loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "NB. The ini config file is not used in this tutorial for clarity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# config = helper.ConfigParser(sys.argv, 'training')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set up the data feeders that will be used in the training, by specifying the folders containing the images and labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reader_moving_image, reader_fixed_image, reader_moving_label, reader_fixed_label = helper.get_data_readers(\n",
    "    './data/train/mr_images',\n",
    "    './data/train/us_images',\n",
    "    './data/train/mr_labels',\n",
    "    './data/train/us_labels')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The placeholders for moving images, fixed images and their associated labels; <br/>\n",
    "The on-the-fly data augmentation use random affine transformation, independently drawn for moving and fixed data; <br/>\n",
    "The minibatch size is 4 here; <br/>\n",
    "The parameters of the 12 degrees-of-freedom affine transformationwill be used for augmenting the input data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ph_moving_image = tf.placeholder(tf.float32, [4]+reader_moving_image.data_shape+[1])\n",
    "ph_fixed_image = tf.placeholder(tf.float32, [4]+reader_fixed_image.data_shape+[1])\n",
    "ph_moving_affine = tf.placeholder(tf.float32, [4]+[1, 12])\n",
    "ph_fixed_affine = tf.placeholder(tf.float32, [4]+[1, 12])\n",
    "input_moving_image = util.warp_image_affine(ph_moving_image, ph_moving_affine)  # data augmentation\n",
    "input_fixed_image = util.warp_image_affine(ph_fixed_image, ph_fixed_affine)  # data augmentation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now load the instance of the \"local\" network, a single U-Net-like encoder-decoder network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reg_net = network.build_network(network_type='local',\n",
    "                                minibatch_size=4,\n",
    "                                image_moving=input_moving_image,\n",
    "                                image_fixed=input_fixed_image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ph_moving_label = tf.placeholder(tf.float32, [4]+reader_moving_image.data_shape+[1])\n",
    "ph_fixed_label = tf.placeholder(tf.float32, [4]+reader_fixed_image.data_shape+[1])\n",
    "input_moving_label = util.warp_image_affine(ph_moving_label, ph_moving_affine)  # data augmentation\n",
    "input_fixed_label = util.warp_image_affine(ph_fixed_label, ph_fixed_affine)  # data augmentation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Warp the moving label with the predicted ddf."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "warped_moving_label = reg_net.warp_image(input_moving_label)  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compute the loss: <br/>\n",
    "* The label similarity between the warped moving labels and fixed labels; <br/>\n",
    "* The weighted (here, 0.5) defomation regularisation on the predicted DDFs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_similarity = tf.reduce_mean(loss.multi_scale_loss(input_fixed_label, warped_moving_label, 'dice', [0, 1, 2, 4, 8]))\n",
    "loss_regulariser = tf.reduce_mean(loss.local_displacement_energy(reg_net.ddf, 'bending', 0.5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Build the training op with a learning rate of 1e-04."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_op = tf.train.AdamOptimizer(learning_rate=1e-04).minimize(loss_similarity+loss_regulariser)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These are the two utility nodes for information during the training iterations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dice = util.compute_binary_dice(warped_moving_label, input_fixed_label)\n",
    "dist = util.compute_centroid_distance(warped_moving_label, input_fixed_label)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set up the training session."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_minibatch = int(reader_moving_label.num_data/4)\n",
    "train_indices = [i for i in range(reader_moving_label.num_data)]\n",
    "\n",
    "saver = tf.train.Saver(max_to_keep=1)\n",
    "sess = tf.Session()\n",
    "sess.run(tf.global_variables_initializer())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Start training! <br/>\n",
    "Save model to ./data/model.ckpt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for step in range(10000):\n",
    "\n",
    "    if step in range(0, 10000, num_minibatch):\n",
    "        random.shuffle(train_indices)\n",
    "\n",
    "    minibatch_idx = step % num_minibatch\n",
    "    case_indices = train_indices[\n",
    "                    minibatch_idx*4:(minibatch_idx+1)*4]\n",
    "    label_indices = [random.randrange(reader_moving_label.num_labels[i]) for i in case_indices]\n",
    "\n",
    "    trainFeed = {ph_moving_image: reader_moving_image.get_data(case_indices),\n",
    "                 ph_fixed_image: reader_fixed_image.get_data(case_indices),\n",
    "                 ph_moving_label: reader_moving_label.get_data(case_indices, label_indices),\n",
    "                 ph_fixed_label: reader_fixed_label.get_data(case_indices, label_indices),\n",
    "                 ph_moving_affine: helper.random_transform_generator(4),\n",
    "                 ph_fixed_affine: helper.random_transform_generator(4)}\n",
    "\n",
    "    sess.run(train_op, feed_dict=trainFeed)\n",
    "\n",
    "    if step in range(0, 10000, 10):  # print info every 10 iterations\n",
    "        current_time = time.asctime(time.gmtime())\n",
    "        loss_similarity_train, loss_regulariser_train, dice_train, dist_train = sess.run(\n",
    "            [loss_similarity,\n",
    "             loss_regulariser,\n",
    "             dice,\n",
    "             dist],\n",
    "            feed_dict=trainFeed)\n",
    "\n",
    "        # print('----- Training -----')\n",
    "        print('Step %d [%s]: Loss=%f (similarity=%f, regulariser=%f)' %\n",
    "              (step,\n",
    "               current_time,\n",
    "               loss_similarity_train+loss_regulariser_train,\n",
    "               1-loss_similarity_train,\n",
    "               loss_regulariser_train))\n",
    "        print('  Dice: %s' % dice_train)\n",
    "        print('  Distance: %s' % dist_train)\n",
    "        print('  Image-label indices: %s - %s' % (case_indices, label_indices))\n",
    "\n",
    "    if step in range(0, 10000, 100):  # save the model every 100 iterations\n",
    "        save_path = saver.save(sess, './data/model.ckpt', write_meta_graph=False)\n",
    "        print(\"Model saved in: %s\" % save_path)"
   ]
  }
 ],
 "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}