Timestack_Swash_Front_Tracing_SAI2020_Unet_MT.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Swash Timestack Tracer\n",
    "This notebook is used to train a U-Net to find the position of the swash front through a timestack. Once segmentation has been done a thresholding technique similar to Otsu's method is used to defined the shoreline.\n",
    "\n",
    "# Your Task:\n",
    "Train the best U-net for tracing the swash front. The default structure used is Wei Dai's example which you can build from by tuning hyperparameters.\n",
    "\n",
    "Code cell 1 imports all the libraries and stores a function used for evaluating your trained model. You don't need to change anything in this code cell.\n",
    "\n",
    "Code cell 2 loads the training images into the notebook and adds any data augmentation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import libraries\n",
    "import tensorflow as tf\n",
    "from tensorflow.keras.models import *\n",
    "from tensorflow.keras.layers import *\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import cv2\n",
    "import os\n",
    "import shutil\n",
    "import random\n",
    "import csv\n",
    "from scipy.stats import norm\n",
    "\n",
    "\n",
    "# https://stackoverflow.com/questions/22579434/python-finding-the-intersection-point-of-two-gaussian-curves\n",
    "def solve(m1,m2,std1,std2):\n",
    "    a = 1/(2*std1**2) - 1/(2*std2**2)\n",
    "    b = m2/(std2**2) - m1/(std1**2)\n",
    "    c = m1**2 /(2*std1**2) - m2**2 / (2*std2**2) - np.log(std2/std1)\n",
    "    return np.roots([a,b,c])\n",
    "\n",
    "# function used to evaluate how good the model is\n",
    "def compare_shoreline(min_uprush, max_uprush, tf_model_version, IMG_SIZE, timestack_name, results_dir):\n",
    "    \n",
    "    folder_name = \"test_timestacks/\"\n",
    "\n",
    "    image_range = max_uprush - min_uprush\n",
    "    \n",
    "    timestack = cv2.imread(folder_name + timestack_name)\n",
    "    timestack_focus = timestack[:, min_uprush:max_uprush]\n",
    "    timestack_labelled = cv2.imread(folder_name + timestack_name + \"_output.png\")\n",
    "    timestack_annotated = timestack_labelled.copy()\n",
    "    mask_annotated = cv2.imread(folder_name + timestack_name + \"_output.png\", cv2.IMREAD_GRAYSCALE)\n",
    "    \n",
    "    timestack_count = 0\n",
    "\n",
    "    IMG_SIZE_WID = int(IMG_SIZE) # can change this if needed\n",
    "    tf_images = []\n",
    "    if tf_model_version == \"michaelThompson_imgSize_64\":\n",
    "        # this way of reading/saving models worked for an older version of tensorflow 2.x\n",
    "        tf_model_xc = tf.keras.models.load_model(tf_model_version + \".model\")\n",
    "    else:\n",
    "        tf_model_xc = tf.keras.models.load_model(tf_model_version + \"_model\")\n",
    "\n",
    "    vertical_coordinates = []\n",
    "    horizontal_coordinates = []\n",
    "    full_height, full_width, _ = timestack_focus.shape\n",
    "    snap_width = full_width\n",
    "    snap_height = full_width\n",
    "    max_snap_start = full_height - snap_height\n",
    "    snap_start = 0\n",
    "    snap_counter = 0\n",
    "    # split the image up into square sections - overlap the last section\n",
    "    while snap_start < max_snap_start:\n",
    "\n",
    "        # get square image\n",
    "        snapshot = timestack_focus[snap_start:(snap_start + snap_height)]\n",
    "        \n",
    "        # prepare as tf input\n",
    "        snapshot_resize = cv2.resize(snapshot, (IMG_SIZE_WID, IMG_SIZE))\n",
    "        snapshot_resize_nm = np.array(snapshot_resize) / 255\n",
    "        snapshot_input = snapshot_resize_nm.reshape(-1, IMG_SIZE, IMG_SIZE_WID, 3)\n",
    "        \n",
    "        # put into model\n",
    "        mask_test = tf_model_xc.predict([snapshot_input])[0]\n",
    "        mask_test = cv2.resize(mask_test, (snap_width, snap_height))\n",
    "        \n",
    "        # add mask to mask image\n",
    "        mask_annotated[snap_start:(snap_start + snap_height), min_uprush:max_uprush] = mask_test * 255\n",
    "        \n",
    "        # apply light and dark gaussian interesection method to extract shoreline points\n",
    "        # use something similar to otsu's method for each row to draw a boundary line\n",
    "        for j in range(0, len(mask_test), 1):\n",
    "            pixel_row = mask_test[j]\n",
    "\n",
    "            light_pixels = []\n",
    "            dark_pixels = []\n",
    "            for k in range(0, len(pixel_row), 1):\n",
    "                if pixel_row[k] >= 0.5:\n",
    "                    light_pixels.extend([k]*int(pixel_row[k]*255))\n",
    "                else:\n",
    "                    dark_pixels.extend([k]*int((0.5 - pixel_row[k])*255))\n",
    "\n",
    "            # find mean and variance of pixels\n",
    "            light_mean, light_std = norm.fit(light_pixels)\n",
    "            dark_mean, dark_std = norm.fit(dark_pixels)\n",
    "\n",
    "            try:\n",
    "                # get the interect of both gaussians\n",
    "                intersect_i = solve(light_mean, dark_mean, light_std, dark_std)\n",
    "\n",
    "                # if intersect has more than 1 point, choose the point closest to the middle of the two means\n",
    "                if len(intersect_i) == 2:\n",
    "                    mean_middle = (light_mean + dark_mean)/2\n",
    "                    intersect = min(intersect_i, key=lambda x:abs(x-mean_middle))\n",
    "                # if 1 point thats great!\n",
    "                elif len(intersect_i) == 1:\n",
    "                    intersect = intersect_i\n",
    "                # if 0 or 3 or more points, simply choose middle of two means\n",
    "                else:\n",
    "                    intersect = (light_mean + dark_mean)/2\n",
    "            except:\n",
    "                # just use middle of means\n",
    "                if np.isnan(light_mean) or np.isnan(dark_mean):\n",
    "                    intersect = 0\n",
    "                else:\n",
    "                    intersect = (light_mean + dark_mean)/2\n",
    "            \n",
    "            # annotate over timestack\n",
    "            cv2.circle(timestack_annotated, (min_uprush + int(intersect), snap_start + j), 1, (255, 0, 0))\n",
    "            \n",
    "            # save vertical and horizontal coordinates for comparison with labels.\n",
    "            horizontal_coordinates.append(snap_start + j)\n",
    "            vertical_coordinates.append(min_uprush + int(intersect))\n",
    "        \n",
    "        timestack_count = timestack_count + 1\n",
    "        snap_start = snap_start + int(snap_height)\n",
    "        snap_counter = snap_counter + 1\n",
    "    \n",
    "    snap_end = snap_counter*snap_height\n",
    "    # after while loop is finished, do one final snapshot that overlaps with last snapshot till the end.\n",
    "    # get square image\n",
    "    snapshot = timestack_focus[(full_height - snap_height):full_height]\n",
    "\n",
    "    # prepare as tf input\n",
    "    snapshot_resize = cv2.resize(snapshot, (IMG_SIZE_WID, IMG_SIZE))\n",
    "    snapshot_resize_nm = np.array(snapshot_resize) / 255\n",
    "    snapshot_input = snapshot_resize_nm.reshape(-1, IMG_SIZE, IMG_SIZE_WID, 3)\n",
    "\n",
    "    # put into model\n",
    "    mask_test = tf_model_xc.predict([snapshot_input])[0]\n",
    "    mask_test = cv2.resize(mask_test, (snap_width, snap_height))\n",
    "    \n",
    "    # add mask to mask image\n",
    "    mask_annotated[(full_height - snap_height):full_height, min_uprush:max_uprush] = mask_test * 255\n",
    "    \n",
    "    # apply light and dark gaussian interesection method to extract shoreline points\n",
    "    # use something similar to otsu's method for each row to draw a boundary line\n",
    "    for j in range(0, len(mask_test), 1):\n",
    "        pixel_row = mask_test[j]\n",
    "\n",
    "        light_pixels = []\n",
    "        dark_pixels = []\n",
    "        for k in range(0, len(pixel_row), 1):\n",
    "            if pixel_row[k] >= 0.5:\n",
    "                light_pixels.extend([k]*int(pixel_row[k]*255))\n",
    "            else:\n",
    "                dark_pixels.extend([k]*int((0.5 - pixel_row[k])*255))\n",
    "\n",
    "        # find mean and variance of pixels\n",
    "        light_mean, light_std = norm.fit(light_pixels)\n",
    "        dark_mean, dark_std = norm.fit(dark_pixels)\n",
    "\n",
    "        try:\n",
    "            # get the interect of both gaussians\n",
    "            intersect_i = solve(light_mean, dark_mean, light_std, dark_std)\n",
    "            \n",
    "            # if intersect has more than 1 point, choose the point closest to the middle of the two means\n",
    "            if len(intersect_i) == 2:\n",
    "                mean_middle = (light_mean + dark_mean)/2\n",
    "                intersect = min(intersect_i, key=lambda x:abs(x-mean_middle))\n",
    "            # if 1 point thats great!\n",
    "            elif len(intersect_i) == 1:\n",
    "                intersect = intersect_i\n",
    "            # if 0 or 3 or more points, simply choose middle of two means\n",
    "            else:\n",
    "                intersect = (light_mean + dark_mean)/2\n",
    "        except:\n",
    "            # just use middle of means\n",
    "            if np.isnan(light_mean) or np.isnan(dark_mean):\n",
    "                intersect = 0\n",
    "            else:\n",
    "                intersect = (light_mean + dark_mean)/2\n",
    "        \n",
    "        # only add to timestack and data if new data\n",
    "        if (full_height - snap_height + j) >= snap_end:\n",
    "            # annotate over timestack\n",
    "            cv2.circle(timestack_annotated, (min_uprush + int(intersect), full_height - snap_height + j), 1, (255, 0, 0))\n",
    "\n",
    "            # save vertical and horizontal coordinates for comparison with labels.\n",
    "            horizontal_coordinates.append(full_height - snap_height + j)\n",
    "            vertical_coordinates.append(min_uprush + int(intersect))\n",
    "        \n",
    "    # get csv values from labelled timestack\n",
    "    manual_csv_points_filename = folder_name + \"runup_data_test\" + timestack_name + \".csv\"\n",
    "    x_vals = []\n",
    "    t_vals = []\n",
    "    with open(manual_csv_points_filename, newline='') as csvfile:\n",
    "        points_reader = csv.reader(csvfile, delimiter=' ', quotechar='|')\n",
    "        initial_row = 5\n",
    "        initial_row_counter = 0\n",
    "        for row in points_reader:\n",
    "            if initial_row_counter > initial_row:\n",
    "                rows = row[0].split(',')\n",
    "                x_vals.append( int(rows[0][1:-1]) ) # need to remove [] from value and turn into a number\n",
    "                t_vals.append( int(rows[1][1:-1]) ) # need to remove [] from value and turn into a number\n",
    "            initial_row_counter = initial_row_counter + 1\n",
    "\n",
    "    print(\"red is hand picked, green is found by model\")\n",
    "    plt.plot(x_vals, t_vals, color = 'r')\n",
    "    plt.plot(vertical_coordinates, horizontal_coordinates, color = 'g')\n",
    "    plt.gca().invert_yaxis() # invert axis so it is like the timestack image coordinates\n",
    "    plt.show()\n",
    "\n",
    "    #full_stack_window = \"timestack_annotated\"\n",
    "    #cv2.namedWindow(full_stack_window, cv2.WINDOW_NORMAL)\n",
    "    #cv2.resizeWindow(full_stack_window, 1080, 1080)\n",
    "    #cv2.imshow(full_stack_window, timestack_annotated)\n",
    "    #cv2.waitKey()\n",
    "    #cv2.destroyAllWindows()\n",
    "\n",
    "    # find MSE for evey row\n",
    "    SE = []\n",
    "    for i in range(0, len(t_vals), 1):\n",
    "        # make sure the comparison values line up correctly\n",
    "        index_matching = np.where(np.array(horizontal_coordinates) == t_vals[i])\n",
    "        vertical_coordinate_matching = vertical_coordinates[ index_matching[0][0] ]\n",
    "        # /image range normalises the pixel values between 0 and 1\n",
    "        SE.append( ((x_vals[i] - vertical_coordinate_matching)/image_range)**2 )\n",
    "    MSE = np.mean(SE)\n",
    "\n",
    "    # save the image\n",
    "    print(\"saving image\")\n",
    "    cv2.imwrite(results_dir + \"/\" + \"annotated_runup_timestack_\" + timestack_name + \"_\" + tf_model_version + \".png\", timestack_annotated)\n",
    "    \n",
    "    # save the image\n",
    "    print(\"saving mask\")\n",
    "    cv2.imwrite(results_dir + \"/\" + \"mask_output_timestack_\" + timestack_name + \"_\" + tf_model_version + \".png\", mask_annotated)\n",
    "    \n",
    "    # save the data\n",
    "    print(\"Saving Data to csv file\")\n",
    "\n",
    "    # create csv file\n",
    "    output_data_name = results_dir + \"/\" + \"runup_data_\" + timestack_name + \"_\" + tf_model_version + \".csv\"\n",
    "    with open(output_data_name, 'w') as writeFile:\n",
    "        writer = csv.writer(writeFile)\n",
    "        writer.writerows([[\"timestack shoreline coordinates\"]])\n",
    "\n",
    "    with open(output_data_name, 'a', newline='') as csvFile:\n",
    "        writer = csv.writer(csvFile)\n",
    "\n",
    "        time_row = list([\"time (horizontal coordinate)\"])\n",
    "        time_row.extend(horizontal_coordinates)\n",
    "        x_row = list([\"x (vertical coordinate)\"])\n",
    "        x_row.extend(vertical_coordinates)\n",
    "        writer.writerow(time_row)\n",
    "        writer.writerow(x_row)\n",
    "\n",
    "    csvFile.close()\n",
    "\n",
    "    return MSE\n",
    "\n",
    "print(\"done\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now upload images and masks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"%%%%%%%%%%%%%%%%%%%% GET IMAGES %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%\"\"\"\n",
    "# Instead get my own images\n",
    "\n",
    "input_images_folder = \"colour_images\"\n",
    "mask_images_folder = \"masks\"\n",
    "\n",
    "IMG_SIZE = 256 # smallest image size is 593x592\n",
    "IMG_SIZE_WID = IMG_SIZE\n",
    "brightness_change = 40 # only necessary if you include data augmentation with brightness variation\n",
    "\n",
    "data = []\n",
    "# first get all images and label accordingly\n",
    "image_names = os.listdir(input_images_folder)\n",
    "mask_names = os.listdir(mask_images_folder)\n",
    "for i in range(0, len(image_names), 1):\n",
    "    \n",
    "    # data should be list of images and masks\n",
    "    inp_img_full = cv2.imread(input_images_folder + \"/\" + image_names[i])\n",
    "    \n",
    "    # get equivalent mask name\n",
    "    # this line matches the timestack name (0th to 22nd character) and snapshot number name (between m and .)\n",
    "    mask_indx = mask_names.index(image_names[i][:23] + \"_mask_\" + image_names[i].split('m')[1].split('.')[0] + \".png\")\n",
    "    mask_label_gray_full = cv2.imread(mask_images_folder + \"/\" + mask_names[mask_indx], cv2.IMREAD_GRAYSCALE)\n",
    "    \n",
    "    # resize masks and colour images\n",
    "    inp_img = cv2.resize(inp_img_full, (IMG_SIZE_WID, IMG_SIZE))\n",
    "    mask_label_gray = cv2.resize(mask_label_gray_full, (IMG_SIZE_WID, IMG_SIZE))\n",
    "    \n",
    "    # convert mask into black and white\n",
    "    mask_label = cv2.threshold(mask_label_gray, 127, 255, cv2.THRESH_BINARY)[1]\n",
    "    \n",
    "    # do some manipulations with brightness\n",
    "    # https://stackoverflow.com/questions/37822375/python-opencv-increasing-image-brightness-without-overflowing-uint8-array\n",
    "    # increase the brightness of image\n",
    "    inp_img_l = inp_img.copy()\n",
    "    inp_img_l = np.where((255 - inp_img_l) < brightness_change, 255, inp_img_l + brightness_change)\n",
    "\n",
    "    # decrease the brightness of image\n",
    "    inp_img_d = inp_img.copy()\n",
    "    inp_img_d = np.where((inp_img_d) < brightness_change, 0, inp_img_d - brightness_change)\n",
    "\n",
    "    data.append([inp_img_d, mask_label])  \n",
    "    data.append([inp_img_l, mask_label])\n",
    "    data.append([inp_img, mask_label])\n",
    "\n",
    "    \n",
    "random.seed(0) # good for comparing hyperparameter performance\n",
    "\n",
    "# randomize rows of data\n",
    "random.shuffle(data)\n",
    "\n",
    "# setup a training set and test set\n",
    "training_test_ratio = 0.90 # 90% training is because the more meaningful test set is another set of images. Otherwise this would be a larger proportion. You could even set this to 1 if you wanted.\n",
    "data_train = data[0:int(training_test_ratio*len(data))]\n",
    "data_test = data[int(training_test_ratio*len(data)):]\n",
    "\n",
    "# split x data from y data\n",
    "x_train = np.array(list(map(list, zip(*data_train)))[0])\n",
    "y_train = np.array(list(map(list, zip(*data_train)))[1])\n",
    "\n",
    "# needs to all be numpy arrays to work in tensorflow\n",
    "x_test = np.array(list(map(list, zip(*data_test)))[0])\n",
    "x_test_show = x_test # keep a copy for visualization later\n",
    "y_test = np.array(list(map(list, zip(*data_test)))[1])\n",
    "\n",
    "# normalize input features (which are the pixel values)\n",
    "#x_train = tf.keras.utils.normalize(x_train, axis=1)\n",
    "#x_test = tf.keras.utils.normalize(x_test, axis=1)\n",
    "#y_train = tf.keras.utils.normalize(y_train, axis=1)\n",
    "#y_test = tf.keras.utils.normalize(y_test, axis=1)\n",
    "x_train = x_train / 255\n",
    "x_test = x_test / 255\n",
    "y_train = y_train / 255\n",
    "y_test = y_test / 255\n",
    "\n",
    "# reshape for convolutional compatability\n",
    "x_train = x_train.reshape(-1, IMG_SIZE, IMG_SIZE_WID, 3) # 3 is for the 3 colour channels\n",
    "x_test = x_test.reshape(-1, IMG_SIZE, IMG_SIZE_WID, 3)\n",
    "y_train = y_train.reshape(-1, IMG_SIZE, IMG_SIZE_WID, 1) # this is for the black and white\n",
    "y_test = y_test.reshape(-1, IMG_SIZE, IMG_SIZE_WID, 1)\n",
    "\n",
    "print(\"done\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Unet Architecture - This is from Wei Dai's example U-net python code from the Summer of AI 2020. Slightly modified for this use case."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "\"\"\"%%%%%%%%%%%%%%%%%%%% MACHINE LEARNING %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%\"\"\"\n",
    "# Every time you train a model, change the name of tf_model_version below. \n",
    "# It is also good practice to record all hyperparameters you used and the performance \n",
    "# of each model trained. Then you can look through all your results and learn something.\n",
    "tf_model_version = \"test_1\"\n",
    "\n",
    "def unet_model(output_channels = 1, f=8):\n",
    "    inputs = tf.keras.layers.Input(shape=(IMG_SIZE, IMG_SIZE_WID, 3))\n",
    "    \n",
    "    # Downsampling through the model\n",
    "    d1 = tf.keras.layers.Conv2D(f, 3, padding='same', activation='relu')(inputs)\n",
    "    d1 = tf.keras.layers.Conv2D(f, 3, padding='same', activation='relu')(d1)\n",
    "\n",
    "    d2 = tf.keras.layers.MaxPooling2D()(d1)\n",
    "    d2 = tf.keras.layers.Conv2D(2*f, 3, padding='same', activation='relu')(d2)\n",
    "    d2 = tf.keras.layers.Conv2D(2*f, 3, padding='same', activation='relu')(d2)\n",
    "    \n",
    "    d3 = tf.keras.layers.MaxPooling2D()(d2)\n",
    "    d3 = tf.keras.layers.Conv2D(4*f, 3, padding='same', activation='relu')(d3)\n",
    "    d3 = tf.keras.layers.Conv2D(4*f, 3, padding='same', activation='relu')(d3)\n",
    "\n",
    "    d4 = tf.keras.layers.MaxPooling2D()(d3)\n",
    "    d4 = tf.keras.layers.Conv2D(8*f, 3, padding='same', activation='relu')(d4)\n",
    "    d4 = tf.keras.layers.Conv2D(8*f, 3, padding='same', activation='relu')(d4)\n",
    "\n",
    "    d5 = tf.keras.layers.MaxPooling2D()(d4)\n",
    "    d5 = tf.keras.layers.Conv2D(16*f, 3, padding='same', activation='relu')(d5)\n",
    "    d5 = tf.keras.layers.Conv2D(16*f, 3, padding='same', activation='relu')(d5)\n",
    "    \n",
    "    # extra layer deep\n",
    "    d6 = tf.keras.layers.MaxPooling2D()(d5)\n",
    "    d6 = tf.keras.layers.Conv2D(32*f, 3, padding='same', activation='relu')(d6)\n",
    "    d6 = tf.keras.layers.Conv2D(32*f, 3, padding='same', activation='relu')(d6)\n",
    "\n",
    "    # Upsampling and establishing the skip connections\n",
    "    u5 = tf.keras.layers.UpSampling2D()(d6)\n",
    "    u5 = tf.keras.layers.concatenate([u5, d5])\n",
    "    u5 = tf.keras.layers.Conv2D(16*f, 3, padding='same', activation='relu')(u5)\n",
    "    u5 = tf.keras.layers.Conv2D(16*f, 3, padding='same', activation='relu')(u5)\n",
    "    \n",
    "    u4 = tf.keras.layers.UpSampling2D()(d5)\n",
    "    u4 = tf.keras.layers.concatenate([u4, d4])\n",
    "    u4 = tf.keras.layers.Conv2D(8*f, 3, padding='same', activation='relu')(u4)\n",
    "    u4 = tf.keras.layers.Conv2D(8*f, 3, padding='same', activation='relu')(u4)\n",
    "\n",
    "    u3 = tf.keras.layers.UpSampling2D()(u4)\n",
    "    u3 = tf.keras.layers.concatenate([u3, d3])\n",
    "    u3 = tf.keras.layers.Conv2D(4*f, 3, padding='same', activation='relu')(u3)\n",
    "    u3 = tf.keras.layers.Conv2D(4*f, 3, padding='same', activation='relu')(u3)\n",
    "\n",
    "    u2 = tf.keras.layers.UpSampling2D()(u3)\n",
    "    u2 = tf.keras.layers.concatenate([u2, d2])\n",
    "    u2 = tf.keras.layers.Conv2D(2*f, 3, padding='same', activation='relu')(u2)\n",
    "    u2 = tf.keras.layers.Conv2D(2*f, 3, padding='same', activation='relu')(u2)\n",
    "\n",
    "    u1 = tf.keras.layers.UpSampling2D()(u2)\n",
    "    u1 = tf.keras.layers.concatenate([u1, d1])\n",
    "    u1 = tf.keras.layers.Conv2D(f, 3, padding='same', activation='relu')(u1)\n",
    "    u1 = tf.keras.layers.Conv2D(f, 3, padding='same', activation='relu')(u1)\n",
    "\n",
    "    # This is the last layer of the model. \n",
    "    outputs = tf.keras.layers.Conv2D(output_channels, 1, activation='sigmoid')(u1)\n",
    "\n",
    "    return tf.keras.Model(inputs=inputs, outputs=outputs)\n",
    "\n",
    "model = unet_model(1,8)\n",
    "model.summary()\n",
    "\n",
    "model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])\n",
    "\n",
    "model.fit(x_train, y_train, batch_size=4, validation_split=0.1, epochs=30)\n",
    "\n",
    "\"\"\"%%%%%%%%%%%%%%%%%%%% EVALUATION %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%\"\"\"\n",
    "\n",
    "# validation loss and validation accuracy should be close to training loss and accuracy (otherwise overfitting!)\n",
    "train_loss, train_acc = model.evaluate(x_train, y_train)\n",
    "val_loss, val_acc = model.evaluate(x_test, y_test)\n",
    "\n",
    "print(\"training loss = \" + str(train_loss) + \", training accuracy = \" + str(train_acc))\n",
    "print(\"test loss     = \" + str(val_loss) + \", test accuracy     = \" + str(val_acc))\n",
    "\n",
    "#model.save(tf_model_version + \".model\") # this worked for an older version of tensorflow 2.x\n",
    "model.save(tf_model_version + \"_model\")\n",
    "\n",
    "predictions = model.predict([x_test])\n",
    "\n",
    "#print(predictions[0])\n",
    "\n",
    "# plot graphs inside the notebook\n",
    "%matplotlib inline \n",
    "\n",
    "# make a new directory for results\n",
    "current_path = os.getcwd()\n",
    "results_dir = tf_model_version\n",
    "results_dir_path = os.path.join(current_path, results_dir)\n",
    "if os.path.exists(results_dir_path):\n",
    "    shutil.rmtree(results_dir_path)\n",
    "    print(\"old results directory removed\")\n",
    "os.mkdir(results_dir_path)\n",
    "print(\"new results directory created called \" + results_dir)\n",
    "\n",
    "# don't change the min_uprush and max_uprush values, \n",
    "# these are fixed values according to the labelled evaluation timestacks\n",
    "error_list = []\n",
    "timestack_name = \"20101109084810_02_25ppm_test_s.png\" # 194 and 925\n",
    "min_uprush = 194\n",
    "max_uprush = 925\n",
    "ground_truth_error = compare_shoreline(min_uprush, max_uprush, tf_model_version, IMG_SIZE, timestack_name, results_dir)\n",
    "print(\"MSE for \" + timestack_name + \" is \" + str(ground_truth_error))\n",
    "error_list.append(ground_truth_error)\n",
    "\n",
    "timestack_name = \"20101110094203_05_25ppm_test_s.png\" # 40 and 580\n",
    "min_uprush = 40\n",
    "max_uprush = 580\n",
    "ground_truth_error = compare_shoreline(min_uprush, max_uprush, tf_model_version, IMG_SIZE, timestack_name, results_dir)\n",
    "print(\"MSE for \" + timestack_name + \" is \" + str(ground_truth_error))\n",
    "error_list.append(ground_truth_error)\n",
    "\n",
    "timestack_name = \"20101111092821_01_25ppm_test_s.png\" # 116 and 800\n",
    "min_uprush = 116\n",
    "max_uprush = 800\n",
    "ground_truth_error = compare_shoreline(min_uprush, max_uprush, tf_model_version, IMG_SIZE, timestack_name, results_dir)\n",
    "print(\"MSE for \" + timestack_name + \" is \" + str(ground_truth_error))\n",
    "error_list.append(ground_truth_error)\n",
    "\n",
    "timestack_name = \"OneMile8_test_s.png\" # 77 and 289\n",
    "min_uprush = 77\n",
    "max_uprush = 289\n",
    "ground_truth_error = compare_shoreline(min_uprush, max_uprush, tf_model_version, IMG_SIZE, timestack_name, results_dir)\n",
    "print(\"MSE for \" + timestack_name + \" is \" + str(ground_truth_error))\n",
    "error_list.append(ground_truth_error)\n",
    "\n",
    "timestack_name = \"FC4_test_s.png\" # 37 and 240\n",
    "min_uprush = 37\n",
    "max_uprush = 240\n",
    "ground_truth_error = compare_shoreline(min_uprush, max_uprush, tf_model_version, IMG_SIZE, timestack_name, results_dir)\n",
    "print(\"MSE for \" + timestack_name + \" is \" + str(ground_truth_error))\n",
    "error_list.append(ground_truth_error)\n",
    "\n",
    "print(\"\\nMSE summary for all timestacks:\")\n",
    "print(\"20101109084810_02_25ppm_test_s.png\" + \" \" + str(error_list[0]))\n",
    "print(\"20101110094203_05_25ppm_test_s.png\" + \" \" + str(error_list[1]))\n",
    "print(\"20101111092821_01_25ppm_test_s.png\" + \" \" + str(error_list[2]))\n",
    "print(\"OneMile8_test_s.png               \" + \" \" + str(error_list[3]))\n",
    "print(\"FC4_test_s.png                    \" + \" \" + str(error_list[4]))\n",
    "\n",
    "print(\"\\nAverage MSE for all seen timestacks = \" + str(np.mean(error_list)))\n",
    "\n",
    "print(\"done\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can now visualize your results in the results directory. You can run the below cell to see what my best results were using a regression CNN."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function used to evaluate how good the model is\n",
    "def compare_shoreline_mt(min_uprush, max_uprush, tf_model_version, IMG_SIZE, timestack_name, results_dir):\n",
    "    \n",
    "    folder_name = \"test_timestacks/\"\n",
    "\n",
    "    image_range = max_uprush - min_uprush\n",
    "    \n",
    "    timestack = cv2.imread(folder_name + timestack_name)\n",
    "    timestack_focus = timestack[:, min_uprush:max_uprush]\n",
    "    timestack_labelled = cv2.imread(folder_name + timestack_name + \"_output.png\")\n",
    "    timestack_annotated = timestack_labelled.copy()\n",
    "\n",
    "    height_of_two_peaks = 104 # normally 104 for 2 peaks\n",
    "    # for every row, find the shoreline location (plot as blue dots). If it is a minimum point plot as red Dot, If maximum plot as blue Dot\n",
    "    padding = height_of_two_peaks\n",
    "    height, width, _ = timestack.shape\n",
    "    width = max_uprush - min_uprush\n",
    "    timestack_count = 0\n",
    "\n",
    "    IMG_SIZE_WID = int(IMG_SIZE) # can change this if needed\n",
    "    tf_images = []\n",
    "    if tf_model_version == \"michaelThompson_imgSize_64\":\n",
    "        # this way of reading/saving models worked for an older version of tensorflow 2.x\n",
    "        tf_model_xc = tf.keras.models.load_model(tf_model_version + \".model\")\n",
    "    else:\n",
    "        tf_model_xc = tf.keras.models.load_model(tf_model_version + \"_model\")\n",
    "\n",
    "    vertical_coordinates = []\n",
    "    horizontal_coordinates = []\n",
    "    for i in range(0, int((height - padding)), 1):\n",
    "\n",
    "        horizontal_coordinate = i + int(padding/2)\n",
    "\n",
    "        # get the small section of the image for the horizontal position i\n",
    "        timestack_snapshot = timestack_focus[i:i + padding, :]\n",
    "        title_window_2 = \"Snapshot of timestack used for processing\"\n",
    "\n",
    "        # classify the timestack snapshot through the convolutional neural network for X_coordinate\n",
    "        #tf_images = cv2.cvtColor(timestack_snapshot, cv2.COLOR_BGR2GRAY)   # should have this for gray\n",
    "        tf_images = cv2.resize(timestack_snapshot, (IMG_SIZE_WID, IMG_SIZE))\n",
    "        tf_images_i = np.array(tf_images)\n",
    "        tf_images_i = tf.keras.utils.normalize(tf_images_i, axis=1)\n",
    "        tf_images_i = tf_images_i.reshape(-1, IMG_SIZE, IMG_SIZE_WID, 3)   #should be changed to 1 for gray\n",
    "        prediction = tf_model_xc.predict([tf_images_i])\n",
    "\n",
    "        vertical_coordinate = int((prediction[0])*width) + min_uprush\n",
    "\n",
    "        cv2.circle(timestack_annotated, (vertical_coordinate, horizontal_coordinate), radius=1, color=[255, 0, 0])\n",
    "        vertical_coordinates.append(vertical_coordinate)\n",
    "        horizontal_coordinates.append(horizontal_coordinate)\n",
    "\n",
    "        timestack_count = timestack_count + 1\n",
    "\n",
    "\n",
    "    # get csv values from labelled timestack\n",
    "    manual_csv_points_filename = folder_name + \"runup_data_test\" + timestack_name + \".csv\"\n",
    "    x_vals = []\n",
    "    t_vals = []\n",
    "    with open(manual_csv_points_filename, newline='') as csvfile:\n",
    "        points_reader = csv.reader(csvfile, delimiter=' ', quotechar='|')\n",
    "        initial_row = 5\n",
    "        initial_row_counter = 0\n",
    "        for row in points_reader:\n",
    "            if initial_row_counter > initial_row:\n",
    "                rows = row[0].split(',')\n",
    "                x_vals.append( int(rows[0][1:-1]) ) # need to remove [] from value and turn into a number\n",
    "                t_vals.append( int(rows[1][1:-1]) ) # need to remove [] from value and turn into a number\n",
    "            initial_row_counter = initial_row_counter + 1\n",
    "\n",
    "    print(\"red is hand picked, green is found by model\")\n",
    "    plt.plot(x_vals, t_vals, color = 'r')\n",
    "    plt.plot(vertical_coordinates, horizontal_coordinates, color = 'g')\n",
    "    plt.gca().invert_yaxis() # invert axis so it is like the timestack image coordinates\n",
    "    plt.show()\n",
    "\n",
    "    #full_stack_window = \"timestack_annotated\"\n",
    "    #cv2.namedWindow(full_stack_window, cv2.WINDOW_NORMAL)\n",
    "    #cv2.resizeWindow(full_stack_window, 1080, 1080)\n",
    "    #cv2.imshow(full_stack_window, timestack_annotated)\n",
    "    #cv2.waitKey()\n",
    "    #cv2.destroyAllWindows()\n",
    "\n",
    "    # find MSE for evey row\n",
    "    SE = []\n",
    "    for i in range(0, len(t_vals), 1):\n",
    "        # make sure the comparison values line up correctly\n",
    "        index_matching = np.where(np.array(horizontal_coordinates) == t_vals[i])\n",
    "        vertical_coordinate_matching = vertical_coordinates[ index_matching[0][0] ]\n",
    "        # /image range normalises the pixel values between 0 and 1\n",
    "        SE.append( ((x_vals[i] - vertical_coordinate_matching)/image_range)**2 )\n",
    "    MSE = np.mean(SE)\n",
    "\n",
    "    # save the image\n",
    "    print(\"saving image\")\n",
    "    cv2.imwrite(results_dir + \"/\" + \"annotated_runup_timestack_\" + timestack_name + \"_\" + tf_model_version + \".png\", timestack_annotated)\n",
    "    \n",
    "    # save the data\n",
    "    print(\"Saving Data to csv file\")\n",
    "\n",
    "    # create csv file\n",
    "    output_data_name = results_dir + \"/\" + \"runup_data_\" + timestack_name + \"_\" + tf_model_version + \".csv\"\n",
    "    with open(output_data_name, 'w') as writeFile:\n",
    "        writer = csv.writer(writeFile)\n",
    "        writer.writerows([[\"timestack shoreline coordinates\"]])\n",
    "\n",
    "    with open(output_data_name, 'a', newline='') as csvFile:\n",
    "        writer = csv.writer(csvFile)\n",
    "\n",
    "        time_row = list([\"time (horizontal coordinate)\"])\n",
    "        time_row.extend(horizontal_coordinates)\n",
    "        x_row = list([\"x (vertical coordinate)\"])\n",
    "        x_row.extend(vertical_coordinates)\n",
    "        writer.writerow(time_row)\n",
    "        writer.writerow(x_row)\n",
    "\n",
    "    csvFile.close()\n",
    "\n",
    "    return MSE\n",
    "\n",
    "print(\"done\")\n",
    "\n",
    "mt_model = \"michaelThompson_imgSize_64\"\n",
    "IMG_SIZE = 64\n",
    "\n",
    "# make a new directory\n",
    "current_path = os.getcwd()\n",
    "results_dir = mt_model\n",
    "results_dir_path = os.path.join(current_path, results_dir)\n",
    "if os.path.exists(results_dir_path):\n",
    "    shutil.rmtree(results_dir_path)\n",
    "    print(\"old results directory removed\")\n",
    "os.mkdir(results_dir_path)\n",
    "print(\"new results directory created called \" + results_dir)\n",
    "\n",
    "# don't change the min_uprush and max_uprush values, \n",
    "# these are fixed values according to the labelled evaluation timestacks\n",
    "error_list = []\n",
    "timestack_name = \"20101109084810_02_25ppm_test_s.png\" # 194 and 925\n",
    "min_uprush = 194\n",
    "max_uprush = 925\n",
    "ground_truth_error = compare_shoreline_mt(min_uprush, max_uprush, mt_model, IMG_SIZE, timestack_name, results_dir)\n",
    "print(\"MSE for \" + timestack_name + \" is \" + str(ground_truth_error))\n",
    "error_list.append(ground_truth_error)\n",
    "\n",
    "timestack_name = \"20101110094203_05_25ppm_test_s.png\" # 40 and 580\n",
    "min_uprush = 40\n",
    "max_uprush = 580\n",
    "ground_truth_error = compare_shoreline_mt(min_uprush, max_uprush, mt_model, IMG_SIZE, timestack_name, results_dir)\n",
    "print(\"MSE for \" + timestack_name + \" is \" + str(ground_truth_error))\n",
    "error_list.append(ground_truth_error)\n",
    "\n",
    "timestack_name = \"20101111092821_01_25ppm_test_s.png\" # 116 and 800\n",
    "min_uprush = 116\n",
    "max_uprush = 800\n",
    "ground_truth_error = compare_shoreline_mt(min_uprush, max_uprush, mt_model, IMG_SIZE, timestack_name, results_dir)\n",
    "print(\"MSE for \" + timestack_name + \" is \" + str(ground_truth_error))\n",
    "error_list.append(ground_truth_error)\n",
    "\n",
    "timestack_name = \"OneMile8_test_s.png\" # 77 and 289\n",
    "min_uprush = 77\n",
    "max_uprush = 289\n",
    "ground_truth_error = compare_shoreline_mt(min_uprush, max_uprush, mt_model, IMG_SIZE, timestack_name, results_dir)\n",
    "print(\"MSE for \" + timestack_name + \" is \" + str(ground_truth_error))\n",
    "error_list.append(ground_truth_error)\n",
    "\n",
    "timestack_name = \"FC4_test_s.png\" # 37 and 240\n",
    "min_uprush = 37\n",
    "max_uprush = 240\n",
    "ground_truth_error = compare_shoreline_mt(min_uprush, max_uprush, mt_model, IMG_SIZE, timestack_name, results_dir)\n",
    "print(\"MSE for \" + timestack_name + \" is \" + str(ground_truth_error))\n",
    "error_list.append(ground_truth_error)\n",
    "\n",
    "print(\"\\nMSE summary for all timestacks:\")\n",
    "print(\"20101109084810_02_25ppm_test_s.png\" + \" \" + str(error_list[0]))\n",
    "print(\"20101110094203_05_25ppm_test_s.png\" + \" \" + str(error_list[1]))\n",
    "print(\"20101111092821_01_25ppm_test_s.png\" + \" \" + str(error_list[2]))\n",
    "print(\"OneMile8_test_s.png               \" + \" \" + str(error_list[3]))\n",
    "print(\"FC4_test_s.png                    \" + \" \" + str(error_list[4]))\n",
    "\n",
    "print(\"\\nAverage MSE for all seen timestacks = \" + str(np.mean(error_list)))\n",
    "\n",
    "print(\"done\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you want to use a GPU, this code is helpful to check that it will work. It can be a bit of a pain to set it up in anaconda, but the efficiency is worth it!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.test.is_gpu_available(\n",
    "    cuda_only=False,\n",
    "    min_cuda_compute_capability=None\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}