flair_2_notebook_firsrt.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/aemrhb/IPI_project-seminar/blob/main/flair_2_notebook_firsrt.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f631ae2d-c169-461c-937a-3446a78abe79",
      "metadata": {
        "id": "f631ae2d-c169-461c-937a-3446a78abe79",
        "tags": []
      },
      "source": [
        "<center><img src=\"https://drive.google.com/uc?export=view&id=1ygAs8EMNlIim2ypwmvQn9yN1LbY3hWHV\" alt=\"Drawing\"  width=\"30%\"/><center>\n",
        "\n",
        "# <center><strong>Data Visualization and Baseline replication</strong></center>\n",
        "<br/>\n",
        "\n",
        "<br/><center>This notebook allows you to visualize the data used in the **FLAIR #2 challenge**.<br/>The code bellow works with the toy dataset (subset) provided in the starting-kit alongside this notebook as well as with the full FLAIR-two dataset accessible after registration to the competition.</center> <br/>\n",
        "<center>**We also strongly advise you to read the data technical description provided in the datapaper.**</center>\n",
        "<br/> <br/>\n",
        "  \n",
        "\n",
        "<hr style=\"height:1.5px;border-width:0;color:red;background-color:red\">    \n",
        "\n",
        "# <font color='red'>PART-1: Data vizualisation with the toy dataset</font>\n",
        "\n",
        "First, let's import relevant functions from the <font color='#D7881C'><em>data_display.py</em></font> file.\n",
        "<br/>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d403bff2-f610-41d6-b511-9cbfaa07a6e1",
      "metadata": {
        "id": "d403bff2-f610-41d6-b511-9cbfaa07a6e1"
      },
      "outputs": [],
      "source": [
        "from pathlib import Path\n",
        "import yaml\n",
        "from os.path import join\n",
        "import numpy as np\n",
        "import sys\n",
        "\n",
        "# Necessary to load from src\n",
        "module_path = str(Path.cwd().parents[0])\n",
        "if module_path not in sys.path:\n",
        "    sys.path.append(module_path)\n",
        "\n",
        "sys.path.append('/content/drive/My Drive/PS_23_24/FLAIR-2-AI-Challenge/notebook')\n",
        "from data_display import (display_nomenclature,\n",
        "                            display_samples,\n",
        "                            display_time_serie,\n",
        "                            display_all_with_semantic_class,\n",
        "                            display_all)\n",
        "from src.load_data import load_data"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "88a3b469-b677-4aa2-826e-c00e4476805b",
      "metadata": {
        "id": "88a3b469-b677-4aa2-826e-c00e4476805b",
        "tags": []
      },
      "source": [
        "## <font color='#90c149'>Nomenclatures</font>\n",
        "\n",
        "<br/><hr>\n",
        "\n",
        "Next, we display the semantic land-cover classes used in the FLAIR #2 datatset. You will see that <font color='#90c149'>two nomenclatures are available </font> :\n",
        "<ul>\n",
        "    <li>the <strong><font color='#90c149'>full nomenclature</font></strong> corresponds to the semantic classes used by experts in photo-interpretation to label the pixels of the ground-truth images.</li>\n",
        "    <li>the <font color='#90c149'><b>main (baseline) nomenclature</b></font> is a simplified version of the full nomenclature. It regroups (into the class 'other') classes that are either strongly under-represented or irrelevant to this challenge.</li>\n",
        "</ul>        \n",
        "See the associated datapaper for additionnal details on these nomenclatures.<br/><br/>\n",
        "\n",
        "<font color='#90c149'>Note:</font> in the data exploration part, we employ the full nomenclature. For the second part related to the challenge baseline, the main nomenclature is used. <br/><hr><br/>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4be159c6-f1d1-4212-84bc-c469dc52d399",
      "metadata": {
        "id": "4be159c6-f1d1-4212-84bc-c469dc52d399"
      },
      "outputs": [],
      "source": [
        "display_nomenclature()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e147039d-d45e-4ab4-a670-7093c603d7ed",
      "metadata": {
        "id": "e147039d-d45e-4ab4-a670-7093c603d7ed",
        "tags": []
      },
      "source": [
        "## <font color='#90c149'>Data display</font>\n",
        "\n",
        "<br/><hr>\n",
        "\n",
        "We start by creating lists containing the paths to the input images (`images`) and supervision masks (`masks`) files of the dataset.<hr><br/>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%cd '/content/drive/My Drive/PS_23_24/FLAIR-2-AI-Challenge'"
      ],
      "metadata": {
        "id": "WVkDFfin_5nS"
      },
      "id": "WVkDFfin_5nS",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a034c90c",
      "metadata": {
        "id": "a034c90c"
      },
      "outputs": [],
      "source": [
        "config_path = \"/content/drive/My Drive/PS_23_24/FLAIR-2-AI-Challenge/flair-2-config.yml\" # Change to yours\n",
        "with open(config_path, \"r\") as f:\n",
        "    config = yaml.safe_load(f)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(config)"
      ],
      "metadata": {
        "id": "VbxpToQ2AUit"
      },
      "id": "VbxpToQ2AUit",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "043963be",
      "metadata": {
        "id": "043963be"
      },
      "outputs": [],
      "source": [
        "# Creation of the train, val and test dictionnaries with the data file paths\n",
        "d_train, d_val, d_test = load_data(config)\n",
        "\n",
        "images = d_train[\"PATH_IMG\"]\n",
        "labels = d_train[\"PATH_LABELS\"]\n",
        "sentinel_images = d_train[\"PATH_SP_DATA\"]\n",
        "sentinel_masks = d_train[\"PATH_SP_MASKS\"] # Cloud masks\n",
        "sentinel_products = d_train[\"PATH_SP_DATES\"] # Needed to get the dates of the sentinel images\n",
        "centroids = d_train[\"SP_COORDS\"] # Position of the aerial image in the sentinel super area\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "56ac1c07",
      "metadata": {
        "id": "56ac1c07"
      },
      "source": [
        "### Visu"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "48cea476-3bd9-4de8-a562-2b325703d3b7",
      "metadata": {
        "id": "48cea476-3bd9-4de8-a562-2b325703d3b7",
        "tags": []
      },
      "source": [
        "<br/><hr>\n",
        "\n",
        "Let's display some random samples of IMG-MSK pairs. <font color='#90c149'>Re-run the cell bellow for a different image.</font> Here we also plot the Sentinel super area, super patch and patch. Even though the last one is not used in practice, it is shown to provide an idea of what the Sentinel data looks like. The red rectangle shows the extent of the RVB image inside the Sentinel image. <hr><br/>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "1a99d4b7-3f63-4321-bf59-b895ecada184",
      "metadata": {
        "id": "1a99d4b7-3f63-4321-bf59-b895ecada184"
      },
      "outputs": [],
      "source": [
        "display_samples(images, labels, sentinel_images, centroids)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "b9ac49c5",
      "metadata": {
        "id": "b9ac49c5"
      },
      "source": [
        "<br/><hr>\n",
        "We can also plot a few images from sentinel time series along with the acquisition date. Here we filter the dates with too much cloud coverage.\n",
        "\n",
        "<hr><br/>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8cb7ffa7",
      "metadata": {
        "id": "8cb7ffa7"
      },
      "outputs": [],
      "source": [
        "display_time_serie(sentinel_images, sentinel_masks, sentinel_products, nb_samples=3)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f8e4fe44-5f9c-4e69-881b-e24b64c205c8",
      "metadata": {
        "id": "f8e4fe44-5f9c-4e69-881b-e24b64c205c8",
        "tags": []
      },
      "source": [
        "<br/><hr>\n",
        "\n",
        "Next let's have a closer look at some specific semantic class.<br/> By setting `semantic_class` to a class number (*e.g.*, `semantic_class`=1 for building or `semantic_class`=5 for water) we can visualize the images containing pixels of this specific class. (the full nomenclature is be used.)<br/>\n",
        "<font color='#90c149'>Note:</font> for Colab users, this can take some time. <hr><br/>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9ce3fcb7-8c9b-4207-bdf5-1227ba6600f0",
      "metadata": {
        "id": "9ce3fcb7-8c9b-4207-bdf5-1227ba6600f0"
      },
      "outputs": [],
      "source": [
        "display_all_with_semantic_class(images, labels, semantic_class=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "da00bce4-c1da-4843-9705-e573a74a7dd9",
      "metadata": {
        "id": "da00bce4-c1da-4843-9705-e573a74a7dd9",
        "tags": []
      },
      "source": [
        "<br/><hr>\n",
        "\n",
        "We can directly display all images (be sure to use the toy dataset!).<br/> <hr><br/>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a3b7d91f-3409-471a-b40c-8ca98be09a3f",
      "metadata": {
        "id": "a3b7d91f-3409-471a-b40c-8ca98be09a3f"
      },
      "outputs": [],
      "source": [
        "display_all(images, labels)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "a0bc22ea-1ef0-4f21-aa8c-171807d11362",
      "metadata": {
        "id": "a0bc22ea-1ef0-4f21-aa8c-171807d11362",
        "tags": []
      },
      "source": [
        "<br><br>\n",
        "<hr style=\"height:3px;border-width:0;color:red;background-color:red\">   \n",
        "\n",
        "# <center><font color='red'>PART-2: Baseline </font></center>\n",
        "\n",
        "<br/><hr>\n",
        "\n",
        "In this second part, we use the toy dataset to train a model similar to the FLAIR #2 baseline provided with the challenge.<br/>\n",
        "<font color='#90c149'>Note:</font> the presented pipeline can also be applied to the full dataset.\n",
        "\n",
        "First, let's check if GPU ressources are available in our execution environment. If not, make sure to set `accelerator = 'cpu'` in the parameters.\n",
        "<hr><br/>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e5b1f9ae-bd82-40e6-94ae-3ce7e8376975",
      "metadata": {
        "id": "e5b1f9ae-bd82-40e6-94ae-3ce7e8376975"
      },
      "outputs": [],
      "source": [
        "gpu_info = !nvidia-smi\n",
        "gpu_info = '\\n'.join(gpu_info)\n",
        "if gpu_info.find('failed') >= 0: print('No GPU found.')\n",
        "else: print(gpu_info)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e368bf75-b5aa-43e8-80ef-2b03b74e7020",
      "metadata": {
        "id": "e368bf75-b5aa-43e8-80ef-2b03b74e7020",
        "tags": []
      },
      "source": [
        "<br/><hr>\n",
        "\n",
        "The cell bellow imports the required libraries, classes and functions, including those provided in the <font color='#D7881C'><em>src</em></font> folder provided with this starting-kit. If you are running this notebook on a local environment, make sure all necessary libraries are installed (refer to the <font color='red'>README.md</font> file).\n",
        "\n",
        "This baseline relies on <font color='#90c149'><em>pytorch-lightning</em></font>, a high-level python framework built on top of Pytorch. It allows multi-GPU training, significantly speeding-up computation of the baseline on the full FLAIR #2 dataset. It is however also possible to train on a single GPU as we demonstrate in this notebook.\n",
        "\n",
        "In this notebook, we also take advantage of the <font color='#90c149'><em>segmentation-models-pytorch</em></font> library, which provides a variery of different pre-trained segmentation models (*e.g.*, U-Net, PSPNet,...).\n",
        "<hr><br/>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install pytorch_lightning==1.5.10"
      ],
      "metadata": {
        "id": "LF2YtekPS-fg"
      },
      "id": "LF2YtekPS-fg",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install torchmetrics==0.9.2"
      ],
      "metadata": {
        "id": "OAlKjD8ZBXA2"
      },
      "id": "OAlKjD8ZBXA2",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torchmetrics\n",
        "print(torchmetrics.__version__)"
      ],
      "metadata": {
        "id": "4VHiqYUxzKbO"
      },
      "id": "4VHiqYUxzKbO",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "metadata": {
        "id": "SKqc5-fyV_LF"
      },
      "id": "SKqc5-fyV_LF",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# %cd /content/drive/My Drive//PS_23_24/FLAIR-2-AI-Challenge\n"
      ],
      "metadata": {
        "id": "YGum9QVPV1fe"
      },
      "id": "YGum9QVPV1fe",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install -U git+https://github.com/qubvel/segmentation_models.pytorch"
      ],
      "metadata": {
        "id": "4DQGWywnXUuB"
      },
      "id": "4DQGWywnXUuB",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# pip install -r requirements.txt"
      ],
      "metadata": {
        "id": "MXJVgPhVgPi2"
      },
      "id": "MXJVgPhVgPi2",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install rasterio>=1.2.8\n",
        "!pip install omegaconf>=2.2.3\n",
        "!pip install numpy>=1.22.4\n",
        "!pip install albumentations>1.2.0\n",
        "!pip install matplotlib==3.7.1\n",
        "!pip install pandas>=1.4.3\n",
        "!pip install Pillow>=9.3.0\n",
        "!pip install pytorch>=1.12.1\n",
        "!pip install PyYAML==6.0\n",
        "!pip install rich==13.4.1\n",
        "!pip install scikit_image>=0.18.3\n",
        "!pip install scikit_learn>=1.1.1\n",
        "#!pip install segmentation_models_pytorch>=0.2.1\n",
        "!pip install torch>=1.11.02.0.0\n",
        "!pip install torchvision>=0.12.0"
      ],
      "metadata": {
        "id": "RsI43WpVgbve"
      },
      "id": "RsI43WpVgbve",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "import albumentations as A"
      ],
      "metadata": {
        "id": "EaJJQTUgL3FK"
      },
      "id": "EaJJQTUgL3FK",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "0e9c1462-91a7-4d3c-b02b-1161380ffcf8",
      "metadata": {
        "id": "0e9c1462-91a7-4d3c-b02b-1161380ffcf8"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from pathlib import Path\n",
        "import sys\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "from pytorch_lightning.loggers import TensorBoardLogger\n",
        "from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint\n",
        "from pytorch_lightning.callbacks.progress.tqdm_progress import TQDMProgressBar\n",
        "from pytorch_lightning import Trainer, seed_everything\n",
        "from pytorch_lightning.utilities.distributed import rank_zero_only\n",
        "\n",
        "# Necessary to load from src\n",
        "module_path = str(Path.cwd().parents[0])\n",
        "if module_path not in sys.path:\n",
        "    sys.path.append(module_path)\n",
        "sys.path.append('/content/drive/My Drive/PS_23_24/FLAIR-2-AI-Challenge')\n",
        "import src\n",
        "from src.backbones.txt_model import TimeTexture_flair\n",
        "from src.datamodule import DataModule\n",
        "from src.task_module import SegmentationTask\n",
        "from src.utils_prints import print_config, print_metrics\n",
        "from src.utils_dataset import read_config\n",
        "from src.load_data import load_data\n",
        "from src.prediction_writer import PredictionWriter\n",
        "#from src.metrics import generate_miou"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "3ccd2527-a222-4ed1-a4e8-899595d0cc5d",
      "metadata": {
        "id": "3ccd2527-a222-4ed1-a4e8-899595d0cc5d",
        "tags": []
      },
      "source": [
        "## <font color='#90c149'>Task and parameters</font>\n",
        "\n",
        "<br/><hr>\n",
        "\n",
        "The toy dataset is composed of $26$ aerial patches of $512x512$ with corresponding semantic masks, and $26$ Sentinel-2 patches. It has $12$ test patches. The full FLAIR #2 dataset contains $61,712$ aerial patches and 41,029 Sentinel-2 acquisitions as training set, and $16,050$ aerial and $10,215$ satellite testing patches.<br/><br/>\n",
        "\n",
        "The next cell loads the configuration file, which defines <font color='#90c149'>the paths and hyper-parameters</font>.\n",
        "\n",
        "\n",
        "We recommand starting with the given default values and test if everything is working (check the datapaper for the baseline hyper-parameters).\n",
        "<hr><br/>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# %pwd"
      ],
      "metadata": {
        "id": "mO5MjaRQj56Z"
      },
      "id": "mO5MjaRQj56Z",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "bcc904b6-e26e-452b-913b-401394512b95",
      "metadata": {
        "id": "bcc904b6-e26e-452b-913b-401394512b95"
      },
      "outputs": [],
      "source": [
        "config_path = \"/content/drive/MyDrive/PS_23_24/FLAIR-2-AI-Challenge/flair-2-config.yml\" # Change to yours\n",
        "config =  read_config(config_path)\n",
        "\n",
        "print_config(config)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "c7933736-da55-4494-ae1b-d4ed0248e684",
      "metadata": {
        "id": "c7933736-da55-4494-ae1b-d4ed0248e684",
        "tags": []
      },
      "source": [
        "## <font color='#90c149'>Dataloaders</font>\n",
        "\n",
        "<br/><hr>\n",
        "\n",
        "The following cell loads the data into a pytorch-lighning DataModule. It takes the dictionaries containing the data paths and the configuration file as input.\n",
        "\n",
        "We fix the global seed (python random, torch, numpy) with `seed_eveything`.\n",
        "\n",
        "<hr><br/>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# %cd /content/drive/My Drive//PS_23_24/FLAIR-2-AI-Challenge/flair2/flair-2-baseline/flair_2_toy_dataset/flair_2_toy_aerial_test"
      ],
      "metadata": {
        "id": "esRYsWbUS3f1"
      },
      "id": "esRYsWbUS3f1",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# !pwd\n",
        "# !ls -a"
      ],
      "metadata": {
        "id": "Mnm2wvH_XGLT"
      },
      "id": "Mnm2wvH_XGLT",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "out_dir = Path(config[\"out_folder\"], config[\"out_model_name\"])\n",
        "out_dir.mkdir(parents=True, exist_ok=True)\n",
        "seed_everything(2022, workers=True)\n",
        "\n",
        "dict_train, dict_val, dict_test = load_data(config)\n",
        "# dict_train, dict_val = load_data(config)\n",
        "\n",
        "# Augmentation\n",
        "if config[\"use_augmentation\"] == True:\n",
        "    transform_set = A.Compose([A.VerticalFlip(p=0.5),\n",
        "                               A.HorizontalFlip(p=0.5),\n",
        "                               A.RandomRotate90(p=0.5)])\n",
        "else:\n",
        "    transform_set = None\n",
        "\n",
        "dm = DataModule(\n",
        "    dict_train = dict_train,\n",
        "    dict_val = dict_val,\n",
        "    dict_test = dict_test,\n",
        "    config=config,\n",
        "    drop_last = True)"
      ],
      "metadata": {
        "id": "6zQen6UyZwLC"
      },
      "id": "6zQen6UyZwLC",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print((dict_train))\n",
        "num_items = len(dict_train)\n",
        "print(num_items)\n",
        "keys = dict_test.keys()\n",
        "print(\"Keys:\", keys)\n",
        "PATH_IMG = dict_train[\"PATH_IMG\"]\n",
        "length_of_PATH_IMG = len(PATH_IMG)\n",
        "print(\"Length of PATH_IMG:\", length_of_PATH_IMG)\n",
        "print(dict_train['PATH_IMG'])\n",
        "PATH_SP_DATA = dict_train[\"PATH_SP_DATA\"]\n",
        "PATH_SP_DATA = len(PATH_SP_DATA)\n",
        "print(\"Length of PATH_IMG:\", PATH_SP_DATA)\n",
        "print(dict_train['PATH_SP_DATA'])\n",
        "# Print the number of items\n",
        "# print(f\"The dictionary contains {num_items} key-value pairs.\")\n",
        "# print(dict_val)\n",
        "# print(dict_test)"
      ],
      "metadata": {
        "id": "3HyooyYXfn1T"
      },
      "id": "3HyooyYXfn1T",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "id": "0c4c744e-9ff1-4acb-8f52-4c063b229915",
      "metadata": {
        "deletable": false,
        "editable": false,
        "id": "0c4c744e-9ff1-4acb-8f52-4c063b229915",
        "tags": []
      },
      "source": [
        "## <font color='#90c149'>Learning setup</font>\n",
        "\n",
        "<br/><hr>\n",
        "\n",
        "Next, we define our <font color='#90c149'>model, criterion, optimizer and callbacks</font>.\n",
        "\n",
        "The model `U-T&T` has `two` branches to extract spatial and temporal information from the very high resolution aerial images and high resolution satellite images. The two architecture which constitute the two branches embedded in the U-T&T model are:\n",
        "- `U-Net` (spatial/texture branch): for the aerial imagery patches, a U-Net architecture is adopted. The encoder is a ResNet34 backbone model  which weights are\n",
        "pre-trained on ImageNet for a total of ≈ 24.4 M parameters. This is similar to the architecture used for the FLAIR #1 baselines.\n",
        "- `U-TAE` (spatio-temporal branch): the spatial and temporal information supplied by the Sentinel-2 time series is explored with a U-TAE architecture. This U-Net based architecture includes a Temporal self-Attention Encoder (TAE) taking as input the lowest resolution features of the convolutional encoder and yielding a set of temporal attention masks further applied to all resolutions upon decoding.\n",
        "\n",
        "The architecture also encompass a fusion module, which takes as input the U-TAE embedding (last feature maps of the U-TAE decoder) and is applied to each stage of the U-\n",
        "Net branch. See the datapaper for more details on the fusion method.\n",
        "\n",
        "\n",
        "If `use_metadata = True`, it adds a custom Multi-layer Perceptron to the U-Net, encoding the metadata.\n",
        "\n",
        "As criterion, we use two `Cross Entropy` losses, one for each branch. They are summed to get the final loss. Each criterion can be initialized with differents weights for the classes to give more or less importance to particular classes.\n",
        "\n",
        "The pytorch-lighning module `SegmentationTask` organizes and manages the different loops and steps (e.g., training, validation), otherwise manually implemented using torch.\n",
        "\n",
        "Finally we define `callbacks` (save model checkpoints, stop if learning is stuck with a patience threshold and display progress) as well as a `logger` (tensorboard logs).\n",
        "\n",
        "<hr><br/>"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f2069ada-8d0a-49e6-82af-6dfb477cdfe7",
      "metadata": {
        "deletable": false,
        "editable": false,
        "id": "f2069ada-8d0a-49e6-82af-6dfb477cdfe7",
        "tags": []
      },
      "source": [
        "### Model\n",
        "<font color='#90c149'>Note:</font> the next cell will trigger the download of ResNet34 (default for U-Net architecture in pytorch-lightning) with pre-trained weights."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "0cdc448b-40fa-4e3f-bfe4-7c506ba60820",
      "metadata": {
        "id": "0cdc448b-40fa-4e3f-bfe4-7c506ba60820"
      },
      "outputs": [],
      "source": [
        "model = TimeTexture_flair(config)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "a7f8c2f2-8241-448e-a534-accbae1b6bf2",
      "metadata": {
        "deletable": false,
        "editable": false,
        "id": "a7f8c2f2-8241-448e-a534-accbae1b6bf2",
        "tags": []
      },
      "source": [
        "### Loss"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4fedcbc1-9da5-4167-9407-9c4dc427cdbb",
      "metadata": {
        "id": "4fedcbc1-9da5-4167-9407-9c4dc427cdbb"
      },
      "outputs": [],
      "source": [
        "with torch.no_grad():\n",
        "    weights_aer = torch.FloatTensor(np.array(list(config['weights_aerial_satellite'].values()))[:,0])\n",
        "    weights_sat = torch.FloatTensor(np.array(list(config['weights_aerial_satellite'].values()))[:,1])\n",
        "criterion_vhr = nn.CrossEntropyLoss(weight=weights_aer)\n",
        "criterion_hr = nn.CrossEntropyLoss(weight=weights_sat)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "2c65e30a-91ce-463b-9168-8b09c827f7a6",
      "metadata": {
        "deletable": false,
        "editable": false,
        "id": "2c65e30a-91ce-463b-9168-8b09c827f7a6",
        "tags": []
      },
      "source": [
        "### Optimizer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c86049ef-0f11-4654-85d9-2797a86e33c7",
      "metadata": {
        "id": "c86049ef-0f11-4654-85d9-2797a86e33c7"
      },
      "outputs": [],
      "source": [
        "optimizer = torch.optim.Adam(model.parameters(), lr=config[\"lr\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "70ee60d0-823d-46fd-957c-e9341fa3aa25",
      "metadata": {
        "deletable": false,
        "editable": false,
        "id": "70ee60d0-823d-46fd-957c-e9341fa3aa25",
        "tags": []
      },
      "source": [
        "### Pytorch lightning module"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4d2f5dc2-4163-490a-8e41-f8742ba5efc9",
      "metadata": {
        "id": "4d2f5dc2-4163-490a-8e41-f8742ba5efc9"
      },
      "outputs": [],
      "source": [
        "seg_module = SegmentationTask(\n",
        "    model=model,\n",
        "    num_classes=config[\"num_classes\"],\n",
        "    criterion=nn.ModuleList([criterion_vhr, criterion_hr]),\n",
        "    optimizer=optimizer,\n",
        "    config=config\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0b9d2cbb-f1ac-4f8d-88f5-f347c0a9e6e8",
      "metadata": {
        "deletable": false,
        "editable": false,
        "id": "0b9d2cbb-f1ac-4f8d-88f5-f347c0a9e6e8",
        "tags": []
      },
      "source": [
        "### Callbacks"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "2041fd30-c39e-49ee-9cb7-f00774ffb1b2",
      "metadata": {
        "id": "2041fd30-c39e-49ee-9cb7-f00774ffb1b2"
      },
      "outputs": [],
      "source": [
        "# Callbacks\n",
        "\n",
        "ckpt_callback = ModelCheckpoint(\n",
        "    monitor=\"val_loss\",\n",
        "    dirpath=os.path.join(out_dir,\"checkpoints\"),\n",
        "    filename=\"ckpt-{epoch:02d}-{val_loss:.2f}\"+'_'+config[\"out_model_name\"],\n",
        "    save_top_k=1,\n",
        "    mode=\"min\",\n",
        "    save_weights_only=True, # can be changed accordingly\n",
        ")\n",
        "\n",
        "# early_stop_callback = EarlyStopping(\n",
        "#     monitor=\"\",\n",
        "#     min_delta=0.00,\n",
        "#     patience=30, # if no improvement after 30 epoch, stop learning.\n",
        "#     mode=\"min\",\n",
        "# )\n",
        "\n",
        "# prog_rate = TQDMProgressBar(refresh_rate=config[\"progress_rate\"])\n",
        "\n",
        "# callbacks = [\n",
        "#     ckpt_callback,\n",
        "#     early_stop_callback,\n",
        "#     prog_rate,\n",
        "# ]\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d63d51b5-249d-4838-bbb6-dca8ba997ff9",
      "metadata": {
        "deletable": false,
        "editable": false,
        "id": "d63d51b5-249d-4838-bbb6-dca8ba997ff9",
        "tags": []
      },
      "source": [
        "### Loggers"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c56fda75-3b52-4a4e-874f-083c06b446e6",
      "metadata": {
        "id": "c56fda75-3b52-4a4e-874f-083c06b446e6"
      },
      "outputs": [],
      "source": [
        "logger = TensorBoardLogger(\n",
        "    save_dir=out_dir,\n",
        "    name=Path(\"tensorboard_logs\"+'_'+config['out_model_name']).as_posix()\n",
        ")\n",
        "\n",
        "loggers = [\n",
        "    logger\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7248b8bc-638e-4085-bd83-7baa94cd094b",
      "metadata": {
        "id": "7248b8bc-638e-4085-bd83-7baa94cd094b",
        "tags": []
      },
      "source": [
        "## <font color='#90c149'>Launch the training</font>\n",
        "\n",
        "<br/><hr>\n",
        "\n",
        "Defining a `Trainer` allows to automate tasks, such as enabling/disabling grads, running the dataloaders or invoking the callbacks when needed.\n",
        "<hr><br/>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "19ceb0de-9a75-43e9-be60-a7a4a1640d7b",
      "metadata": {
        "id": "19ceb0de-9a75-43e9-be60-a7a4a1640d7b"
      },
      "outputs": [],
      "source": [
        "#### instanciation of  Trainer\n",
        "trainer = Trainer(\n",
        "        accelerator=config[\"accelerator\"],\n",
        "        devices=config[\"gpus_per_node\"],\n",
        "        strategy=config[\"strategy\"],\n",
        "        num_nodes=config[\"num_nodes\"],\n",
        "        max_epochs=config[\"num_epochs\"],\n",
        "        num_sanity_val_steps=0,\n",
        "        # callbacks = callbacks,\n",
        "        logger=loggers,\n",
        "        enable_progress_bar = config[\"enable_progress_bar\"],\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "cFE9DTm3A8df",
      "metadata": {
        "deletable": false,
        "editable": false,
        "id": "cFE9DTm3A8df",
        "tags": []
      },
      "source": [
        "<br/><hr>\n",
        "\n",
        "<font color='#90c149'>Let's launch the training.</font>\n",
        "<br/><hr>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5afd841a-9468-4d68-b2a8-aa5dac2d6f60",
      "metadata": {
        "id": "5afd841a-9468-4d68-b2a8-aa5dac2d6f60"
      },
      "outputs": [],
      "source": [
        "trainer.fit(seg_module, datamodule=dm)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "48978bf6-4d34-4308-ad8a-105efb45ea79",
      "metadata": {
        "deletable": false,
        "editable": false,
        "id": "48978bf6-4d34-4308-ad8a-105efb45ea79",
        "tags": []
      },
      "source": [
        "## <font color='#90c149'>Check metrics on the validation dataset</font>\n",
        "\n",
        "<br/><hr>\n",
        "\n",
        "To give an idea on the training results, we call validate on the trainer to print some metrics. <hr><br/>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8d96d9d7-da3c-4266-9f4b-712f4c923aa6",
      "metadata": {
        "id": "8d96d9d7-da3c-4266-9f4b-712f4c923aa6"
      },
      "outputs": [],
      "source": [
        "trainer.validate(seg_module, datamodule=dm)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "a3984d2d-7ed0-4d9b-8d47-47fce2d88987",
      "metadata": {
        "deletable": false,
        "editable": false,
        "id": "a3984d2d-7ed0-4d9b-8d47-47fce2d88987",
        "tags": []
      },
      "source": [
        "## <font color='#90c149'>Inference and predictions export</font>\n",
        "\n",
        "<br/><hr>\n",
        "\n",
        "For inference, we define a new callback, `PredictionWriter`, which is used to export the predictions on the test dataset.<br/><br/>\n",
        "<font color='#90c149'>Note:</font> the callback exports the files with the mandotary formatting of outputs (files named <font color='red'><b> PRED_{ID].tif</b></font>, with datatype <font color='red'><b>uint8</b></font> and <font color='red'><b>LZW</b></font> compression), using Pillow.\n",
        "Check the <font color='#D7881C'><em>writer.py</em></font> file for details.<br/><br/>\n",
        "\n",
        "We instantiate a new `Trainer` with this newly defined callback and call predict.\n",
        "<hr><br/>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "329190c2-040d-4a77-a27d-3239286e01e4",
      "metadata": {
        "id": "329190c2-040d-4a77-a27d-3239286e01e4"
      },
      "outputs": [],
      "source": [
        "# Predict\n",
        "writer_callback = PredictionWriter(\n",
        "    output_dir = os.path.join(out_dir, \"predictions\"+\"_\"+config[\"out_model_name\"]),\n",
        "    write_interval = \"batch\",\n",
        ")\n",
        "\n",
        "#### instanciation of prediction Trainer\n",
        "trainer = Trainer(\n",
        "    accelerator = config[\"accelerator\"],\n",
        "    devices = config[\"gpus_per_node\"],\n",
        "    strategy = config[\"strategy\"],\n",
        "    num_nodes = config[\"num_nodes\"],\n",
        "    callbacks = [writer_callback],\n",
        "    enable_progress_bar = config[\"enable_progress_bar\"],\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e36febba-49d1-410f-a60f-1fba724e9f76",
      "metadata": {
        "id": "e36febba-49d1-410f-a60f-1fba724e9f76"
      },
      "outputs": [],
      "source": [
        "trainer.predict(seg_module, datamodule=dm)\n",
        "\n",
        "@rank_zero_only\n",
        "def print_finish():\n",
        "    print('--  [FINISHED.]  --', f'output dir : {out_dir}', sep='\\n')\n",
        "print_finish()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "09b7e1da-f928-4b73-8014-a5ed49fd38ea",
      "metadata": {
        "deletable": false,
        "editable": false,
        "id": "09b7e1da-f928-4b73-8014-a5ed49fd38ea",
        "tags": []
      },
      "source": [
        "## <font color='#90c149'>Visual checking of predictions</font>\n",
        "\n",
        "<br/><hr>\n",
        "\n",
        "<font color='#90c149'>For the test set, obviously, you do not have access to the masks.</font> Nevertheless, we can visually display some predictions alongside the RGB images.<br/><br/>\n",
        "\n",
        "First, we create lists containing the paths to the test RGB images (`images_test`) as well as the predicted semantic segmentation masks (`predictions`).<br/><br/>\n",
        "\n",
        "\n",
        "\n",
        "We then display some random couples of predictions together with their corresponding aerial RGB images.<br/><br/>\n",
        "\n",
        "<font color='#90c149'><em>Note 1</em></font>: if you are using the toy dataset, don't expect accurate predictions. A set of $200$ training samples will give limited results.<br/>\n",
        "<font color='#90c149'><em>Note 2</em></font>: rasterio will yield a <em>NotGeoreferencedWarning</em> regarding the predictions files. This is normal as the prediction files have been written without any geographical information, which is expected by rasterio. This kind of information is not important for assessing the model outputs, so we can just omit the warning.\n",
        "<hr><br/>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a1e31b62-590a-4a61-9b95-4be828159cff",
      "metadata": {
        "id": "a1e31b62-590a-4a61-9b95-4be828159cff"
      },
      "outputs": [],
      "source": [
        "from data_display import display_predictions, get_data_paths\n",
        "\n",
        "images_test = dict_test[\"PATH_IMG\"]\n",
        "predictions = sorted(list(get_data_paths(Path(os.path.join(out_dir, \"predictions\"+\"_\"+config[\"out_model_name\"])), 'PRED*.tif')), key=lambda x: int(x.split('_')[-1][:-4]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "520c9ed4-6ff2-4d07-a15b-86d039e8f2dc",
      "metadata": {
        "id": "520c9ed4-6ff2-4d07-a15b-86d039e8f2dc"
      },
      "outputs": [],
      "source": [
        "display_predictions(images_test, predictions, nb_samples=2)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "84bb236d-2c92-4140-88af-6a67d74e9b05",
      "metadata": {
        "id": "84bb236d-2c92-4140-88af-6a67d74e9b05",
        "tags": []
      },
      "source": [
        "## <font color='#90c149'>Metric calculation: mIoU</font>\n",
        "\n",
        "<br/><hr>\n",
        "\n",
        "As mentioned before, the masks of the test set are not available. However, the following cell describes the code that is used to calculate the metric used over the test set and to consequently rank the best models. Again, the toy dataset contains $50$ test pastches, while the full FLAIR-two dataset contains $16,050$ test patches.<br/><br/>\n",
        "\n",
        "The calculation of the mean Intersection-over-Union (`mIou`) is based on the confusion matrix $C$, which is determined for each test patch. The confusion matrices are subsequently summed providing the confusion matrix describing the test set. Per-class IoU, defined as the ratio between true positives divided by the sum of false positives, false negatives and true positives is calculated from the summed confusion matrix as follows: <br/><br/>\n",
        "    $$\n",
        "    IoU_i = \\frac{C_{i,i}}\n",
        "    {C_{i,i} + \\sum_{j \\neq i}\\left(C_{i,j} + C_{j,i} \\right)} = \\frac{TP}{TP+FP+FN}\n",
        "    $$\n",
        "<br>\n",
        "The final `mIou` is then the average of the per-class IoUs.\n",
        "\n",
        "\n",
        "<font color='#90c149'><em>Note:</em></font> as the <font color='#90c149'><em>'other'</em></font> class is <font color='#90c149'>not well defined (void)</font>, its IoU is <font color='#90c149'>removed</font> and therefore does not contribute to the calculation of the `mIou`. In other words,  the remaining per-class IoUs (all except 'other') are averaged by 12 and not 13 to obtain the final `mIou`.</font>\n",
        "\n",
        "<hr><br/>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "77290841-332f-4163-8a4e-d395d04626f4",
      "metadata": {
        "id": "77290841-332f-4163-8a4e-d395d04626f4"
      },
      "outputs": [],
      "source": [
        "#truth_msk = config['data']['path_labels_test']\n",
        "#pred_msk  = os.path.join(out_dir, \"predictions\"+\"_\"+config[\"out_model_name\"])\n",
        "#mIou, ious = generate_miou(truth_msk, pred_msk)\n",
        "#print_metrics(mIou, ious)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "03bd09e5-284a-4f0d-b297-111b6bd87bb7",
      "metadata": {
        "id": "03bd09e5-284a-4f0d-b297-111b6bd87bb7",
        "tags": []
      },
      "source": [
        "<br/><br/><br/><br/>\n",
        "\n",
        "### <center><strong>For any feedback, request, suggestion or simply to say hi, we are reachable at : ai-challenge@ign.fr !</strong></center>\n",
        "<br/>\n",
        "<font size=2.5> <b>@IGN, Mai 2023</b></font>\n",
        "<img src=\"https://drive.google.com/uc?export=view&id=14clxUsTGj7i6oXt6q9FQeaxzjIi3biI2\" alt=\"Drawing\"  width=\"100%\"/>"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "gpuType": "T4",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "torch",
      "language": "python",
      "name": "torch"
    },
    "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"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 5
}