Shoreline_detection_train.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train a HED CNN for shoreline detection\n",
    "This shoreline detector uses a Holistically-Nested Edge Detection framework (Xie and Tu, 2015) -http://openaccess.thecvf.com/content_iccv_2015/papers/Xie_Holistically-Nested_Edge_Detection_ICCV_2015_paper.pdf.\n",
    "\n",
    "This notebook contains the code to train a HED CNN.\n",
    "\n",
    "Please note, at the current time the image and shoreline dataset used to train the network is not publicly available. As such, unless you have your own dataset you should try running the \"HED_edge_detection_predict_unseen\" notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Automatic pdb calling has been turned ON\n"
     ]
    }
   ],
   "source": [
    "# magic\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%pdb 1\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# imports\n",
    "from functions.pytorch_models import hed_cnn, Trainer, pretrained_weights\n",
    "from functions.data_preprocessing import load_images, augment_images_kp, mask_to_uv\n",
    "from functions.data_preprocessing import load_train_test_imagedata, save_train_test_imagedata\n",
    "from functions.data_visualisation import plot_predictions, plot_refined_predictions\n",
    "\n",
    "import scipy.io as sio\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "import os\n",
    "\n",
    "import torch\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "import imgaug as ia\n",
    "import imgaug.augmenters as iaa\n",
    "\n",
    "from ipywidgets import interact, fixed, IntSlider, FloatSlider, interact_manual"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup HED model\n",
    "1. Import the model from model.py\n",
    "2. Define and apply training parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "basePath = './data/'\n",
    "partition, labels = load_train_test_imagedata(basePath)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# load the model\n",
    "applyWeights = True\n",
    "modelSave = 'A'\n",
    "\n",
    "#pretrained model\n",
    "weightsPath = './pytorch/pretrained_models/vgg16-397923af.pth'\n",
    "\n",
    "hedModel = hed_cnn()\n",
    "hedModel = pretrained_weights(hedModel, weightsPath, applyWeights)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train the model\n",
    "1. Specify and train the model\n",
    "2. Save the model\n",
    "2. Make predictions for training and test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# setup training\n",
    "modelParams = {\n",
    "    'epochs': 25,\n",
    "    'batchSize': 8,\n",
    "    'lr': 3e-4,\n",
    "    'lrDecay': 5e-1,\n",
    "    'lossFunction': 'weightedBCE', # weightedBCEReg\n",
    "    'cuda': False,\n",
    "    'basePath': basePath,\n",
    "}\n",
    "\n",
    "fuse_params = list(map(id, hedModel.fuse.parameters()))\n",
    "base_params = filter(lambda p: id(p) not in fuse_params, hedModel.parameters())\n",
    "\n",
    "modelParams['optimiser'] = torch.optim.Adam(filter(lambda p: p.requires_grad, hedModel.parameters()), lr=modelParams['lr'])\n",
    "\n",
    "# initialize trainer class\n",
    "trainer = Trainer(hedModel, partition, labels, modelParams)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "#train\n",
    "history = trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the model\n",
    "trainedHedModel = trainer.model\n",
    "torch.save(trainedHedModel.state_dict(), './models/shorelineDetectModel_{}.pt'.format(modelSave))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# predict\n",
    "trainX, trainPred, trainY = trainer.predict('train')\n",
    "valX, valPred, valY  = trainer.predict('validation')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## View model output\n",
    "For a small sample (one batch) of the training and validation datasets\n",
    "\n",
    "### Raw output\n",
    "From each layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# training output\n",
    "print('Training output...')\n",
    "interact(plot_predictions,\n",
    "         prntNum =IntSlider(\n",
    "             value=0,\n",
    "             min=0,\n",
    "             step=1,\n",
    "             max=trainX.shape[0]-1,\n",
    "             continuous_update=False,\n",
    "         ),\n",
    "         dataX=fixed(trainX),\n",
    "         dataY=fixed(trainY),\n",
    "         dataPred=fixed(trainPred),\n",
    "         jj=IntSlider(\n",
    "             value=5,\n",
    "             min=0,\n",
    "             step=1,\n",
    "             max=5,\n",
    "             continuous_update=False,\n",
    "         ),\n",
    "         thres=FloatSlider(\n",
    "             value=0.5,\n",
    "             min=0.05,\n",
    "             step=0.05,\n",
    "             max=0.95,\n",
    "             continuous_update=False,\n",
    "         ),\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Validation output...')\n",
    "prntNum = 3\n",
    "interact(plot_predictions,\n",
    "         prntNum =IntSlider(\n",
    "             value=0,\n",
    "             min=0,\n",
    "             step=1,\n",
    "             max=valX.shape[0]-1,\n",
    "             continuous_update=False,\n",
    "         ),\n",
    "         dataX=fixed(valX),\n",
    "         dataY=fixed(valY),\n",
    "         dataPred=fixed(valPred),\n",
    "         jj=IntSlider(\n",
    "             value=5,\n",
    "             min=0,\n",
    "             step=1,\n",
    "             max=5,\n",
    "             continuous_update=False,\n",
    "         ),\n",
    "         thres=FloatSlider(\n",
    "             value=0.5,\n",
    "             min=0.05,\n",
    "             step=0.05,\n",
    "             max=1,\n",
    "             continuous_update=False,\n",
    "         ),\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Final output\n",
    "From weighted combination."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Training output...')\n",
    "interact(plot_refined_predictions,\n",
    "         prntNum =IntSlider(\n",
    "             value=0,\n",
    "             min=0,\n",
    "             step=1,\n",
    "             max=trainX.shape[0]-1,\n",
    "             continuous_update=False,\n",
    "         ),\n",
    "         dataX=fixed(trainX),\n",
    "         dataY=fixed(trainY),\n",
    "         dataPred=fixed(trainPred),\n",
    "         thres=FloatSlider(\n",
    "             value=0.7,\n",
    "             min=0.05,\n",
    "             step=0.05,\n",
    "             max=0.95,\n",
    "             continuous_update=False,\n",
    "         ),\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Validation output...')\n",
    "interact(plot_refined_predictions,\n",
    "         prntNum =IntSlider(\n",
    "             value=0,\n",
    "             min=0,\n",
    "             step=1,\n",
    "             max=valX.shape[0]-1,\n",
    "             continuous_update=False,\n",
    "         ),\n",
    "         dataX=fixed(valX),\n",
    "         dataY=fixed(valY),\n",
    "         dataPred=fixed(valPred),\n",
    "         thres=FloatSlider(\n",
    "             value=0.7,\n",
    "             min=0.05,\n",
    "             step=0.05,\n",
    "             max=0.95,\n",
    "             continuous_update=False,\n",
    "         ),\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (plaid)",
   "language": "python",
   "name": "plaid"
  },
  "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}