chapter-3/Burn Scars/GFM Burn Scars.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "505e8f00-813a-4359-a06b-a47aedf90e2f",
   "metadata": {
    "tags": []
   },
   "source": [
    "![NASA logo](https://upload.wikimedia.org/wikipedia/commons/thumb/e/e5/NASA_logo.svg/110px-NASA_logo.svg.png) ![IBM Research logo](https://encrypted-tbn0.gstatic.com/images?q=tbn:ANd9GcSwHxsDwxcOHsQUD2pghQ32j90pzsZLcOujpGCyU1yE&s)\n",
    "\n",
    "# Geospatial Foundation Model: Burn scar fine-tuning\n",
    "\n",
    "This is an example of how to fine-tune a model to map burn scars from HLS data using the IBM Geospatial Foundation models as a starting point.  \n",
    "\n",
    "To run a fine-tuning experiment for flood mapping we will use the MMSegmentation library (https://github.ibm.com/GeoFM-Finetuning/mmsegmentation) to fine-tune a model starting from the geospatial foundation model trained on HLS data.\n",
    "\n",
    "The following notebook assumes that you project files are placed in folder on the shared volume in the following folder structure:\n",
    "```\n",
    "configs                   - folder to place experiment configuration files\n",
    "fine-tune-checkpoints     - folder where training outputs will be generated\n",
    "GFM-Models                - folder containing the checkpoint files from the pre-trained GFM\n",
    "inference                 - folder where we will carry out our inference tasks\n",
    "training_data             - folder containing the training dataset (including labels and test/train splits etc)\n",
    "```\n",
    "\n",
    "You then create you configuration script, before submitting to the cluster to run.  The notebook will then guide you to: \n",
    "* monitor and visualise the training, \n",
    "* run the test tasks\n",
    "* use the trained model for local inference.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b5989bd-51aa-4de6-a225-a93a25b9e7f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import pandas as pd\n",
    "import glob\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import subprocess\n",
    "from pprint import pprint\n",
    "from dotenv import load_dotenv\n",
    "import datetime\n",
    "import string\n",
    "import sys\n",
    "\n",
    "sys.path.append('../')\n",
    "import geoft\n",
    "\n",
    "# Load environment variables\n",
    "load_dotenv()\n",
    "\n",
    "# Grab cluster details\n",
    "login_url, namespace, path_to_shared_volume = geoft.get_cluster_details()\n",
    "\n",
    "# Create S3 client (for pulling data and model weights)\n",
    "aws_access_key_id = '19ef4682acca46c4a8c2405d34b4fc3f'\n",
    "aws_access_key_secret = 'd57f50701a34d33fcfadd95411174fe0390f8234f83863a0'\n",
    "\n",
    "s3 = geoft.create_s3_client(aws_access_key_id, aws_access_key_secret)\n",
    "\n",
    "# S3 bucket where data and model weights reside\n",
    "bucket_name = \"nasa-gfm-summer-school\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ec23b6a-8d55-4c7e-ab9a-3505279571fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "#------- Define the project name you wish to use\n",
    "\n",
    "project_name = \"burn\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28d8c6a4-31b2-4e4a-ad11-eec0e2d9ff94",
   "metadata": {},
   "source": [
    "## Connect to cluster for task submission\n",
    "\n",
    "The first thing you need to do in order to submit a training job to the cluster is login to the cluster.  This will only need to be done once per 24 hours.\n",
    "\n",
    "Run the cell below (`login_url`), and click on the generated url.\n",
    "\n",
    "Authenticate, then copy and paste the `oc login` command into the cell below (with `%%sh` at the top) and this will log you in to the cluster and allow you submit and monitor jobs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfa7c5c8-1243-48f7-848e-31158ad49591",
   "metadata": {},
   "outputs": [],
   "source": [
    "login_url"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e774ba1a-b7d3-40ee-9c51-0118206d030d",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%sh\n",
    "# <Paste oc command here>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41713036-c595-457c-acfd-a82eb48cbc21",
   "metadata": {},
   "source": [
    "## Project setup\n",
    "\n",
    "If we are starting a new fine-tuning project, we can create a new set of folders and download the training data+labels and the pre-train foundation model weights.  We create the folder structure described above, then pull the data and weights from an S3 bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "856d7d45-3695-43ac-8f22-65978f68dd1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#------- Create project folder structure \n",
    "geoft.create_project_folders(project_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bd8e0ab-1b1a-4bd7-859a-a7aa710f137b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#------- Download the pre-trained model weights\n",
    "model_name = 'epoch-832-loss-0.0473.pt' # best for burn scar mapping\n",
    "s3.download_file(bucket_name, 'gfm-models/' + model_name, path_to_shared_volume + project_name + '/gfm-models/' + model_name)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7646968-6db7-46a4-9c9e-d6e9fe0ed53a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#------- Download the training data\n",
    "dataset = 'burn-scars'\n",
    "number_of_train_files = 200\n",
    "\n",
    "training_data_path = path_to_shared_volume + project_name + '/training-data/'\n",
    "\n",
    "# Download training data\n",
    "subfolder = 'training/'\n",
    "geoft.download_s3_dir(dataset + '/' + subfolder, training_data_path, bucket_name, client=s3, number_of_files=number_of_train_files)\n",
    "\n",
    "# Download training data labels\n",
    "subfolder = 'validation/'\n",
    "geoft.download_s3_dir(dataset + '/' + subfolder, training_data_path, bucket_name, client=s3, number_of_files=50)\n",
    "\n",
    "# Download data for batch inference\n",
    "inference_data_path = path_to_shared_volume + project_name + '/inference/'\n",
    "subfolder = 'inference/'\n",
    "geoft.download_s3_dir(dataset + '/' + subfolder, inference_data_path, bucket_name, client=s3, key_dir_levels=0)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e05d4b6-e6c1-4299-8346-42fa5e0f3169",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Creating fine-tuning configuration\n",
    "\n",
    "![Fine-tune architecture](../images/finetune_arch.png)\n",
    "\n",
    "### Brief introduction to the hyperparameters we will adapt as part of this session\n",
    "**Loss function**:\n",
    "    Both tasks we will solve as part of this exercise are binary semantic segmentation task (e.g., pixelwise classification of flood vs. background). There will be two available loss functions for the task:\n",
    "| Loss functions | Description | Code |\n",
    "| -------------- | ----------- | ---- |\n",
    "| CrossEntropyLoss | is sensitive to class imbalance but very general and a good choice for an initial training | `type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1, class_weight=[0.15, 0.85]` |\n",
    "| DiceLoss | is invariant against class imbalance but tends to be more sensitive to other hyperparameters | `type='DiceLoss', use_sigmoid=False, loss_weight=1` |\n",
    "\n",
    "We observed a good performance of unweighted dice loss in our experiments.\n",
    "\n",
    "**Weighting classes** in the loss function:\n",
    "    As described above some loss functions (like CE loss) are sensitive to class imbalance. We can counter class imbalance by weighting the classes in the loss. For example, for flood mapping ~5-10% of the pixels represent parts of flood events while the rest is background. To meet the class imbalance, we can set the class weight of the flood class to, e.g., 90%, while the background class will be assigned a class weight of 10%. For segmentation on burn scars, there are only two classes, so only two class weights need to be specified.\n",
    "\n",
    "\n",
    "Example of Cross Entropy Loss options for burn scars:\n",
    "| Weight land class | Weight burn scar class | Code | \n",
    "| ------------------ | ----------------- | ---- |\n",
    "| 0.3 | 0.7 | `[0.7, 0.3]` |\n",
    "| 0.1 | 0.9 | `[0.9, 0.1]` |\n",
    "<!--     * cross entropy loss with weight water class = 0.7, weight land class = 0.3, weight cloud class = 0.0,\n",
    "    * cross entropy loss with weight water class = 0.9, weight land class = 0.1, weight cloud class = 0.0\n",
    "     -->\n",
    "\n",
    "**Learning rate**: Defines how much we want the model to change in response to the estimated error each time the model weights are updated. Options: `6e-4`, `6e-5`, `6e-6`\n",
    "\n",
    "**Auxiliary head**: To stabilize the finetuning process, the model not only includes an encoder and a decode head for segmentation, but also an auxiliary head. This part of the architecture helps to make the model more robust during finetuning. You can add and remove the auxiliary head using the boolean option: `aux_head=True`, `aux_head=False`\n",
    "\n",
    "**Depth of the decoder**: Generally, the decoder is quite light-weight compared to the GeoFM encoder. A default choice would be one or two layers of convolutions. Increasing this value will result in more parameters that the model can leverage to adapt to the downstream task -- at the cost of heavier computations (finetuning will take more time!). Options: `decode_head_conv = 1`, `decode_head_conv = 2`\n",
    "\n",
    "**Number of epochs**: Deep neural networks are typically require a certain number of epochs to converge. For example, in our experiments, we observed that the finetuning for flood mapping achieves a desirable level of fitness after ~40-50 epochs. *Please do not extend the number of epochs to more than 50 epochs to have a managable time for computations. :-)*\n",
    "\n",
    "### Setting up your experiment\n",
    "Now to set up your experiment populate the dictionary below with the options you wish to chose (based on the description above and discussions).  Don't edit the `gfm_ckpt`, `num_epochs`, `batch_size`, `number_training_files` or `project_name`.\n",
    "\n",
    "You generate your config which places the options you have chosen into a configuration file, which you can then view using the next cell.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39044e54-d093-4bed-992a-676a12ebd0ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "conf = {'gfm_ckpt': 'epoch-832-loss-0.0473.pt',\n",
    "        'loss_function': '''type='DiceLoss', use_sigmoid=False, loss_weight=1''',\n",
    "        'learning_rate': '6e-5',\n",
    "        'aux_head': 'True',\n",
    "        'decode_head_conv': '1',        \n",
    "        'num_epochs': 15,\n",
    "        'batch_size': '2',\n",
    "        'number_training_files': number_of_train_files,\n",
    "        'project_name': project_name}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40bd364c-1e64-4c70-9593-aaca3588ca31",
   "metadata": {},
   "outputs": [],
   "source": [
    "experiment_name, experiment_filepath = geoft.generate_config(project_name, conf, \"burn_config.py.template\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4277ba77-fea7-4442-8610-3b65d0310256",
   "metadata": {},
   "outputs": [],
   "source": [
    "geoft.view_config(experiment_filepath)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0364b664-699f-41ab-9bed-809b19011734",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Submitting fine-tuning job to run\n",
    "Now we have the configuration script ready, we can just send it to the cluster to run.  The next cell will submit the job to the cluster using TorchX.  This will spin up a now pod/container where the fine-tuning will run.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7df8d4d-3309-4a28-9565-cb1afa534a6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "mcad_id = geoft.submit_tune(project_name,\n",
    "                namespace,\n",
    "                experiment_name,\n",
    "                image='quay.io/bedwards-ibm/mmsegmentation-geo:latest',\n",
    "                num_gpus=1,\n",
    "                memory_mb=24000)\n",
    "\n",
    "print(mcad_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "195bd999-0402-4545-8c22-742584326926",
   "metadata": {},
   "source": [
    "## Monitoring training job\n",
    "Once you have submitted the job to the cluster, we can monitor it using the following commands.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47827ce7-38e2-4a6a-9a29-23c769dd3b54",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%sh\n",
    "torchx list -s kubernetes_mcad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52b36f8c-6b41-4cc5-bd66-ed9ef495c7ae",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "check_log_cmd = '''torchx log ''' + str(mcad_id) +  ''' | tail -n20'''\n",
    "os.system(check_log_cmd)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "009687b5-9f81-4c35-9bb0-3adb9963b9dc",
   "metadata": {},
   "source": [
    "## Viewing the training metrics\n",
    "\n",
    "Now that we have run (or at least are running) the experiment, we can view the training metrics.  To do this we will load the log file and extract the metrics to a dataframe (`val_df`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "125a922b-6f56-4d65-bb39-af505eb609b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df, val_df = geoft.load_tune_metrics(project_name, experiment_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66741529-f6a5-4a78-8cfb-aea06f0772e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure().set_figwidth(15)\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.plot(train_df.index, train_df.loss, '-r');\n",
    "plt.ylabel('Training Loss');\n",
    "# plt.yscale('log')\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.plot(train_df.index, train_df.loss_val, '-b');\n",
    "plt.ylabel('Validation Loss');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fdab586-cfa2-439e-b0aa-97f8308b401d",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Test output model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f10de3e8-7958-44ca-9ff4-9ceef79f4c69",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_mcad_id = geoft.submit_test(project_name,\n",
    "                        namespace,\n",
    "                        experiment_name,\n",
    "                        checkpoint='latest.pth',\n",
    "                        num_gpus=1,\n",
    "                        memory_mb=8000,\n",
    "                        bands='[0, 1, 2, 3, 4, 5]')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61d40fb9-5d35-4133-8e50-cc6960676027",
   "metadata": {},
   "outputs": [],
   "source": [
    "check_log_cmd = '''torchx log ''' + str(test_mcad_id) +  ''' | tail -n20'''\n",
    "os.system(check_log_cmd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f83b22b5-5b91-49db-9a43-748c27c70571",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_metrics = geoft.get_test_metrics(project_name, experiment_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2b7fe89-5f24-46ee-bc99-4dc02ac3857b",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Running inference using the trained model\n",
    "\n",
    "Once we have a trained model, we can use it to run inference on other images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ceb2fee-684f-44bb-aa3e-790f69aa0545",
   "metadata": {},
   "outputs": [],
   "source": [
    "infer_mcad_id = geoft.submit_inference(project_name,\n",
    "                namespace,\n",
    "                experiment_name,\n",
    "                checkpoint='latest.pth',\n",
    "                image='quay.io/bedwards-ibm/mmsegmentation-geo:latest',\n",
    "                num_gpus=1,\n",
    "                memory_mb=8000)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9e0c875-ebe5-4d4e-a3cf-a842b313f145",
   "metadata": {},
   "source": [
    "## Visualizing the predicitons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "842f63a4-6167-47d6-aeef-e10036095826",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install rasterio folium"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56a65c3d-df4a-4c8f-a5be-c5d511333cc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import folium\n",
    "import folium.plugins as plugins\n",
    "import rasterio\n",
    "from rasterio.plot import show\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "def colorize(array, cmax, cmin=0, cmap=\"rainbow\"):\n",
    "    \"\"\"Converts a 2D numpy array of values into an RGBA array given a colour map and range.\n",
    "    Args:\n",
    "        array (ndarray):\n",
    "        cmax (float): Max value for colour range\n",
    "        cmin (float): Min value for colour range\n",
    "        cmap (string): Colour map to use (from matplotlib colourmaps)\n",
    "    Returns:\n",
    "            rgba_array (ndarray): 3D RGBA array which can be plotted.\n",
    "    \"\"\"\n",
    "    normed_data = (array - cmin) / (array.max() - cmin)\n",
    "    cm = plt.cm.get_cmap(cmap)\n",
    "    return cm(normed_data)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae1cbd93-9ab5-4dc2-8c60-67f6a6cfb052",
   "metadata": {},
   "outputs": [],
   "source": [
    "inference_files = sorted(glob.glob('/opt/app-root/src/data/' + project_name + '/inference/*.tif'))\n",
    "inference_files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb4ef2e4-cdf6-495a-8b1b-e6792b1e0336",
   "metadata": {},
   "outputs": [],
   "source": [
    "for filenum in range(len(inference_files)):\n",
    "    original_file = inference_files[filenum]\n",
    "    predict_file = inference_files[filenum].replace('/inference','/inference/pred/' + experiment_name).replace('.tif','_pred.tif')\n",
    "    # Load the original image layer\n",
    "    with rasterio.open(original_file) as src:\n",
    "        redArray = src.read(1)\n",
    "        greenArray = src.read(2)\n",
    "        blueArray = src.read(3)\n",
    "        bounds = src.bounds\n",
    "        nd = src.nodata\n",
    "        midLat = (bounds[3] + bounds[1]) / 2\n",
    "        midLon = (bounds[2] + bounds[0]) / 2\n",
    "        im_rgb = np.moveaxis(np.array([redArray,greenArray,blueArray]), 0, -1)\n",
    "        im_rgb = im_rgb/np.max(im_rgb)\n",
    "    # Create the map\n",
    "    if filenum == 0:\n",
    "        m = folium.Map(location=[midLat, midLon], tiles='openstreetmap', max_zoom=22)\n",
    "    # Add the prediciton layer to the map\n",
    "    with rasterio.open(predict_file) as src:\n",
    "        dataArray = src.read(1)\n",
    "        bounds = src.bounds\n",
    "        nd = src.nodata\n",
    "    # cmax = np.max(dataArray)\n",
    "    cmax = 1000\n",
    "    dataArrayMasked = np.ma.masked_where(dataArray == nd, dataArray)\n",
    "    dataArrayMasked = np.ma.masked_where(dataArray == 0, dataArrayMasked)\n",
    "    imc = colorize(dataArrayMasked, cmax, cmin=0, cmap=\"viridis\")\n",
    "    # Add the layers to the map\n",
    "    pred = folium.raster_layers.ImageOverlay(imc, [[bounds[1], bounds[0]], [bounds[3], bounds[2]]], name=f\"Prediction-{filenum}\", opacity=0.8)\n",
    "    orig = folium.raster_layers.ImageOverlay(im_rgb, [[bounds[1], bounds[0]], [bounds[3], bounds[2]]], name=f\"Original image{filenum}\", opacity=1.0)\n",
    "    orig.add_to(m)\n",
    "    pred.add_to(m)\n",
    "folium.LayerControl().add_to(m)\n",
    "m.fit_bounds(bounds)\n",
    "m    "
   ]
  }
 ],
 "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}