UNET_regression_sugarcane.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "_MJ4kW1pEhwP"
   },
   "source": [
    "# Setup software libraries\n",
    "\n",
    "Authenticate and import as necessary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "neIa46CpciXq"
   },
   "outputs": [],
   "source": [
    "# Cloud authentication.\n",
    "from google.colab import auth\n",
    "auth.authenticate_user()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use `%tensorflow_version 1.x` to keep using v1, otherwise the default tensorflow version is 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "8RnZzcYhcpsQ"
   },
   "outputs": [],
   "source": [
    "# Tensorflow setup.\n",
    "import tensorflow as tf\n",
    "%tensorflow_version 1.x\n",
    "tf.enable_eager_execution()\n",
    "print(tf.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "iT8ycmzClYwf"
   },
   "source": [
    "# Variables\n",
    "\n",
    "Declare the variables that will be in use throughout the notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "qKs6HuxOzjMl"
   },
   "source": [
    "## Specify the Cloud Storage Bucket"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "obDDH1eDzsch"
   },
   "outputs": [],
   "source": [
    "# INSERT YOUR BUCKET HERE:\n",
    "BUCKET = 'sugarcane-unet-model'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "wmfKLl9XcnGJ"
   },
   "source": [
    "## Define the loss function and the one of the metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "bi8lguY2voEz"
   },
   "outputs": [],
   "source": [
    "from tensorflow.python.keras import backend\n",
    "\n",
    "def dice_coeff(y_true, y_pred, smooth=1):\n",
    "    y_true_f = backend.flatten(y_true)\n",
    "    y_pred_f = backend.flatten(y_pred)\n",
    "    intersection = backend.sum(y_true_f * y_pred_f)\n",
    "    return (2. * intersection + smooth) / (backend.sum(y_true_f) + backend.sum(y_pred_f) + smooth)\n",
    "\n",
    "def dice_loss(y_true, y_pred):\n",
    "    loss = 1 - dice_coeff(y_true, y_pred)\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "bSHhmG5atJJG"
   },
   "source": [
    "## Set other global variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "psz7wJKalaoj"
   },
   "outputs": [],
   "source": [
    "from tensorflow.python.keras import losses\n",
    "from tensorflow.python.keras import metrics\n",
    "\n",
    "# Specify names of output locations in Cloud Storage.\n",
    "JOB_DIR = 'gs://' + BUCKET + '/modified_trainer_128_25_Adam_16'\n",
    "MODEL_DIR = JOB_DIR + '/model'\n",
    "LOGS_DIR = JOB_DIR + '/logs'\n",
    "DATA_DIR = 'gs://' + BUCKET + '/data_128_NormalizedData'\n",
    "\n",
    "# Pre-computed training and eval data.\n",
    "TRAINING_BASE = 'training_patches'\n",
    "EVAL_BASE = 'eval_patches'\n",
    "\n",
    "# Specify inputs to the model and the response variable.\n",
    "opticalBands = ['blue', 'green', 'red', 'nir', 'swir1', 'swir2']\n",
    "thermalBands = ['thermal'] # currently has some issues so not including\n",
    "thermalBands = []\n",
    "otherBands = [\n",
    "  'p20_green', 'p20_nir', 'p20_blue', 'p20_red', 'p20_swir1', 'p20_swir2',\n",
    "  'p80_green', 'p80_nir', 'p80_blue', 'p80_red', 'p80_swir1', 'p80_swir2',\n",
    "]\n",
    "indices = [\n",
    "  'NDVI', 'EVI', 'SAVI', 'IBI', 'GRatio',\n",
    "  'brightness', 'greenness', 'wetness', 'fourth', 'fifth', 'sixth',\n",
    "]\n",
    "BANDS = opticalBands + thermalBands + otherBands + indices\n",
    "RESPONSE = 'sugarcane'\n",
    "FEATURES = BANDS + [RESPONSE]\n",
    "\n",
    "# Specify the size and shape of patches expected by the model.\n",
    "KERNEL_SIZE = 128\n",
    "KERNEL_SHAPE = [KERNEL_SIZE, KERNEL_SIZE]\n",
    "COLUMNS = [\n",
    "  tf.io.FixedLenFeature(shape=KERNEL_SHAPE, dtype=tf.float32) for k in FEATURES\n",
    "]\n",
    "FEATURES_DICT = dict(zip(FEATURES, COLUMNS))\n",
    "\n",
    "# Sizes of the training and evaluation datasets.\n",
    "TRAIN_SIZE = 16000\n",
    "EVAL_SIZE = 8000\n",
    "\n",
    "# Specify model training parameters.\n",
    "BATCH_SIZE = 16\n",
    "EPOCHS = 25\n",
    "BUFFER_SIZE = 2000\n",
    "OPTIMIZER = 'Adam'\n",
    "LOSS = dice_loss\n",
    "METRICS = [\n",
    "    metrics.get('RootMeanSquaredError'),\n",
    "    metrics.get('MeanAbsoluteError'),\n",
    "    metrics.get('Accuracy'),\n",
    "    dice_coeff,\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "rWXrvBE4607G"
   },
   "source": [
    "# Training data\n",
    "\n",
    "Load the data exported from Earth Engine into a `tf.data.Dataset`.  The following are helper functions for that."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "WWZ0UXCVMyJP"
   },
   "outputs": [],
   "source": [
    "def parse_tfrecord(example_proto):\n",
    "  \"\"\"The parsing function.\n",
    "  Read a serialized example into the structure defined by FEATURES_DICT.\n",
    "  Args:\n",
    "    example_proto: a serialized Example.\n",
    "  Returns: \n",
    "    A dictionary of tensors, keyed by feature name.\n",
    "  \"\"\"\n",
    "  return tf.io.parse_single_example(example_proto, FEATURES_DICT)\n",
    "\n",
    "\n",
    "def to_tuple(inputs):\n",
    "  \"\"\"Function to convert a dictionary of tensors to a tuple of (inputs, outputs).\n",
    "  Turn the tensors returned by parse_tfrecord into a stack in HWC shape.\n",
    "  Args:\n",
    "    inputs: A dictionary of tensors, keyed by feature name.\n",
    "  Returns: \n",
    "    A dtuple of (inputs, outputs).\n",
    "  \"\"\"\n",
    "  inputsList = [inputs.get(key) for key in FEATURES]\n",
    "  stacked = tf.stack(inputsList, axis=0)\n",
    "  # Convert from CHW to HWC\n",
    "  stacked = tf.transpose(stacked, [1, 2, 0])\n",
    "  return stacked[:,:,:len(BANDS)], stacked[:,:,len(BANDS):]\n",
    "\n",
    "\n",
    "def get_dataset(pattern):\n",
    "  \"\"\"Function to read, parse and format to tuple a set of input tfrecord files.\n",
    "  Get all the files matching the pattern, parse and convert to tuple.\n",
    "  Args:\n",
    "    pattern: A file pattern to match in a Cloud Storage bucket.\n",
    "  Returns: \n",
    "    A tf.data.Dataset\n",
    "  \"\"\"\n",
    "  glob = tf.gfile.Glob(pattern)\n",
    "  dataset = tf.data.TFRecordDataset(glob, compression_type='GZIP')\n",
    "  dataset = dataset.map(parse_tfrecord, num_parallel_calls=5)\n",
    "  dataset = dataset.map(to_tuple, num_parallel_calls=5)\n",
    "  return dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Xg1fa18336D2"
   },
   "source": [
    "Use the helpers to read in the training dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "rm0qRF0fAYcC"
   },
   "outputs": [],
   "source": [
    "def get_training_dataset():\n",
    "\t\"\"\"Get the preprocessed training dataset\n",
    "  Returns: \n",
    "    A tf.data.Dataset of training data.\n",
    "  \"\"\"\n",
    "\tglob = DATA_DIR + '/' + TRAINING_BASE + '*'\n",
    "\tdataset = get_dataset(glob)\n",
    "\tdataset = dataset.shuffle(BUFFER_SIZE).batch(BATCH_SIZE).repeat()\n",
    "\treturn dataset\n",
    "\n",
    "training = get_training_dataset()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "j-cQO5RL6vob"
   },
   "source": [
    "# Evaluation data\n",
    "\n",
    "Now do the same thing to get an evaluation dataset.  Note that unlike the training dataset, the evaluation dataset has a batch size of 1, is not repeated and is not shuffled."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ieKTCGiJ6xzo"
   },
   "outputs": [],
   "source": [
    "def get_eval_dataset():\n",
    "\t\"\"\"Get the preprocessed evaluation dataset\n",
    "  Returns: \n",
    "    A tf.data.Dataset of evaluation data.\n",
    "  \"\"\"\n",
    "\tglob = DATA_DIR + '/' + EVAL_BASE + '*'\n",
    "\tdataset = get_dataset(glob)\n",
    "\tdataset = dataset.batch(1).repeat()\n",
    "\treturn dataset\n",
    "\n",
    "evaluation = get_eval_dataset()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "OmvvOGg6sbJH"
   },
   "source": [
    "# Data Augmentation\n",
    "\n",
    "Performing the data augmentation to increase the number of training datasets "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "jz5gZp89slDM"
   },
   "outputs": [],
   "source": [
    "def upDownFlipInputs(inputs, labels):\n",
    "    return tf.image.flip_up_down(inputs), tf.image.flip_up_down(labels)\n",
    "\n",
    "def leftRightFlipInputs(inputs, labels):\n",
    "    return tf.image.flip_left_right(inputs), tf.image.flip_left_right(labels)\n",
    "\n",
    "def transposeInputs(inputs, labels):\n",
    "    upDown = tf.image.flip_up_down(inputs)\n",
    "    transpose = tf.image.flip_left_right(upDown)\n",
    "    upDownLabel = tf.image.flip_up_down(labels)\n",
    "    transposeLabel = tf.image.flip_left_right(upDownLabel)\n",
    "    return transpose, transposeLabel\n",
    "\n",
    "def rotateInputs(inputs, labels):        \n",
    "    return tf.image.rot90(inputs), tf.image.rot90(labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "8ubGMjrls4gQ"
   },
   "outputs": [],
   "source": [
    "training = training.concatenate(training.map(upDownFlipInputs))\\\n",
    "                   .concatenate(training.map(leftRightFlipInputs))\\\n",
    "                   .concatenate(training.map(transposeInputs))\\\n",
    "                   .concatenate(training.map(rotateInputs))\n",
    "\n",
    "training = training.shuffle(buffer_size = BATCH_SIZE * 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "9JIE7Yl87lgU"
   },
   "source": [
    "# Model\n",
    "\n",
    "Here we use the Keras implementation of the U-Net model.  The U-Net model takes any pixel patches as input and outputs per-pixel class probability, label or a continuous output. Add in the regularizer and/or dropout as required. Currently they are commented out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "wsnnnz56yS3l"
   },
   "outputs": [],
   "source": [
    "from tensorflow.python.keras import layers\n",
    "from tensorflow.python.keras import losses\n",
    "from tensorflow.python.keras import models\n",
    "from tensorflow.python.keras import metrics\n",
    "from tensorflow.python.keras import optimizers\n",
    "from tensorflow.python.keras import regularizers\n",
    "\n",
    "def conv_block(input_tensor, num_filters):\n",
    "\tencoder = layers.Conv2D(num_filters, (3, 3), padding='same')(input_tensor)\n",
    "\t#encoder = layers.Conv2D(num_filters, (3, 3), padding='same', kernel_regularizer=regularizers.l2(0.01))(input_tensor)\n",
    "\tencoder = layers.BatchNormalization()(encoder)\n",
    "\tencoder = layers.Activation('relu')(encoder)\n",
    "\t#encoder = layers.Dropout(0.2)(encoder)\n",
    "\n",
    "\tencoder = layers.Conv2D(num_filters, (3, 3), padding='same')(encoder)\n",
    "\t#encoder = layers.Conv2D(num_filters, (3, 3), padding='same', kernel_regularizer=regularizers.l2(0.01))(encoder)\n",
    "\tencoder = layers.BatchNormalization()(encoder)\n",
    "\tencoder = layers.Activation('relu')(encoder)\n",
    "\t#encoder = layers.Dropout(0.2)(encoder)\n",
    "\n",
    "    return encoder\n",
    "\n",
    "def encoder_block(input_tensor, num_filters):\n",
    "\tencoder = conv_block(input_tensor, num_filters)\n",
    "\tencoder_pool = layers.MaxPooling2D((2, 2), strides=(2, 2))(encoder)\n",
    "\treturn encoder_pool, encoder\n",
    "\n",
    "def decoder_block(input_tensor, concat_tensor, num_filters):\n",
    "\tdecoder = layers.Conv2DTranspose(num_filters, (2, 2), strides=(2, 2), padding='same')(input_tensor)\n",
    "\t#decoder = layers.Conv2DTranspose(num_filters, (2, 2), strides=(2, 2), padding='same', kernel_regularizer=regularizers.l2(0.01))(input_tensor)\n",
    "\tdecoder = layers.concatenate([concat_tensor, decoder], axis=-1)\n",
    "\tdecoder = layers.BatchNormalization()(decoder)\n",
    "\tdecoder = layers.Activation('relu')(decoder)\n",
    "    #decoder = layers.Dropout(0.2)(decoder)\n",
    "    \n",
    "    decoder = layers.Conv2D(num_filters, (3, 3), padding='same')(decoder)\n",
    "\t#decoder = layers.Conv2D(num_filters, (3, 3), padding='same', kernel_regularizer=regularizers.l2(0.01))(decoder)\n",
    "\tdecoder = layers.BatchNormalization()(decoder)\n",
    "\tdecoder = layers.Activation('relu')(decoder)\n",
    "    #decoder = layers.Dropout(0.2)(decoder)\n",
    "\n",
    "    decoder = layers.Conv2D(num_filters, (3, 3), padding='same')(decoder)\n",
    "\t#decoder = layers.Conv2D(num_filters, (3, 3), padding='same', kernel_regularizer=regularizers.l2(0.01))(decoder)\n",
    "\tdecoder = layers.BatchNormalization()(decoder)\n",
    "\tdecoder = layers.Activation('relu')(decoder)\n",
    "    #decoder = layers.Dropout(0.2)(decoder)\n",
    "\n",
    "    return decoder\n",
    "\n",
    "def get_model():\n",
    "\tinputs = layers.Input(shape=[None, None, len(BANDS)])\n",
    "\tencoder0_pool, encoder0 = encoder_block(inputs, 16)\n",
    "\tencoder1_pool, encoder1 = encoder_block(encoder0_pool, 32)\n",
    "\tencoder2_pool, encoder2 = encoder_block(encoder1_pool, 64)\n",
    "\tencoder3_pool, encoder3 = encoder_block(encoder2_pool, 128)\n",
    "\tcenter = conv_block(encoder3_pool, 256) # center\n",
    "\tdecoder3 = decoder_block(center, encoder3, 128)\n",
    "\tdecoder2 = decoder_block(decoder3, encoder2, 64)\n",
    "\tdecoder1 = decoder_block(decoder2, encoder1, 32)\n",
    "\tdecoder0 = decoder_block(decoder1, encoder0, 16)\n",
    "\toutputs = layers.Conv2D(1, (1, 1), activation='sigmoid')(decoder0)\n",
    "\n",
    "\tmodel = models.Model(inputs=[inputs], outputs=[outputs])\n",
    "\n",
    "\tmodel.compile(\n",
    "\t\toptimizer=optimizers.get(OPTIMIZER), \n",
    "\t\tloss=LOSS,\n",
    "\t\tmetrics=METRICS\n",
    "\t)\n",
    "\n",
    "\treturn model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "mIu9oDW9yACs"
   },
   "outputs": [],
   "source": [
    "model = get_model()\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "uu_E7OTDBCoS"
   },
   "source": [
    "# Training the model\n",
    "\n",
    "You train a Keras model by calling `.fit()` on it.  Here we're going to train for EPOCHS epochs, which is suitable for demonstration purposes.\n",
    "Some callbacks have been commented out. Use them as required."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "NzzaWxOhSxBy"
   },
   "outputs": [],
   "source": [
    "m = get_model()\n",
    "m.fit(\n",
    "    x=training,\n",
    "    epochs=EPOCHS, \n",
    "    steps_per_epoch=int(TRAIN_SIZE / BATCH_SIZE), \n",
    "    validation_data=evaluation,\n",
    "    validation_steps=EVAL_SIZE,\n",
    "    callbacks=[\n",
    "        tf.keras.callbacks.TensorBoard(LOGS_DIR),\n",
    "        #tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=8, verbose=1, min_delta=1e-4),\n",
    "        #tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=4, verbose=1)\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Save the model for future use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.contrib.saved_model.save_keras_model(m, MODEL_DIR)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "machine_shape": "hm",
   "name": "Modified_UNET_regression_sugarcane.ipynb",
   "private_outputs": true,
   "provenance": [],
   "toc_visible": true
  },
  "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.5.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}