Nasa_Harvest_Model_MLHub.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "27f7fd9c",
   "metadata": {},
   "source": [
    "<img src='https://radiant-assets.s3-us-west-2.amazonaws.com/PrimaryRadiantMLHubLogo.png' alt='Radiant MLHub Logo' width='300'/>\n",
    "\n",
    "# 2021 NASA Harvest Rwanda Baseline Model\n",
    "\n",
    "This notebook walks you through the steps to create a baseline field delineation model for detecting boundaries from sentinel-2 time-series satellite imagery using a spatio-temporal U-Net model on the 2021 NASA Harvest dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6337cc89",
   "metadata": {},
   "source": [
    "### Dependencies\n",
    "\n",
    "All the dependencies for this notebook are included in the `requirements.txt` file included in this folder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3752bc39",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Importing the needed libraries\n",
    "import getpass\n",
    "import glob\n",
    "import keras\n",
    "import os\n",
    "import pickle\n",
    "import random\n",
    "from radiant_mlhub import Dataset\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import rasterio as rio\n",
    "\n",
    "import tensorflow as tf\n",
    "import segmentation_models as sm\n",
    "from segmentation_models import Unet\n",
    "\n",
    "from pathlib import Path\n",
    "from random import choice\n",
    "from scipy.ndimage import gaussian_filter\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "\n",
    "from keras import backend as K\n",
    "from keras.layers import *\n",
    "from keras.models import *\n",
    "from keras.models import load_model\n",
    "from keras.optimizers import *\n",
    "from keras.preprocessing import image\n",
    "\n",
    "from tensorflow.keras.layers import *\n",
    "from tensorflow.keras.losses import *\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "from typing import List, Any, Callable, Tuple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87d838a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.config.list_physical_devices(\"GPU\") #activate GPU resource for model training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bc382f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_id = 'nasa_rwanda_field_boundary_competition'\n",
    "assets = ['labels']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6aa90238",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Append your MLHUB_API_KEY after this cell is executed to download dataset\n",
    "os.environ['MLHUB_API_KEY'] =  getpass.getpass(prompt=\"MLHub API Key: \")\n",
    "dataset = Dataset.fetch(dataset_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa7a8a79",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset.download(if_exists='overwrite')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6e40e30",
   "metadata": {},
   "outputs": [],
   "source": [
    "#image snapshot dimensions\n",
    "IMG_WIDTH = 256 \n",
    "IMG_HEIGHT = 256 \n",
    "IMG_CHANNELS = 4 #we have the rgba bands"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "764ddebe",
   "metadata": {},
   "source": [
    "We have two sets of data: the train and test dataset, each having a list of file ids belonging to them.\n",
    "For model development purposes, we will use the training set(`train_tiles`) and use the test set for model prediction/evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ee46009",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_source_items = f\"{dataset_id}/{dataset_id}_source_train\"\n",
    "train_label_items = f\"{dataset_id}/{dataset_id}_labels_train\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3962103",
   "metadata": {},
   "outputs": [],
   "source": [
    "next(os.walk(train_source_items))[1][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62b6f279",
   "metadata": {},
   "outputs": [],
   "source": [
    "def clean_string(s: str) -> str:\n",
    "    \"\"\"\n",
    "    extract the tile id and timestamp from a source image folder\n",
    "    e.g extract 'ID_YYYY_MM' from 'nasa_rwanda_field_boundary_competition_source_train_ID_YYYY_MM'\n",
    "    \"\"\"\n",
    "    s = s.replace(f\"{dataset_id}_source_\", '').split('_')[1:]\n",
    "    return '_'.join(s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5dcde75e",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_tiles = [clean_string(s) for s in next(os.walk(train_source_items))[1]]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bcbd418",
   "metadata": {},
   "source": [
    "Our source images have pixel values > 255, hence we need to apply normalisation on our images to generate a normalised image. We apply the min-max normalisation for this which is simply: $${\\text{all pixel values - minimum pixel value} \\over \\text{maximum pixel value - minimum pixel value}}$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2e534f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def normalize(\n",
    "    array: np.ndarray\n",
    "):\n",
    "    \"\"\" normalise image to give a meaningful output \"\"\"\n",
    "    array_min, array_max = array.min(), array.max()\n",
    "    return (array - array_min) / (array_max - array_min)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40e56a84",
   "metadata": {},
   "source": [
    "### Data Augmentation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e2060bd",
   "metadata": {},
   "source": [
    "In this notebook, we will perform data augmentation on our normalised images. This will be used to populate the model with data to obtain even more accurate results.\n",
    "\n",
    "We will employ the following data augmentation techniques on the dataset:\n",
    "- rotation, flipping, blurring.\n",
    "\n",
    "These techniques were thanks to the radix-ai GitHub repository, which can be accessed [here](https://github.com/radix-ai/agoro-field-boundary-detector).\n",
    "We will observe the results on a random source image and its associated label below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1aa9927",
   "metadata": {},
   "outputs": [],
   "source": [
    "#loading the 4 bands of the image\n",
    "tile = random.choice(train_tiles)\n",
    "print(tile)\n",
    "bd1 = rio.open(f\"{train_source_items}/{dataset_id}_source_train_{tile}/B01.tif\")\n",
    "bd1_array = bd1.read(1)\n",
    "bd2 = rio.open(f\"{train_source_items}/{dataset_id}_source_train_{tile}/B02.tif\")\n",
    "bd2_array = bd2.read(1)\n",
    "bd3 = rio.open(f\"{train_source_items}/{dataset_id}_source_train_{tile}/B03.tif\")\n",
    "bd3_array = bd3.read(1)\n",
    "bd4 = rio.open(f\"{train_source_items}/{dataset_id}_source_train_{tile}/B04.tif\")\n",
    "bd4_array = bd4.read(1)\n",
    "b01_norm = normalize(bd1_array)\n",
    "b02_norm = normalize(bd2_array)\n",
    "b03_norm = normalize(bd3_array)\n",
    "b04_norm = normalize(bd4_array)\n",
    "\n",
    "field = np.dstack((b04_norm, b03_norm, b02_norm, b01_norm))\n",
    "mask  = rio.open(Path.cwd() / f\"{train_label_items}/{dataset_id}_labels_train_{tile.split('_')[0]}/raster_labels.tif\").read(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "206812fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "#https://github.com/radix-ai/agoro-field-boundary-detector/tree/master/src/agoro_field_boundary_detector\n",
    "def t_linear(\n",
    "    field: np.ndarray,\n",
    "    mask: np.ndarray,\n",
    "    _: int = 0,\n",
    ") -> Tuple[np.ndarray, np.ndarray]:\n",
    "    \"\"\"Apply a linear (i.e. no) transformation and save.\"\"\"\n",
    "    return field, mask\n",
    "\n",
    "def t_rotation(\n",
    "    field: np.ndarray,\n",
    "    mask: np.ndarray,\n",
    "    rot: int,\n",
    ") -> Tuple[np.ndarray, np.ndarray]:\n",
    "    \"\"\"Rotate the data.\"\"\"\n",
    "    assert rot in range(0, 3 + 1)\n",
    "    for _ in range(rot):\n",
    "        field = np.rot90(field)\n",
    "        mask = np.rot90(mask)\n",
    "    return field, mask\n",
    "\n",
    "def t_flip(\n",
    "    field: np.ndarray,\n",
    "    mask: np.ndarray,\n",
    "    idx: int,\n",
    ") -> Tuple[np.ndarray, np.ndarray]:\n",
    "    \"\"\"Flip the data.\"\"\"\n",
    "    assert idx in range(0, 2 + 1)\n",
    "    if idx == 0:  # Diagonal\n",
    "        field = np.rot90(np.fliplr(field))\n",
    "        mask = np.rot90(np.fliplr(mask))\n",
    "    if idx == 1:  # Horizontal\n",
    "        field = np.flip(field, axis=0)\n",
    "        mask = np.flip(mask, axis=0)\n",
    "    if idx == 2:  # Vertical\n",
    "        field = np.flip(field, axis=1)\n",
    "        mask = np.flip(mask, axis=1)\n",
    "    return field, mask\n",
    "\n",
    "def t_blur(\n",
    "    field: np.ndarray,\n",
    "    mask: np.ndarray,\n",
    "    sigma: int,\n",
    ") -> Tuple[np.ndarray, np.ndarray]:\n",
    "    \"\"\"Blur the image by applying a Gaussian filter.\"\"\"\n",
    "    assert 0 <= sigma <= 10\n",
    "    sigma_f = 1.0 + (sigma / 10)\n",
    "    field = np.copy(field)\n",
    "    for i in range(3):\n",
    "        field[:, :, i] = gaussian_filter(field[:, :, i], sigma=sigma_f)\n",
    "    return field, mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e399abb8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_image(field:np.ndarray, mask:np.ndarray): \n",
    "    \"\"\"Show the field and corresponding mask.\"\"\"\n",
    "    fig = plt.figure(figsize=(8,6))\n",
    "    ax1 = fig.add_subplot(121)  # left side\n",
    "    ax2 = fig.add_subplot(122)  # right side\n",
    "    ax1.imshow(field[:,:,0:3])  # rgb band\n",
    "    plt.gray()\n",
    "    ax2.imshow(mask)\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "658d7fa3",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_image(field, mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98f36ccc",
   "metadata": {},
   "outputs": [],
   "source": [
    "f,m = t_rotation(field, mask, rot=1) #rotation\n",
    "show_image(f,m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bdefc22",
   "metadata": {},
   "outputs": [],
   "source": [
    "f,m = t_flip(field, mask, idx=0) #flipping\n",
    "show_image(f,m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3f46a71",
   "metadata": {},
   "outputs": [],
   "source": [
    "f,m = t_blur(field, mask, sigma=5) #blur\n",
    "show_image(f,m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2dbd77cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate(\n",
    "    field: np.ndarray,\n",
    "    mask: np.ndarray,\n",
    "    write_folder: Path,\n",
    "    prefix: str = \"\",\n",
    ") -> None:\n",
    "    \"\"\"\n",
    "    Generate data augmentations of the provided field and corresponding mask which includes:\n",
    "     - Linear (no) transformation\n",
    "     - Rotation\n",
    "     - Horizontal or vertical flip\n",
    "     - Gaussian filter (blur)\n",
    "    :param field: Input array of the field to augment\n",
    "    :param mask: Input array of the corresponding mask to augment\n",
    "    :param write_folder: Folder (path) to write the results (augmentations) to\n",
    "    :param prefix: Field-specific prefix used when writing the augmentation results\n",
    "    \"\"\"\n",
    "    # Generate transformations\n",
    "    f, m = [0,1,2,3], [0,1,2,3] #dummy data. will be replaced\n",
    "    f[0],m[0] = t_linear(field, mask) #no augmentation\n",
    "    f[1],m[1] = t_rotation(field, mask, rot=1) #rotation\n",
    "    f[2],m[2] = t_flip(field, mask, idx=0) #flipping\n",
    "    f[3],m[3] = t_blur(field, mask, sigma=5) #blur\n",
    "    for i in range(len(f)):        \n",
    "        with open(write_folder +'/'+ f\"fields/{str(prefix).zfill(2)}_{i}.pkl\", 'wb') as handle:\n",
    "            pickle.dump(f[i], handle, protocol=pickle.HIGHEST_PROTOCOL)\n",
    "\n",
    "        with open(write_folder +'/'+ f\"masks/{str(prefix).zfill(2)}_{i}.pkl\", 'wb') as handle:\n",
    "            pickle.dump(m[i], handle, protocol=pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19c76a1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def main(\n",
    "    field: List[np.ndarray],\n",
    "    mask: List[np.ndarray],\n",
    "    prefix: List[str],\n",
    "    write_folder: Path,\n",
    ") -> None:\n",
    "    \"\"\"\n",
    "    Generate and save data augmentations for all the fields and corresponding masks with the following:\n",
    "     - Linear (no) transformation\n",
    "     - Rotation\n",
    "     - Horizontal or vertical flip\n",
    "     - Gaussian filter (blur)\n",
    "     - Gamma filter (brightness)\n",
    "    :param fields: Fields to augment\n",
    "    :param masks: Corresponding masks to augment\n",
    "    :param prefixes: Field-specific prefixes corresponding each field\n",
    "    :param write_folder: Path to write the results (augmentations) to\n",
    "    \"\"\"\n",
    "    generate(\n",
    "        field=field,\n",
    "        mask=mask,\n",
    "        prefix=prefix,\n",
    "        write_folder=write_folder,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e379cf5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#apply augmentation effects to training set\n",
    "for tile in train_tiles:\n",
    "    bd1 = rio.open(f\"{train_source_items}/{dataset_id}_source_train_{tile}/B01.tif\")\n",
    "    bd1_array = bd1.read(1)\n",
    "    bd2 = rio.open(f\"{train_source_items}/{dataset_id}_source_train_{tile}/B02.tif\")\n",
    "    bd2_array = bd2.read(1)\n",
    "    bd3 = rio.open(f\"{train_source_items}/{dataset_id}_source_train_{tile}/B03.tif\")\n",
    "    bd3_array = bd3.read(1)\n",
    "    bd4 = rio.open(f\"{train_source_items}/{dataset_id}_source_train_{tile}/B04.tif\")\n",
    "    bd4_array = bd4.read(1)\n",
    "    b01_norm = normalize(bd1_array)\n",
    "    b02_norm = normalize(bd2_array)\n",
    "    b03_norm = normalize(bd3_array)\n",
    "    b04_norm = normalize(bd4_array)\n",
    "\n",
    "    ids_list  = tile.split('_') # XX_YYYY_MM where XX is the training file id and YYYY_MM is the timestamp\n",
    "    tile_id   = ids_list[0]\n",
    "    timestamp = f\"{ids_list[1]}_{ids_list[2]}\"\n",
    "\n",
    "    field = np.dstack((b04_norm, b03_norm, b02_norm, b01_norm))\n",
    "    mask  = rio.open(Path.cwd() / f\"{train_label_items}/{dataset_id}_labels_train_{tile_id}/raster_labels.tif\").read(1) \n",
    "\n",
    "    #create a folder for the augmented images\n",
    "    if not os.path.isdir(f\"./augmented_data/{timestamp}\"):\n",
    "        os.makedirs(f\"./augmented_data/{timestamp}\")\n",
    "    if not os.path.isdir(f\"./augmented_data/{timestamp}/fields\"):\n",
    "        os.makedirs(f\"./augmented_data/{timestamp}/fields\")\n",
    "    if not os.path.isdir(f\"./augmented_data/{timestamp}/masks\"):\n",
    "        os.makedirs(f\"./augmented_data/{timestamp}/masks\")\n",
    "\n",
    "    main( #applying augmentation effects\n",
    "        field  = field,\n",
    "        mask   = mask,\n",
    "        prefix = tile_id,\n",
    "        write_folder = f\"./augmented_data/{timestamp}\"\n",
    "    ) #approximately 30 seconds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0db7813f",
   "metadata": {},
   "outputs": [],
   "source": [
    "timestamps = next(os.walk(f\"./augmented_data\"))[1] #Get all timestamps\n",
    "augmented_files = next(os.walk(f\"./augmented_data/{timestamps[0]}/fields\"))[2] #Get all augmented tile ids. can just use one timestamp\n",
    "X = np.empty((len(augmented_files), IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS*len(timestamps)), dtype=np.float32) #time-series image\n",
    "y = np.empty((len(augmented_files), IMG_HEIGHT, IMG_WIDTH, 1), dtype=np.uint8) #mask for each scene\n",
    "i = 0\n",
    "for file in augmented_files:\n",
    "    idx = 0\n",
    "    augmented_id = file.split('.pkl')[0] #id without .pkl extension\n",
    "    temporal_fields = []\n",
    "    for timestamp in timestamps:\n",
    "        with open(f\"./augmented_data/{timestamp}/fields/{augmented_id}.pkl\", 'rb') as field:\n",
    "            field = pickle.load(field) \n",
    "        X[i][:,:,idx:idx+IMG_CHANNELS] = field\n",
    "        idx += IMG_CHANNELS\n",
    "    with open(f\"./augmented_data/{timestamp}/masks/{augmented_id}.pkl\", 'rb') as mask:\n",
    "        mask = pickle.load(mask)\n",
    "    y[i] = mask.reshape(IMG_HEIGHT, IMG_WIDTH, 1)\n",
    "    i+=1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c88dd7f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "random.randint(0, len(augmented_files)) #sanity check\n",
    "random_image = random.randint(0, len(augmented_files)-1)\n",
    "show_image(X[random_image][:,:,0:3], y[random_image])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bdf9520",
   "metadata": {},
   "source": [
    "### Model Training\n",
    "\n",
    "We decided to use U-Net as it has shown impressive results over multiple domains in image segmentation.\n",
    "We will employ a ResNet34 backbone with our spatio-temporal U-Net model.\n",
    "\n",
    "This uses our 24 channels (6 timestamps * 4 bands per timestamp) and generates the predicted field boundary per scene.\n",
    "\n",
    "We will also use an 80:20 train:validation set split for model training.\n",
    "\n",
    "Since this is a binary segmentation problem (field boundary or no field boundary), we will use the `binary cross_entropy` loss."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6fd5235",
   "metadata": {},
   "outputs": [],
   "source": [
    "BACKBONE = 'resnet34'\n",
    "preprocess_input = sm.get_preprocessing(BACKBONE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2ffb90d",
   "metadata": {},
   "outputs": [],
   "source": [
    "X = preprocess_input(X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77c631af",
   "metadata": {},
   "outputs": [],
   "source": [
    "#https://github.com/sustainlab-group/ParcelDelineation/blob/master/models/unet.py\n",
    "def unet(pretrained_weights = None,input_size = (256,256,4)):\n",
    "    inputs = Input(input_size)\n",
    "    conv1 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(inputs)\n",
    "    conv1 = BatchNormalization()(conv1)\n",
    "    conv1 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv1)\n",
    "    conv1 = BatchNormalization()(conv1)\n",
    "    pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)\n",
    "    conv2 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool1)\n",
    "    conv2 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv2)\n",
    "    pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)\n",
    "    conv3 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool2)\n",
    "    conv3 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv3)\n",
    "    pool3 = MaxPooling2D(pool_size=(2, 2))(conv3)\n",
    "    conv4 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool3)\n",
    "    conv4 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv4)\n",
    "    drop4 = Dropout(0.5)(conv4)\n",
    "    pool4 = MaxPooling2D(pool_size=(2, 2))(drop4)\n",
    "\n",
    "    conv5 = Conv2D(1024, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool4)\n",
    "    conv5 = Conv2D(1024, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv5)\n",
    "    drop5 = Dropout(0.5)(conv5)\n",
    "\n",
    "    up6 = Conv2D(512, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(drop5))\n",
    "    merge6 = concatenate([drop4,up6], axis = 3)\n",
    "    conv6 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge6)\n",
    "    conv6 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv6)\n",
    "\n",
    "    up7 = Conv2D(256, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(conv6))\n",
    "    merge7 = concatenate([conv3,up7], axis = 3)\n",
    "    conv7 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge7)\n",
    "    conv7 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv7)\n",
    "\n",
    "    up8 = Conv2D(128, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(conv7))\n",
    "    merge8 = concatenate([conv2,up8], axis = 3)\n",
    "    conv8 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge8)\n",
    "    conv8 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv8)\n",
    "\n",
    "    up9 = Conv2D(64, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(conv8))\n",
    "    merge9 = concatenate([conv1,up9], axis = 3)\n",
    "    conv9 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge9)\n",
    "    conv9 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv9)\n",
    "    conv9 = Conv2D(2, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv9)\n",
    "    conv10 = Conv2D(1, 1, activation = 'sigmoid')(conv9)\n",
    "\n",
    "    model = Model(input = inputs, output = conv10)\n",
    "\n",
    "    if(pretrained_weights):\n",
    "    \tmodel.load_weights(pretrained_weights)\n",
    "\n",
    "    return "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72ed3b45",
   "metadata": {},
   "outputs": [],
   "source": [
    "def learning_rate_scheduler(epoch):\n",
    "    lr = 1e-4\n",
    "    '''\n",
    "    if epoch > 180:\n",
    "        lr *= 0.5e-3\n",
    "    elif epoch > 150:\n",
    "        lr *= 1e-3\n",
    "    elif epoch > 120:\n",
    "        lr *= 1e-2\n",
    "    elif epoch > 80:\n",
    "        lr *= 1e-1\n",
    "    '''\n",
    "    print(\"Set Learning Rate : {}\".format(lr))\n",
    "    return lr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "643843fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_channels = 24\n",
    "input_shape = (256,256,num_channels)\n",
    "batch_size = 4"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15275933",
   "metadata": {},
   "source": [
    "For the model, we will make use of two key metrics: **Recall** and **F1-score**.\n",
    "\n",
    "**Recall** evaluates how much of the field boundaries which were labelled were actually predicted as well while the **f1-score** combines the precision and recall by evaluating the harmonic mean.\n",
    "\n",
    "**NOTE** that the recall is the more important metric for this case as we are mostly concerned about the retrieved field boundaries out of the labelled field boundaries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38e689d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "#https://datascience.stackexchange.com/questions/45165/how-to-get-accuracy-f1-precision-and-recall-for-a-keras-model\n",
    "def recall(y_true, y_pred):\n",
    "        true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1)))\n",
    "        possible_positives = K.sum(K.round(K.clip(y_true, 0, 1)))\n",
    "        recall = true_positives / (possible_positives + K.epsilon())\n",
    "        return recall\n",
    "def f1(y_true, y_pred):\n",
    "    def precision(y_true, y_pred):\n",
    "        true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1)))\n",
    "        predicted_positives = K.sum(K.round(K.clip(y_pred, 0, 1)))\n",
    "        precision = true_positives / (predicted_positives + K.epsilon())\n",
    "        return precision\n",
    "    precision = precision(y_true, y_pred)\n",
    "    recall_   = recall(y_true, y_pred)\n",
    "    return 2*((precision*recall_)/(precision+recall_+K.epsilon()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5654af4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = None \n",
    "model_unet = Unet(BACKBONE, encoder_weights='imagenet')\n",
    "new_model = keras.models.Sequential()\n",
    "new_model.add(Conv2D(3, (1,1), padding='same', activation='relu', input_shape=input_shape))\n",
    "new_model.add(model_unet)\n",
    "model = new_model \n",
    "#sm.metrics.FScore(beta=1)\n",
    "model.compile(loss='binary_crossentropy',\n",
    "              optimizer=Adam(learning_rate=learning_rate_scheduler(0)),\n",
    "              metrics=[recall,f1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "beee012b",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f2a8ce6",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_train, x_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c00a8a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.config.experimental_run_functions_eagerly(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46e51a4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "history = model.fit(x=x_train, y=y_train,\n",
    "              validation_data=(x_val, y_val),\n",
    "              steps_per_epoch = len(x_train)//batch_size,\n",
    "              validation_steps = len(x_val)//batch_size,\n",
    "              batch_size=batch_size, epochs=200)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0feb8f34",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save(f\"./unet_model.h5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ba35983",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = load_model(f\"./unet_model.h5\", custom_objects ={\"recall\":sm.metrics.Recall(threshold=0.5), \"f1\": f1})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23753d38",
   "metadata": {},
   "source": [
    "### Loading test chips to run predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ff02d6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_source_items = f\"{dataset_id}/{dataset_id}_source_test\"\n",
    "test_tiles = [clean_string(s) for s in next(os.walk(test_source_items))[1]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a880ed4",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_tile_ids = set()\n",
    "for tile in test_tiles:\n",
    "    test_tile_ids.add(tile.split('_')[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29312c09",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_test = np.empty((len(test_tile_ids), IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS*len(timestamps)), dtype=np.float32)\n",
    "i = 0\n",
    "loaded_tiles = []\n",
    "for tile_id in test_tile_ids:\n",
    "    idx = 0\n",
    "    for timestamp in timestamps:\n",
    "        bd1 = rio.open(f\"{test_source_items}/{dataset_id}_source_test_{tile_id}_{timestamp}/B01.tif\")\n",
    "        bd1_array = bd1.read(1)\n",
    "        bd2 = rio.open(f\"{test_source_items}/{dataset_id}_source_test_{tile_id}_{timestamp}/B02.tif\")\n",
    "        bd2_array = bd2.read(1)\n",
    "        bd3 = rio.open(f\"{test_source_items}/{dataset_id}_source_test_{tile_id}_{timestamp}/B03.tif\")\n",
    "        bd3_array = bd3.read(1)\n",
    "        bd4 = rio.open(f\"{test_source_items}/{dataset_id}_source_test_{tile_id}_{timestamp}/B04.tif\")\n",
    "        bd4_array = bd4.read(1)\n",
    "        b01_norm = normalize(bd1_array)\n",
    "        b02_norm = normalize(bd2_array)\n",
    "        b03_norm = normalize(bd3_array)\n",
    "        b04_norm = normalize(bd4_array)\n",
    "        \n",
    "        field = np.dstack((b04_norm, b03_norm, b02_norm, b01_norm))\n",
    "        X_test[i][:,:,idx:idx+IMG_CHANNELS] = field\n",
    "        idx+=IMG_CHANNELS\n",
    "    loaded_tiles.append(str(tile_id).zfill(2)) #track order test tiles are loaded into X to make sure tile id matches \n",
    "    i+=1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "469cabb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions_dictionary = {}\n",
    "for i in range(len(test_tile_ids)):\n",
    "    model_pred = model.predict(np.expand_dims(X_test[i], 0))\n",
    "    model_pred = model_pred[0]\n",
    "    model_pred = (model_pred >= 0.5).astype(np.uint8)\n",
    "    model_pred = model_pred.reshape(IMG_HEIGHT, IMG_WIDTH)\n",
    "    predictions_dictionary.update([(str(loaded_tiles[i]), pd.DataFrame(model_pred))])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33186bf0",
   "metadata": {},
   "outputs": [],
   "source": [
    "dfs = []\n",
    "for key, value in predictions_dictionary.items():\n",
    "    ftd = value.unstack().reset_index().rename(columns={'level_0': 'row', 'level_1': 'column', 0: 'label'})\n",
    "    ftd['tile_row_column'] = f'Tile{key}_' + ftd['row'].astype(str) + '_' + ftd['column'].astype(str)\n",
    "    ftd = ftd[['tile_row_column', 'label']]\n",
    "    dfs.append(ftd)\n",
    "\n",
    "sub = pd.concat(dfs)\n",
    "sub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5bac12a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "sub.to_csv(f\"./harvest_sample_submission.csv\", index = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26b4d879",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "1456b81e0739e370391a1194b27bc839a9f5597c29befb4a579ec9fcd0f2acde"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}