u-net-all_classes.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### U-Net architecture\n",
    "\n",
    "See [here](https://github.com/jocicmarko/ultrasound-nerve-segmentation/blob/master/train.py#L19) for code and [here](https://arxiv.org/pdf/1505.04597.pdf) for the original literature."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def restart_kernel(restart=False):\n",
    "    import IPython\n",
    "    app = IPython.Application.instance()\n",
    "    if restart:\n",
    "        app.kernel.do_shutdown(True)\n",
    "\n",
    "restart_kernel(False)\n",
    "\n",
    "#import os\n",
    "#os.environ[\"THEANO_FLAGS\"] = \"device=cpu\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from keras.models import Model\n",
    "from keras.layers import Reshape, Input, merge, Activation, Convolution2D, MaxPooling2D, UpSampling2D, ZeroPadding2D, Cropping2D, BatchNormalization, Dropout\n",
    "from keras.optimizers import Adadelta, Adam\n",
    "from keras.callbacks import CSVLogger, ModelCheckpoint, EarlyStopping, TensorBoard, ReduceLROnPlateau\n",
    "from keras import backend as K\n",
    "from keras.regularizers import l2\n",
    "from keras.constraints import maxnorm\n",
    "from keras.layers.advanced_activations import LeakyReLU\n",
    "\n",
    "K.set_image_dim_ordering('th')  # Theano dimension ordering in this code\n",
    "# \"tf\" assumes (rows, cols, channels) while \"th\" assumes (channels, rows, cols)\n",
    "# Possibly change this around natively in the data so the backend doesn't have to switch them\n",
    "# Only necessary if I use TF!\n",
    "\n",
    "from matplotlib import pyplot as plt\n",
    "from pushbullet import Pushbullet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Pushbullet notifier\n",
    "def push(title='Done!',text=''):\n",
    "    Pushbullet('o.YFPNNPfGRekivaCGHa4qMSgjZt8zJ6FL').devices[0].push_note(title,text)\n",
    "    \n",
    "# Import the training data\n",
    "def import_data(img_h,img_w):\n",
    "    x = np.load('./data/training-data/train_imgs_{}x{}_augmented_array.npy'.format(img_h,img_w),'r')\n",
    "    y = np.load('./data/training-data/train_msks_{}x{}_augmented_array.npy'.format(img_h,img_w),'r')\n",
    "    #y = y.reshape((y.shape[0],y.shape[1],y.shape[2]*y.shape[3]))\n",
    "    '''\n",
    "    Classes:\n",
    "    0 Buildings - large building, residential, non-residential, fuel storage facility, fortified building\n",
    "    1 Misc. Manmade structures \n",
    "    2 Road \n",
    "    3 Track - poor/dirt/cart track, footpath/trail\n",
    "    4 Trees - woodland, hedgerows, groups of trees, standalone trees\n",
    "    5 Crops - contour ploughing/cropland, grain (wheat) crops, row (potatoes, turnips) crops\n",
    "    6 Waterway \n",
    "    7 Standing water\n",
    "    8 Vehicle Large - large vehicle (e.g. lorry, truck,bus), logistics vehicle\n",
    "    9 Vehicle Small - small vehicle (car, van), motorbike\n",
    "    '''    \n",
    "    return x, y\n",
    "\n",
    "img_h,img_w = 160,160\n",
    "x, y = import_data(img_h,img_w)\n",
    "\n",
    "# Increment the counter\n",
    "def counter():\n",
    "    run = np.load('./data/misc/run_counter.npy')\n",
    "    run += 1\n",
    "    np.save('./data/misc/run_counter.npy',run)\n",
    "    return run\n",
    "run = counter()\n",
    "\n",
    "# Set the counter to a specific value\n",
    "def set_counter(run):\n",
    "    run = run\n",
    "    np.save('./data/misc/run_counter.npy',run)\n",
    "    return run\n",
    "# Uncomment the next line to manually set the counter if something goes wrong\n",
    "run = set_counter(101)\n",
    "print('This is run # %i' %run)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "def compiler(img_rows = x.shape[2],img_cols = x.shape[3],\n",
    "            nfilters = 64,activation = 'relu',init = 'he_normal',\n",
    "            lr=1.0,decay=0.0,momentum=0.0, nesterov=False,reg=0.01,p=[0.2,0.2,0.2,0.2,0.2,0.2,0.2,0.2,0.2]):\n",
    "    \n",
    "    weights = np.load('./data/misc/weights.npy')\n",
    "    \n",
    "    def weighted_binary_crossentropy(y_pred, y_true, weights=weights,epsilon=1e-7):\n",
    "        y_pred = K.clip(y_pred, epsilon, 1 - epsilon)\n",
    "        entropy = (-y_true * K.log(y_pred) - (1 - y_true) * K.log(1 - y_pred)) * weights[:,np.newaxis,np.newaxis]\n",
    "        return entropy\n",
    "    \n",
    "    def jaccard(y_true, y_pred,smooth=1.):\n",
    "        y_true_f = K.flatten(y_true)\n",
    "        y_pred_f = K.flatten(y_pred)\n",
    "        intersection = K.sum(y_true_f * y_pred_f)\n",
    "        return (intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) - intersection + smooth)\n",
    "    \n",
    "    def jaccard_0(y_true, y_pred,smooth=1.):\n",
    "        y_true_f = K.flatten(y_true[0])\n",
    "        y_pred_f = K.flatten(y_pred[0])\n",
    "        intersection = K.sum(y_true_f * y_pred_f)\n",
    "        return (intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) - intersection + smooth)\n",
    "    \n",
    "    def jaccard_1(y_true, y_pred,smooth=1.):\n",
    "        y_true_f = K.flatten(y_true[1])\n",
    "        y_pred_f = K.flatten(y_pred[1])\n",
    "        intersection = K.sum(y_true_f * y_pred_f)\n",
    "        return (intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) - intersection + smooth)\n",
    "    \n",
    "    def jaccard_2(y_true, y_pred,smooth=1.):\n",
    "        y_true_f = K.flatten(y_true[2])\n",
    "        y_pred_f = K.flatten(y_pred[2])\n",
    "        intersection = K.sum(y_true_f * y_pred_f)\n",
    "        return (intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) - intersection + smooth)\n",
    "    \n",
    "    def jaccard_3(y_true, y_pred,smooth=1.):\n",
    "        y_true_f = K.flatten(y_true[3])\n",
    "        y_pred_f = K.flatten(y_pred[3])\n",
    "        intersection = K.sum(y_true_f * y_pred_f)\n",
    "        return (intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) - intersection + smooth)\n",
    "    \n",
    "    def jaccard_4(y_true, y_pred,smooth=1.):\n",
    "        y_true_f = K.flatten(y_true[4])\n",
    "        y_pred_f = K.flatten(y_pred[4])\n",
    "        intersection = K.sum(y_true_f * y_pred_f)\n",
    "        return (intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) - intersection + smooth)\n",
    "    \n",
    "    def jaccard_5(y_true, y_pred,smooth=1.):\n",
    "        y_true_f = K.flatten(y_true[5])\n",
    "        y_pred_f = K.flatten(y_pred[5])\n",
    "        intersection = K.sum(y_true_f * y_pred_f)\n",
    "        return (intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) - intersection + smooth)\n",
    "    \n",
    "    def jaccard_6(y_true, y_pred,smooth=1.):\n",
    "        y_true_f = K.flatten(y_true[6])\n",
    "        y_pred_f = K.flatten(y_pred[6])\n",
    "        intersection = K.sum(y_true_f * y_pred_f)\n",
    "        return (intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) - intersection + smooth)\n",
    "    \n",
    "    def jaccard_7(y_true, y_pred,smooth=1.):\n",
    "        y_true_f = K.flatten(y_true[7])\n",
    "        y_pred_f = K.flatten(y_pred[7])\n",
    "        intersection = K.sum(y_true_f * y_pred_f)\n",
    "        return (intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) - intersection + smooth)\n",
    "    \n",
    "    def jaccard_8(y_true, y_pred,smooth=1.):\n",
    "        y_true_f = K.flatten(y_true[8])\n",
    "        y_pred_f = K.flatten(y_pred[8])\n",
    "        intersection = K.sum(y_true_f * y_pred_f)\n",
    "        return (intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) - intersection + smooth)\n",
    "    \n",
    "    def jaccard_9(y_true, y_pred,smooth=1.):\n",
    "        y_true_f = K.flatten(y_true[9])\n",
    "        y_pred_f = K.flatten(y_pred[9])\n",
    "        intersection = K.sum(y_true_f * y_pred_f)\n",
    "        return (intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) - intersection + smooth)\n",
    "    \n",
    "    def Conv2DReluBatchNorm(n_filter, w_filter, h_filter, inputs, activation, init='he_uniform',dropout=0.2):\n",
    "        # Batch norm after activation / leakyrelu\n",
    "        #return BatchNormalization(mode=2, axis=1)(LeakyReLU()((Convolution2D(n_filter, w_filter, h_filter, border_mode='same',init=init,W_regularizer=l2(reg),W_constraint = maxnorm(3))(inputs))))\n",
    "        \n",
    "        # Batch norm before activation\n",
    "        #return LeakyReLU()(BatchNormalization(mode=0, axis=1)((Convolution2D(n_filter, w_filter, h_filter, border_mode='same',init=init,W_regularizer=l2(reg),W_constraint = maxnorm(3))(inputs))))\n",
    "        \n",
    "        # Batch norm after activation / relu\n",
    "        return BatchNormalization(mode=2, axis=1)(Activation(activation=activation)((Convolution2D(n_filter, w_filter, h_filter, border_mode='same',init=init,W_regularizer=l2(reg),W_constraint = maxnorm(3))(inputs))))\n",
    "        \n",
    "    def up_conv(nfilters,filter_factor,inputs,init=init,activation=activation):\n",
    "        # No batch norm\n",
    "        #return LeakyReLU()(Convolution2D(nfilters*filter_factor, 2, 2, border_mode='same',init=init,W_regularizer=l2(reg),W_constraint = maxnorm(3))(UpSampling2D(size=(2, 2))(inputs)))\n",
    "        \n",
    "        # Batch norm after activation\n",
    "        #return BatchNormalization(mode=2, axis=1)(LeakyReLU()(Convolution2D(nfilters*filter_factor, 2, 2, border_mode='same',init=init,W_regularizer=l2(reg),W_constraint = maxnorm(3))(UpSampling2D(size=(2, 2))(inputs))))\n",
    "        \n",
    "        # Batch norm after activation / relu\n",
    "        return BatchNormalization(mode=2, axis=1)(Activation(activation=activation)(Convolution2D(nfilters*filter_factor, 2, 2, border_mode='same',init=init,W_regularizer=l2(reg),W_constraint = maxnorm(3))(UpSampling2D(size=(2, 2))(inputs))))\n",
    "\n",
    "    inputs = Input((20, img_rows, img_cols))\n",
    "    #padded = ZeroPadding2D(padding=(12,12))(inputs)\n",
    "    \n",
    "    conv1 = Conv2DReluBatchNorm(nfilters, 3, 3, inputs, activation=activation,init=init)\n",
    "    conv1 = Conv2DReluBatchNorm(nfilters, 3, 3, conv1, activation=activation,init=init)\n",
    "    pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)\n",
    "    pool1 = Dropout(p=p[0])(pool1)\n",
    "\n",
    "    conv2 = Conv2DReluBatchNorm(nfilters*2, 3, 3, pool1, activation=activation,init=init)\n",
    "    conv2 = Conv2DReluBatchNorm(nfilters*2, 3, 3, conv2, activation=activation,init=init)\n",
    "    pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)\n",
    "    pool2 = Dropout(p=p[1])(pool2)\n",
    "\n",
    "    conv3 = Conv2DReluBatchNorm(nfilters*4, 3, 3, pool2, activation=activation,init=init)\n",
    "    conv3 = Conv2DReluBatchNorm(nfilters*4, 3, 3, conv3, activation=activation,init=init)\n",
    "    pool3 = MaxPooling2D(pool_size=(2, 2))(conv3)\n",
    "    pool3 = Dropout(p=p[2])(pool3)\n",
    "\n",
    "    conv4 = Conv2DReluBatchNorm(nfilters*8, 3, 3, pool3, activation=activation,init=init)\n",
    "    conv4 = Conv2DReluBatchNorm(nfilters*8, 3, 3, conv4, activation=activation,init=init)\n",
    "    pool4 = MaxPooling2D(pool_size=(2, 2))(conv4)\n",
    "    pool4 = Dropout(p=p[3])(pool4)\n",
    "\n",
    "    conv5 = Conv2DReluBatchNorm(nfilters*16, 3, 3, pool4, activation=activation,init=init)\n",
    "    conv5 = Conv2DReluBatchNorm(nfilters*16, 3, 3, conv5, activation=activation,init=init)\n",
    "    conv5 = Dropout(p=p[4])(conv5)\n",
    "        \n",
    "    up6 = merge([up_conv(nfilters,8,conv5), conv4], mode='concat', concat_axis=1)\n",
    "    conv6 = Conv2DReluBatchNorm(nfilters*8, 3, 3, up6, activation=activation,init=init)\n",
    "    conv6 = Conv2DReluBatchNorm(nfilters*8, 3, 3, conv6, activation=activation,init=init)\n",
    "    conv6 = Dropout(p=p[5])(conv6)\n",
    "\n",
    "    up7 = merge([up_conv(nfilters,4,conv6), conv3], mode='concat', concat_axis=1)\n",
    "    conv7 = Conv2DReluBatchNorm(nfilters*4, 3, 3, up7, activation=activation,init=init)\n",
    "    conv7 = Conv2DReluBatchNorm(nfilters*4, 3, 3, conv7, activation=activation,init=init)\n",
    "    conv7 = Dropout(p=p[6])(conv7)\n",
    "\n",
    "    up8 = merge([up_conv(nfilters,2,conv7), conv2], mode='concat', concat_axis=1)\n",
    "    conv8 = Conv2DReluBatchNorm(nfilters*2, 3, 3, up8, activation=activation,init=init)\n",
    "    conv8 = Conv2DReluBatchNorm(nfilters*2, 3, 3, conv8, activation=activation,init=init)\n",
    "    conv8 = Dropout(p=p[7])(conv8)\n",
    "\n",
    "    up9 = merge([up_conv(nfilters,1,conv8), conv1], mode='concat', concat_axis=1)\n",
    "    conv9 = Conv2DReluBatchNorm(nfilters, 3, 3, up9, activation=activation,init=init)\n",
    "    conv9 = Conv2DReluBatchNorm(nfilters, 3, 3, conv9, activation=activation,init=init)\n",
    "    conv9 = Dropout(p=p[8])(conv9)\n",
    "    \n",
    "    conv10 = Conv2DReluBatchNorm(10, 1, 1, conv9, activation='relu',init=init)\n",
    "    #cropped = Cropping2D(cropping=((12,12), (12,12)))(conv10)\n",
    "    output = Activation(activation='sigmoid')(conv10)\n",
    "    #output = Reshape((10,img_rows*img_cols))(output)\n",
    "    \n",
    "    model = Model(input=inputs, output=output)\n",
    "    \n",
    "    '''model.compile(optimizer=Adam(lr=lr,decay=decay), loss=weighted_binary_crossentropy,\n",
    "                  metrics=[jaccard,jaccard_0,jaccard_1,jaccard_2,jaccard_3,\n",
    "                           jaccard_4,jaccard_5,jaccard_6,jaccard_7,jaccard_8,jaccard_9])'''\n",
    "    \n",
    "    model.compile(optimizer=Adam(lr=lr,decay=decay), loss=weighted_binary_crossentropy,\n",
    "                  metrics=[jaccard])\n",
    "    \n",
    "    \n",
    "    return model\n",
    "\n",
    "p = [0]*9\n",
    "# p=[0.1,0.2,0.3,0.4,0.5,0.4,0.3,0.2,0.1] # current version\n",
    "#p=[0.2,0.3,0.4,0.5,0.5,0.5,0.4,0.3,0.2] # symmetric but more dropout\n",
    "#p=[0.2,0.2,0.3,0.3,0.4,0.4,0.5,0.5,0.6] # increasing\n",
    "\n",
    "model = compiler(img_rows=x.shape[2],img_cols=x.shape[3],\n",
    "            nfilters=64,activation='relu',init='he_normal',\n",
    "            lr=0.001,decay=0,momentum=0,reg=0,p=p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def trainer(model,fit=True,use_existing=False):\n",
    "    print('This is run # %i' %run)\n",
    "    \n",
    "    if use_existing:\n",
    "        model.load_weights('./data/weights/model_weights_all-classes_run_{}.hdf5'.format(run))\n",
    "        \n",
    "    if fit:\n",
    "        quitter = EarlyStopping(monitor='val_loss', min_delta=0.001, patience=50, verbose=1, mode='auto')\n",
    "        lrreducer = ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=20, verbose=1, mode='auto', epsilon=0.001, cooldown=2, min_lr=0)\n",
    "        model_checkpoint = ModelCheckpoint('./data/weights/model_weights_all-classes_run_{}.hdf5'.format(run), monitor='val_loss', save_best_only=True)\n",
    "        csvlogger = CSVLogger('./data/logs/training_log_run_'+str(run), separator=',', append=True)\n",
    "        # tensorboard = TensorBoard(log_dir='./data/logs/'+'tensorboard_all-classes-run_{:04d}'.format(run), histogram_freq=0, write_graph=True, write_images=False)\n",
    "        # tensorboard --logdir=data/logs\n",
    "        \n",
    "        model.fit(x, y,\n",
    "                  batch_size=5,\n",
    "                  nb_epoch=1000,\n",
    "                  verbose=1,\n",
    "                  shuffle=True,\n",
    "                  callbacks=[model_checkpoint,csvlogger],\n",
    "                  validation_split=0.2,\n",
    "                  initial_epoch=0)\n",
    "            \n",
    "    preds = model.predict(x, verbose=1)\n",
    "    np.save('./data/predictions/predictions_run_{}.npy'.format(run), preds)\n",
    "    \n",
    "    return model\n",
    "\n",
    "model = trainer(model,fit=True,use_existing=False)\n",
    "model.save('./data/models/u-net-complete-model-run_{}_all-classes.hdf5'.format(run))\n",
    "push('Training is done',\n",
    "     'Train loss: %f, train jaccard: %f, val loss %f, val jaccard%f' %(model.history.history['loss'][-1],model.history.history['jaccard'][-1],model.history.history['val_loss'][-1],model.history.history['val_jaccard'][-1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [default]",
   "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}