CS380-UNet.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# CS380 UNet\n",
    "\n",
    "**Author:** Will Hewitt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Purpose\n",
    "\n",
    "Trains UNet for segmenting echocardiography images\n",
    "\n",
    "### Conda Environment\n",
    "\n",
    "`conda activate heartlab`\n",
    "\n",
    "### Data Sources\n",
    "\n",
    "`/data/camus/v5/{x_train/test/val | y_train/test/val}.npy`\n",
    "\n",
    "### References\n",
    "\n",
    "CS380-Prepare-Data.ipynb, CS380-UNet-Evaluate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "import datetime\n",
    "\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from tensorflow.keras import Model\n",
    "from tensorflow.keras.metrics import MeanIoU\n",
    "from tensorflow.keras.optimizers import Adam\n",
    "from tensorflow.keras.applications import VGG16\n",
    "from tensorflow.keras.callbacks import TensorBoard\n",
    "from tensorflow.keras.layers import Input, Conv2D, Conv2DTranspose, Dropout, Concatenate\n",
    "\n",
    "sys.path.append('../segmentation_utils')\n",
    "\n",
    "from segmentation_utils import jaccard_distance_loss, dice_loss, weighted_categorical_crossentropy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###### Import Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = \"/srv/jupyterhub_notebooks/data/camus/v5/\"\n",
    "\n",
    "x_train, x_val = np.load(path + 'x_train.npy'), np.load(path + 'x_val.npy')\n",
    "y_train, y_val = np.load(path + 'y_train.npy'), np.load(path + 'y_val.npy')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###### Hyperparameters\n",
    "\n",
    "Make sure to recompile the whole model after changing these"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "bottleneck_dropout_thresh = 0.5\n",
    "adam_optimizer = Adam(learning_rate=1e-5)\n",
    "freeze_encoder_weights = False # If True, encoder weights ARE frozen. If False, encoder weights ARE NOT frozen.\n",
    "n_epochs=50\n",
    "batch_size=16\n",
    "loss_function= weighted_categorical_crossentropy(np.array([0.5, 2, 2, 2]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###### Build Encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_layer = Input(shape=(256,256,1))\n",
    "input_concat = Concatenate()([input_layer, input_layer, input_layer])\n",
    "\n",
    "encoder_model = VGG16(input_tensor=input_concat, include_top=False, weights='imagenet', pooling=None)\n",
    "\n",
    "# Cut off the last conv block\n",
    "\n",
    "encoder_model = Model(inputs=input_layer, outputs=encoder_model.get_layer(\"block4_pool\").output)\n",
    "\n",
    "# Set weights frozen\n",
    "\n",
    "for layer in encoder_model.layers:\n",
    "    layer.trainable = !freeze_encoder_weights"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###### Build Bottleneck and Decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# bottleneck\n",
    "\n",
    "bottleneck1_1 = Dropout(bottleneck_dropout_thresh)(Conv2D(1024, (3,3), activation='relu', padding='same')(encoder_model.output))\n",
    "bottleneck1_2 = Dropout(bottleneck_dropout_thresh)(Conv2D(1024, (3,3), activation='relu', padding='same')(bottleneck1_1))\n",
    "\n",
    "# up block 1\n",
    "\n",
    "upconv1_1 = Conv2DTranspose(512, 2, (2,2), padding='same', activation='relu')(bottleneck1_2)\n",
    "upconv1_2 = Concatenate(axis=3)(inputs=[upconv1_1, encoder_model.get_layer(name='block4_conv3').output])\n",
    "upconv1_3 = Conv2D(512, (3,3), activation='relu', padding='same')(upconv1_2)\n",
    "upconv1_4 = Conv2D(512, (3,3), activation='relu', padding='same')(upconv1_3)\n",
    "\n",
    "# up block 2\n",
    "\n",
    "upconv2_1 = Conv2DTranspose(256, 2, (2,2), padding='same', activation='relu')(upconv1_4)\n",
    "upconv2_2 = Concatenate(axis=3)(inputs=[upconv2_1, encoder_model.get_layer(name='block3_conv3').output])\n",
    "upconv2_3 = Conv2D(256, (3,3), activation='relu', padding='same')(upconv2_2)\n",
    "upconv2_4 = Conv2D(256, (3,3), activation='relu', padding='same')(upconv2_3)\n",
    "\n",
    "# up block 3\n",
    "\n",
    "upconv3_1 = Conv2DTranspose(128, 2, (2,2), padding='same', activation='relu')(upconv2_4)\n",
    "upconv3_2 = Concatenate(axis=3)(inputs=[upconv3_1, encoder_model.get_layer(name='block2_conv2').output])\n",
    "upconv3_3 = Conv2D(128, (3,3), activation='relu', padding='same')(upconv3_2)\n",
    "upconv3_4 = Conv2D(128, (3,3), activation='relu', padding='same')(upconv3_3)\n",
    "\n",
    "# up block 4\n",
    "\n",
    "upconv4_1 = Conv2DTranspose(64, 2, (2,2), padding='same', activation='relu')(upconv3_4)\n",
    "upconv4_2 = Concatenate(axis=3)(inputs=[upconv4_1, encoder_model.get_layer(name='block1_conv2').output])\n",
    "upconv4_3 = Conv2D(64, (3,3), activation='relu', padding='same')(upconv4_2)\n",
    "upconv4_4 = Conv2D(64, (3,3), activation='relu', padding='same')(upconv4_3)\n",
    "\n",
    "# output layer\n",
    "output_layer = Conv2D(4, (1,1), activation='softmax')(upconv4_4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###### Setup TensorBoard logging"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_name = datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
    "log_dir = \"logs/fit/\" + run_name\n",
    "\n",
    "tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###### Compile final model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Model(inputs=input_layer, outputs=output_layer)\n",
    "\n",
    "model.compile(optimizer=adam_optimizer, \n",
    "              loss=loss_function, \n",
    "              metrics=[MeanIoU(num_classes=4)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###### Train model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.fit(\n",
    "    x=x_train,\n",
    "    y=y_train,\n",
    "    batch_size=batch_size,\n",
    "    epochs=n_epochs,\n",
    "    validation_data=(x_val, y_val),\n",
    "    callbacks=[tensorboard_callback]\n",
    ")\n",
    "\n",
    "model.save('models/' + run_name + '.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:heartlab]",
   "language": "python",
   "name": "conda-env-heartlab-py"
  },
  "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.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}