Cityscapes_modified_R2UNET.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.5"
    },
    "colab": {
      "name": "Cityscapes_modified_R2UNET.ipynb",
      "provenance": []
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cDYIt1kcmg3m"
      },
      "source": [
        "\n",
        "\n",
        "#### Task Structure\n",
        "\n",
        "```\n",
        "1. Imports \n",
        "2. Hyperparameters\n",
        "       -- Includes learning rate, dataset path, batch size and more.\n",
        "3. Loading Dataset  \n",
        "4. Dataset and Dataloader\n",
        "       -- Initialising dataloader\n",
        "5. Model Architecture  \n",
        "       -- Defining model architecture (R2U-Net and modifications for task 3)\n",
        "6. Loss Function and Optimizers  \n",
        "       -- Adam Optimiser and Cross Entropy Loss\n",
        "7. Functions for Metrics Calculation  \n",
        "       -- Calculation of 5 metrics: Specificity, Senstivity, F1 Score, Accuracy, Jaccard Score\n",
        "8. Train Function \n",
        "9. Validation Function  \n",
        "10. Training Epochs and Validation\n",
        "       -- Perform training and calculate metrics\n",
        "11. Plotting Loss over training Epochs  \n",
        "       -- Show loss values over different parts of iteration\n",
        "12. Plot Evaluation Metrics  \n",
        "       -- Plot of 5 metrics\n",
        "13. Visualizing the results\n",
        "       -- Displays 5 predicted segmentation masks along with original segmentation masks\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jGJCvGefmg3o"
      },
      "source": [
        "### 1. Imports"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "q1G3ntsymg3p"
      },
      "source": [
        "import torch\n",
        "import os\n",
        "import numpy as np\n",
        "import scipy.misc as m\n",
        "from torch.utils import data\n",
        "from torch.utils.data import DataLoader\n",
        "import torch.nn as nn\n",
        "import sklearn.metrics as skm\n",
        "import torch.optim as optim\n",
        "from tqdm import tqdm\n",
        "import matplotlib.pyplot as plt  \n",
        "import torch.nn.functional as F\n",
        "import time"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JeJU8zhtmg3q"
      },
      "source": [
        "### 2. Hyperparameters"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gJy0Yqszmg3q"
      },
      "source": [
        "# replace device accordingly\n",
        "device = torch.device('cuda:7')\n",
        "\n",
        "# replace with location of folder containing \"gtFine\" and \"leftImg8bit\"\n",
        "path_data = \"/raid14/\"\n",
        "\n",
        "learning_rate = 1e-6\n",
        "train_epochs = 8\n",
        "n_classes = 19\n",
        "batch_size = 1\n",
        "num_workers = 1"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XcTqmMh_mg3q"
      },
      "source": [
        "### 3. Loading the dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Be4eJP8Nmg3r"
      },
      "source": [
        "# Adapted from dataset loader written by meetshah1995 with modifications\n",
        "# https://github.com/meetshah1995/pytorch-semseg/blob/master/ptsemseg/loader/cityscapes_loader.py\n",
        "\n",
        "def recursive_glob(rootdir=\".\", suffix=\"\"):\n",
        "    return [\n",
        "        os.path.join(looproot, filename)\n",
        "        for looproot, _, filenames in os.walk(rootdir)\n",
        "        for filename in filenames\n",
        "        if filename.endswith(suffix)\n",
        "    ]\n",
        "\n",
        "\n",
        "class cityscapesLoader(data.Dataset):\n",
        "    colors = [  # [  0,   0,   0],\n",
        "        [128, 64, 128],\n",
        "        [244, 35, 232],\n",
        "        [70, 70, 70],\n",
        "        [102, 102, 156],\n",
        "        [190, 153, 153],\n",
        "        [153, 153, 153],\n",
        "        [250, 170, 30],\n",
        "        [220, 220, 0],\n",
        "        [107, 142, 35],\n",
        "        [152, 251, 152],\n",
        "        [0, 130, 180],\n",
        "        [220, 20, 60],\n",
        "        [255, 0, 0],\n",
        "        [0, 0, 142],\n",
        "        [0, 0, 70],\n",
        "        [0, 60, 100],\n",
        "        [0, 80, 100],\n",
        "        [0, 0, 230],\n",
        "        [119, 11, 32],\n",
        "    ]\n",
        "\n",
        "    # makes a dictionary with key:value. For example 0:[128, 64, 128]\n",
        "    label_colours = dict(zip(range(19), colors))\n",
        "\n",
        "    def __init__(\n",
        "        self,\n",
        "        root,\n",
        "        # which data split to use\n",
        "        split=\"train\",\n",
        "        # transform function activation\n",
        "        is_transform=True,\n",
        "        # image_size to use in transform function\n",
        "        img_size=(512, 1024),\n",
        "    ):\n",
        "        self.root = root\n",
        "        self.split = split\n",
        "        self.is_transform = is_transform\n",
        "        self.n_classes = 19\n",
        "        self.img_size = img_size if isinstance(img_size, tuple) else (img_size, img_size)\n",
        "        self.files = {}\n",
        "\n",
        "        # makes it: /raid11/cityscapes/ + leftImg8bit + train (as we named the split folder this)\n",
        "        self.images_base = os.path.join(self.root, \"leftImg8bit\", self.split)\n",
        "        self.annotations_base = os.path.join(self.root, \"gtFine\", self.split)\n",
        "        \n",
        "        # contains list of all pngs inside all different folders. Recursively iterates \n",
        "        self.files[split] = recursive_glob(rootdir=self.images_base, suffix=\".png\")\n",
        "\n",
        "        self.void_classes = [0, 1, 2, 3, 4, 5, 6, 9, 10, 14, 15, 16, 18, 29, 30, -1]\n",
        "        \n",
        "        # these are 19\n",
        "        self.valid_classes = [7, 8, 11, 12, 13, 17, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 31, 32, 33,\n",
        "        ]\n",
        "        \n",
        "        # these are 19 + 1; \"unlabelled\" is extra\n",
        "        self.class_names = [\n",
        "            \"unlabelled\",\n",
        "            \"road\",\n",
        "            \"sidewalk\",\n",
        "            \"building\",\n",
        "            \"wall\",\n",
        "            \"fence\",\n",
        "            \"pole\",\n",
        "            \"traffic_light\",\n",
        "            \"traffic_sign\",\n",
        "            \"vegetation\",\n",
        "            \"terrain\",\n",
        "            \"sky\",\n",
        "            \"person\",\n",
        "            \"rider\",\n",
        "            \"car\",\n",
        "            \"truck\",\n",
        "            \"bus\",\n",
        "            \"train\",\n",
        "            \"motorcycle\",\n",
        "            \"bicycle\",\n",
        "        ]\n",
        "        \n",
        "        # for void_classes; useful for loss function\n",
        "        self.ignore_index = 250\n",
        "        \n",
        "        # dictionary of valid classes 7:0, 8:1, 11:2\n",
        "        self.class_map = dict(zip(self.valid_classes, range(19)))\n",
        "\n",
        "        if not self.files[split]:\n",
        "            raise Exception(\"No files for split=[%s] found in %s\" % (split, self.images_base))\n",
        "        \n",
        "        # prints number of images found\n",
        "        print(\"Found %d %s images\" % (len(self.files[split]), split))\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.files[self.split])\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "        # path of image\n",
        "        img_path = self.files[self.split][index].rstrip()\n",
        "        \n",
        "        # path of label\n",
        "        lbl_path = os.path.join(\n",
        "            self.annotations_base,\n",
        "            img_path.split(os.sep)[-2],\n",
        "            os.path.basename(img_path)[:-15] + \"gtFine_labelIds.png\",\n",
        "        )\n",
        "\n",
        "        # read image\n",
        "        img = m.imread(img_path)\n",
        "        # convert to numpy array\n",
        "        img = np.array(img, dtype=np.uint8)\n",
        "\n",
        "        # read label\n",
        "        lbl = m.imread(lbl_path)\n",
        "        # encode using encode_segmap function: 0...18 and 250\n",
        "        lbl = self.encode_segmap(np.array(lbl, dtype=np.uint8))\n",
        "\n",
        "        if self.is_transform:\n",
        "            img, lbl = self.transform(img, lbl)\n",
        "        \n",
        "        return img, lbl\n",
        "\n",
        "    def transform(self, img, lbl):       \n",
        "        # Image resize; I think imresize outputs in different format than what it received\n",
        "        img = m.imresize(img, (self.img_size[0], self.img_size[1]))  # uint8 with RGB mode\n",
        "        # change to BGR\n",
        "        img = img[:, :, ::-1]  # RGB -> BGR\n",
        "        # change data type to float64\n",
        "        img = img.astype(np.float64)\n",
        "        # subtract mean\n",
        "        # NHWC -> NCHW\n",
        "        img = img.transpose(2, 0, 1)\n",
        "\n",
        "        \n",
        "        classes = np.unique(lbl)\n",
        "        lbl = lbl.astype(float)\n",
        "        lbl = m.imresize(lbl, (self.img_size[0], self.img_size[1]), \"nearest\", mode=\"F\")\n",
        "        lbl = lbl.astype(int)\n",
        "\n",
        "        if not np.all(classes == np.unique(lbl)):\n",
        "            print(\"WARN: resizing labels yielded fewer classes\")\n",
        "\n",
        "        if not np.all(np.unique(lbl[lbl != self.ignore_index]) < self.n_classes):\n",
        "            print(\"after det\", classes, np.unique(lbl))\n",
        "            raise ValueError(\"Segmentation map contained invalid class values\")\n",
        "\n",
        "        img = torch.from_numpy(img).float()\n",
        "        lbl = torch.from_numpy(lbl).long()\n",
        "\n",
        "        return img, lbl\n",
        "      \n",
        "    def decode_segmap(self, temp):\n",
        "        r = temp.copy()\n",
        "        g = temp.copy()\n",
        "        b = temp.copy()\n",
        "        for l in range(0, self.n_classes):\n",
        "            r[temp == l] = self.label_colours[l][0]\n",
        "            g[temp == l] = self.label_colours[l][1]\n",
        "            b[temp == l] = self.label_colours[l][2]\n",
        "\n",
        "        rgb = np.zeros((temp.shape[0], temp.shape[1], 3))\n",
        "        rgb[:, :, 0] = r / 255.0\n",
        "        rgb[:, :, 1] = g / 255.0\n",
        "        rgb[:, :, 2] = b / 255.0\n",
        "        return rgb\n",
        "\n",
        "    # there are different class 0...33\n",
        "    # we are converting that info to 0....18; and 250 for void classes\n",
        "    # final mask has values 0...18 and 250\n",
        "    def encode_segmap(self, mask):\n",
        "        # !! Comment in code had wrong informtion\n",
        "        # Put all void classes to ignore_index\n",
        "        for _voidc in self.void_classes:\n",
        "            mask[mask == _voidc] = self.ignore_index\n",
        "        for _validc in self.valid_classes:\n",
        "            mask[mask == _validc] = self.class_map[_validc]\n",
        "        return mask"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JfosXqIcmg3t"
      },
      "source": [
        "### 4. Dataset and Dataloader"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Df3FGi--mg3u",
        "outputId": "b7f6aba2-106b-4d57-8283-dabe61b6f806"
      },
      "source": [
        "train_data = cityscapesLoader(\n",
        "    root = path_data, \n",
        "    split='train'\n",
        "    )\n",
        "\n",
        "val_data = cityscapesLoader(\n",
        "    root = path_data, \n",
        "    split='val'\n",
        "    )\n",
        "\n",
        "train_loader = DataLoader(\n",
        "    train_data,\n",
        "    batch_size = batch_size,\n",
        "    shuffle=True,\n",
        "    num_workers = num_workers,\n",
        "    #pin_memory = pin_memory  # gave no significant advantage\n",
        ")\n",
        "\n",
        "val_loader = DataLoader(\n",
        "    val_data,\n",
        "    batch_size = batch_size,\n",
        "    num_workers = num_workers,\n",
        "    #pin_memory = pin_memory  # gave no significant advantage\n",
        ")"
      ],
      "execution_count": null,
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Found 2975 train images\n",
            "Found 500 val images\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MWhmDKFgmg3v"
      },
      "source": [
        "### 5. Model Architecture"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IjR1J0zMmg3w"
      },
      "source": [
        "class Up_Sample_Conv(nn.Module):\n",
        "    def __init__(self, ch_in, ch_out):\n",
        "        super(Up_Sample_Conv, self).__init__()\n",
        "        self.up = nn.Sequential(\n",
        "            nn.Upsample(scale_factor=2), # Nearest neighbour for upsampling are two \n",
        "            nn.Conv2d(ch_in, ch_out, kernel_size=3, stride=1, padding=1, bias=True),\n",
        "            nn.BatchNorm2d(ch_out),\n",
        "            nn.ReLU(inplace=True)\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.up(x)\n",
        "        return x\n",
        "\n",
        "    \n",
        "class Repeat(nn.Module):\n",
        "    def __init__(self, ch_out):\n",
        "        super(Repeat, self).__init__()\n",
        "        self.conv = nn.Sequential(\n",
        "            nn.Conv2d(ch_out, ch_out, kernel_size=3, stride=1, padding=1, bias=True),\n",
        "            nn.BatchNorm2d(ch_out),\n",
        "            nn.ReLU(inplace=True)) \n",
        "#Inplace has been set to TRUE so that it modifies the input directly, without allocating any additional output.\n",
        "\n",
        "    def forward(self, x):\n",
        "        for i in range(2):\n",
        "            if i == 0:\n",
        "                x_rec = self.conv(x)\n",
        "            x_rec = self.conv(x + x_rec)\n",
        "        return x_rec\n",
        "\n",
        "class RR_Conv(nn.Module):\n",
        "    def __init__(self, ch_in, ch_out):\n",
        "        super(RR_Conv, self).__init__()\n",
        "        self.Repeat_block = nn.Sequential(Repeat(ch_out),Repeat(ch_out))\n",
        "        self.Conv = nn.Conv2d(ch_in, ch_out, 1, 1, 0)\n",
        "\n",
        "    def forward(self, input_img):\n",
        "        input_img = self.Conv(input_img)\n",
        "        conv_input_img = self.Repeat_block(input_img)\n",
        "        return input_img + conv_input_img \n",
        "    \n",
        "############\n",
        "############\n",
        "\n",
        "class R2U_Net(nn.Module):\n",
        "    def __init__(self, img_ch=3, output_ch=19):\n",
        "        super(R2U_Net, self).__init__()\n",
        "        \n",
        "        self.MaxPool = nn.MaxPool2d(kernel_size=2, stride=2)\n",
        "        \n",
        "        self.channel_1 = 64 # R2U-net activation maps in first layer\n",
        "        self.channel_2 = 2*self.channel_1\n",
        "        self.channel_3 = 2*self.channel_2\n",
        "        self.channel_4 = 2*self.channel_3\n",
        "        self.channel_5 = 2*self.channel_4\n",
        "        \n",
        "        # For new layer added\n",
        "        self.channel_6 = 2*self.channel_5\n",
        "        \n",
        "        self.channels = [self.channel_1, self.channel_2, self.channel_3, self.channel_4, self.channel_5, self.channel_6]\n",
        "            \n",
        "        '''Performs Convolution and responsible for the encoding part of the architecture'''    \n",
        "        self.Layer1 = RR_Conv(img_ch, self.channels[0])\n",
        "        self.Layer2 = RR_Conv(self.channels[0], self.channels[1])\n",
        "        self.Layer3 = RR_Conv(self.channels[1], self.channels[2])\n",
        "        self.Layer4 = RR_Conv(self.channels[2], self.channels[3])\n",
        "        self.Layer5 = RR_Conv(self.channels[3], self.channels[4])\n",
        "        '''Addition of convolutional layer (Depth increased)'''\n",
        "        self.Layer6 = RR_Conv(self.channels[4], self.channels[5]) # For extra depth\n",
        "\n",
        "        '''Below function calls are responsible for the decoding part of the architeture'''\n",
        "        \n",
        "        '''Upsamples the input and then performs convolution followed by ReLU'''\n",
        "        self.DeConvLayer6 = Up_Sample_Conv(self.channels[5], self.channels[4]) # For extra depth\n",
        "        self.DeConvLayer5 = Up_Sample_Conv(self.channels[4], self.channels[3])\n",
        "        self.DeConvLayer4 = Up_Sample_Conv(self.channels[3],self.channels[2])\n",
        "        self.DeConvLayer3 = Up_Sample_Conv(self.channels[2], self.channels[1])\n",
        "        self.DeConvLayer2 = Up_Sample_Conv(self.channels[1], self.channels[0])\n",
        "        \n",
        "        '''Responsible for computation in Recurrent Residual Blocks'''\n",
        "        self.Up_Layer6 = RR_Conv(self.channels[5], self.channels[4]) # For extra depth\n",
        "        self.Up_Layer5 = RR_Conv(self.channels[4], self.channels[3])\n",
        "        self.Up_Layer4 = RR_Conv(self.channels[3], self.channels[2])\n",
        "        self.Up_Layer3 = RR_Conv(self.channels[2], self.channels[1])\n",
        "        self.Up_Layer2 = RR_Conv(self.channels[1], self.channels[0])\n",
        "        \n",
        "        '''Final output of the architecture needs to have output channels=number of class labels(19)'''\n",
        "        self.Conv = nn.Conv2d(self.channels[0], output_ch, kernel_size=1, stride=1, padding=0)        \n",
        "        \n",
        "    def forward(self, x):\n",
        "        '''Recurrent Convolution'''\n",
        "        conv1 = self.Layer1(x)\n",
        "        mp1 = self.MaxPool(conv1)\n",
        "        conv2 = self.Layer2(mp1)\n",
        "        mp2 = self.MaxPool(conv2)\n",
        "        conv3 = self.Layer3(mp2)\n",
        "        mp3 = self.MaxPool(conv3)\n",
        "        conv4 = self.Layer4(mp3)\n",
        "        mp4 = self.MaxPool(conv4)\n",
        "        conv5 = self.Layer5(mp4)\n",
        "        \n",
        "        '''For one extra depth'''\n",
        "        mp5 = self.MaxPool(conv5)\n",
        "        conv6 = self.Layer6(mp5)\n",
        "        '''--------------------'''\n",
        "\n",
        "        ''' \n",
        "        Decoder part of the architecture which performs \n",
        "        Recurrent up convolution as well as concatention from previous layers \n",
        "        '''\n",
        "        \n",
        "        '''For one extra depth'''\n",
        "        deconv6 = self.DeConvLayer6(conv6)\n",
        "        deconv6 = torch.cat((conv5, deconv6), dim=1)\n",
        "        deconv6 = self.Up_Layer6(deconv6)\n",
        "        '''--------------------'''\n",
        "        \n",
        "        deconv5 = self.DeConvLayer5(deconv6)\n",
        "        deconv5 = torch.cat((conv4, deconv5), dim=1)\n",
        "        deconv5 = self.Up_Layer5(deconv5)\n",
        "        deconv4 = self.DeConvLayer4(deconv5)\n",
        "        deconv4 = torch.cat((conv3, deconv4), dim=1)\n",
        "        deconv4 = self.Up_Layer4(deconv4)\n",
        "        deconv3 = self.DeConvLayer3(deconv4)\n",
        "        deconv3 = torch.cat((conv2, deconv3), dim=1)\n",
        "        deconv3 = self.Up_Layer3(deconv3)\n",
        "        deconv2 = self.DeConvLayer2(deconv3)\n",
        "        deconv2 = torch.cat((conv1, deconv2), dim=1)\n",
        "        deconv2 = self.Up_Layer2(deconv2)\n",
        "        deconv1 = self.Conv(deconv2)\n",
        "\n",
        "        return deconv1"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2GL-3Fw3mg3y"
      },
      "source": [
        "# Instance of the model defined above.\n",
        "model = R2U_Net().to(device)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JS_5_nkfmg3z"
      },
      "source": [
        "### 6. Loss Function and Optimiser"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6Wi2Ppa9mg30"
      },
      "source": [
        "optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n",
        "\n",
        "# Cross Entropy Loss adapted from meetshah1995 to prevent size inconsistencies between model precition \n",
        "# and target label\n",
        "# https://github.com/meetshah1995/pytorch-semseg/blob/master/ptsemseg/loss/loss.py\n",
        "\n",
        "def cross_entropy2d(input, target, weight=None, size_average=True):\n",
        "    n, c, h, w = input.size()\n",
        "    nt, ht, wt = target.size()\n",
        "\n",
        "    # Handle inconsistent size between input and target\n",
        "    if h != ht and w != wt:  # upsample labels\n",
        "        input = F.interpolate(input, size=(ht, wt), mode=\"bilinear\", align_corners=True)\n",
        "\n",
        "    input = input.transpose(1, 2).transpose(2, 3).contiguous().view(-1, c)\n",
        "    target = target.view(-1)\n",
        "    loss = F.cross_entropy(\n",
        "        input, target, weight=weight, size_average=size_average, ignore_index=250\n",
        "    )\n",
        "    return loss"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LvC3SLnjmg30"
      },
      "source": [
        "### 7. Functions for Metrics Calculation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TVmCM6wnmg30"
      },
      "source": [
        "'''We have used skelarn libraries to calculate Accuracy and Jaccard Score'''\n",
        "\n",
        "def get_metrics(gt_label, pred_label):\n",
        "    #Accuracy Score\n",
        "    acc = skm.accuracy_score(gt_label, pred_label, normalize=True)\n",
        "    \n",
        "    #Jaccard Score/IoU\n",
        "    js = skm.jaccard_score(gt_label, pred_label, average='micro')\n",
        "    \n",
        "    result_gm_sh = [acc, js]\n",
        "    return(result_gm_sh)\n",
        "\n",
        "'''\n",
        "Calculation of confusion matrix from :\n",
        "https://github.com/meetshah1995/pytorch-semseg/blob/master/ptsemseg/metrics.py\n",
        "\n",
        "Added modifications to calculate 3 evaluation metrics - \n",
        "Specificity, Senstivity, F1 Score\n",
        "'''\n",
        "\n",
        "class runningScore(object):\n",
        "    def __init__(self, n_classes):\n",
        "        self.n_classes = n_classes\n",
        "        self.confusion_matrix = np.zeros((n_classes, n_classes))\n",
        "\n",
        "    def _fast_hist(self, label_true, label_pred, n_class):\n",
        "        mask = (label_true >= 0) & (label_true < n_class)\n",
        "        hist = np.bincount(\n",
        "            n_class * label_true[mask].astype(int) + label_pred[mask], minlength=n_class ** 2\n",
        "        ).reshape(n_class, n_class)\n",
        "        return hist\n",
        "\n",
        "    def update(self, label_trues, label_preds):\n",
        "        for lt, lp in zip(label_trues, label_preds):\n",
        "            self.confusion_matrix += self._fast_hist(lt.flatten(), lp.flatten(), self.n_classes)\n",
        "\n",
        "    def get_scores(self):\n",
        "        # confusion matrix\n",
        "        hist = self.confusion_matrix\n",
        "        \n",
        "        #              T\n",
        "        #         0    1    2\n",
        "        #    0   TP   FP   FP\n",
        "        #  P 1   FN   TN   TN       This is wrt to class 0\n",
        "        #    2   FN   TN   TN\n",
        "\n",
        "        #         0    1    2\n",
        "        #    0   TP   FP   FP\n",
        "        #  P 1   FP   TP   FP       This is wrt prediction classes; AXIS = 1\n",
        "        #    2   FP   FP   TP \n",
        "\n",
        "        #         0    1    2\n",
        "        #    0   TP   FN   FN\n",
        "        #  P 1   FN   TP   FN       This is wrt true classes; AXIS = 0\n",
        "        #    2   FN   FN   TP   \n",
        "\n",
        "        TP = np.diag(hist)\n",
        "        TN = hist.sum() - hist.sum(axis = 1) - hist.sum(axis = 0) + np.diag(hist)\n",
        "        FP = hist.sum(axis = 1) - TP\n",
        "        FN = hist.sum(axis = 0) - TP\n",
        "        \n",
        "        # 1e-6 was added to prevent corner cases where denominator = 0\n",
        "        \n",
        "        # Specificity: TN / TN + FP\n",
        "        specif_cls = (TN) / (TN + FP + 1e-6)\n",
        "        specif = np.nanmean(specif_cls)\n",
        "        \n",
        "        # Senstivity/Recall: TP / TP + FN\n",
        "        sensti_cls = (TP) / (TP + FN + 1e-6)\n",
        "        sensti = np.nanmean(sensti_cls)\n",
        "        \n",
        "        # Precision: TP / (TP + FP)\n",
        "        prec_cls = (TP) / (TP + FP + 1e-6)\n",
        "        prec = np.nanmean(prec_cls)\n",
        "        \n",
        "        # F1 = 2 * Precision * Recall / Precision + Recall\n",
        "        f1 = (2 * prec * sensti) / (prec + sensti + 1e-6)\n",
        "        \n",
        "        return (\n",
        "            {\n",
        "                \"Specificity\": specif,\n",
        "                \"Senstivity\": sensti,\n",
        "                \"F1\": f1,\n",
        "            }\n",
        "        )\n",
        "\n",
        "    def reset(self):\n",
        "        self.confusion_matrix = np.zeros((self.n_classes, self.n_classes))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e0Vko8IWmg31"
      },
      "source": [
        "### 8. Train Function "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6mBf3dJhmg31"
      },
      "source": [
        "def train(train_loader, model, optimizer, epoch_i, epoch_total):\n",
        "        count = 0\n",
        "        \n",
        "        # List to cumulate loss during iterations\n",
        "        loss_list = []\n",
        "        for (images, labels) in train_loader:\n",
        "            count += 1\n",
        "            \n",
        "            # we used model.eval() below. This is to bring model back to training mood.\n",
        "            model.train()\n",
        "\n",
        "            images = images.to(device)\n",
        "            labels = labels.to(device)\n",
        "            \n",
        "            # Model Prediction\n",
        "            pred = model(images)\n",
        "            \n",
        "            # Loss Calculation\n",
        "            loss = cross_entropy2d(pred, labels)\n",
        "            loss_list.append(loss)\n",
        "\n",
        "            # optimiser\n",
        "            optimizer.zero_grad()\n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "            \n",
        "            # interval to print train statistics\n",
        "            if count % 50 == 0:\n",
        "                fmt_str = \"Image: {:d} in epoch: [{:d}/{:d}]  and Loss: {:.4f}\"\n",
        "                print_str = fmt_str.format(\n",
        "                    count,\n",
        "                    epoch_i + 1,\n",
        "                    epoch_total,\n",
        "                    loss.item()\n",
        "                )\n",
        "                print(print_str)\n",
        "                   \n",
        "#           # break for testing purpose\n",
        "#             if count == 10:\n",
        "#                 break\n",
        "        return(loss_list)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vpHXrraimg31"
      },
      "source": [
        "### 9. Validation Function  "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JpMGYGXcmg32"
      },
      "source": [
        "def validate(val_loader, model, epoch_i):\n",
        "    \n",
        "    # tldr: to make layers behave differently during inference (vs training)\n",
        "    model.eval()\n",
        "    \n",
        "    # enable calculation of confusion matrix for n_classes = 19\n",
        "    running_metrics_val = runningScore(19)\n",
        "    \n",
        "    # empty list to add Accuracy and Jaccard Score Calculations\n",
        "    acc_sh = []\n",
        "    js_sh = []\n",
        "    \n",
        "    with torch.no_grad():\n",
        "        for image_num, (val_images, val_labels) in tqdm(enumerate(val_loader)):\n",
        "            \n",
        "            val_images = val_images.to(device)\n",
        "            val_labels = val_labels.to(device)\n",
        "            \n",
        "            # Model prediction\n",
        "            val_pred = model(val_images)\n",
        "            \n",
        "            # Coverting val_pred from (1, 19, 512, 1024) to (1, 512, 1024)\n",
        "            # considering predictions with highest scores for each pixel among 19 classes\n",
        "            pred = val_pred.data.max(1)[1].cpu().numpy()\n",
        "            gt = val_labels.data.cpu().numpy()\n",
        "            \n",
        "            # Updating Mertics\n",
        "            running_metrics_val.update(gt, pred)\n",
        "            sh_metrics = get_metrics(gt.flatten(), pred.flatten())\n",
        "            acc_sh.append(sh_metrics[0])\n",
        "            js_sh.append(sh_metrics[1])\n",
        "                               \n",
        "#            # break for testing purpose\n",
        "#             if image_num == 10:\n",
        "#                 break                \n",
        "\n",
        "    score = running_metrics_val.get_scores()\n",
        "    running_metrics_val.reset()\n",
        "    \n",
        "    acc_s = sum(acc_sh)/len(acc_sh)\n",
        "    js_s = sum(js_sh)/len(js_sh)\n",
        "    score[\"acc\"] = acc_s\n",
        "    score[\"js\"] = js_s\n",
        "    \n",
        "    print(\"Different Metrics were: \", score)  \n",
        "    return(score)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "90-dhONlmg32"
      },
      "source": [
        "### 10. Training Epochs and Validation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IhQaQQJLmg32"
      },
      "source": [
        "if __name__ == \"__main__\":\n",
        "\n",
        "    # to hold loss values after each epoch\n",
        "    loss_all_epochs = []\n",
        "    \n",
        "    # to hold different metrics after each epoch\n",
        "    Specificity_ = []\n",
        "    Senstivity_ = []\n",
        "    F1_ = []\n",
        "    acc_ = []\n",
        "    js_ = []\n",
        "    \n",
        "    for epoch_i in range(train_epochs):\n",
        "        # training\n",
        "        print(f\"Epoch {epoch_i + 1}\\n-------------------------------\")\n",
        "        t1 = time.time()\n",
        "        loss_i = train(train_loader, model, optimizer, epoch_i, train_epochs)\n",
        "        loss_all_epochs.append(loss_i)\n",
        "        t2 = time.time()\n",
        "        print(\"It took: \", t2-t1, \" unit time\")\n",
        "\n",
        "        # metrics calculation on validation data\n",
        "        dummy_list = validate(val_loader, model, epoch_i)   \n",
        "        \n",
        "        # Add metrics to empty list above\n",
        "        Specificity_.append(dummy_list[\"Specificity\"])\n",
        "        Senstivity_.append(dummy_list[\"Senstivity\"])\n",
        "        F1_.append(dummy_list[\"F1\"])\n",
        "        acc_.append(dummy_list[\"acc\"])\n",
        "        js_.append(dummy_list[\"js\"])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X9t_MAShmg32"
      },
      "source": [
        "### 11. Plotting Loss over training Epochs"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ns2HGe6Xmg33"
      },
      "source": [
        "# loss_all_epochs: contains 2d list of tensors with: (epoch, loss tensor)\n",
        "# converting to 1d list for plotting\n",
        "loss_1d_list = [item for sublist in loss_all_epochs for item in sublist]\n",
        "loss_list_numpy = []\n",
        "for i in range(len(loss_1d_list)):\n",
        "    z = loss_1d_list[i].cpu().detach().numpy()\n",
        "    loss_list_numpy.append(z)\n",
        "plt.xlabel(\"Images used in training epochs\")\n",
        "plt.ylabel(\"Cross Entropy Loss\")\n",
        "plt.plot(loss_list_numpy)\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gXFG63lZmg33"
      },
      "source": [
        "### 12. Plot Evaluation Metrics"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7Fze7SO0mg33"
      },
      "source": [
        "plt.clf()\n",
        "\n",
        "x = [i for i in range(1, train_epochs + 1)]\n",
        "\n",
        "# plot 5 metrics: Specificity, Senstivity, F1 Score, Accuracy, Jaccard Score\n",
        "plt.plot(x,Specificity_, label='Specificity')\n",
        "plt.plot(x,Senstivity_, label='Senstivity')\n",
        "plt.plot(x,F1_, label='F1 Score')\n",
        "plt.plot(x,acc_, label='Accuracy')\n",
        "plt.plot(x,js_, label='Jaccard Score')\n",
        "\n",
        "plt.grid(linestyle = '--', linewidth = 0.5)\n",
        "\n",
        "plt.xlabel(\"Epochs\")\n",
        "plt.ylabel(\"Score\")\n",
        "plt.legend()\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UVfSUC_vmg33"
      },
      "source": [
        "### 13. Visualizing the Results"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QR4Ca3H_mg33"
      },
      "source": [
        "# tldr: to make layers behave differently during inference (vs training)\n",
        "model.eval()\n",
        "\n",
        "with torch.no_grad():\n",
        "    for image_num, (val_images, val_labels) in tqdm(enumerate(val_loader)):\n",
        "\n",
        "        val_images = val_images.to(device)\n",
        "        val_labels = val_labels.to(device)\n",
        "        \n",
        "        # model prediction\n",
        "        val_pred = model(val_images)\n",
        "\n",
        "        # Coverting val_pred from (1, 19, 512, 1024) to (1, 512, 1024)\n",
        "        # considering predictions with highest scores for each pixel among 19 classes        \n",
        "        prediction = val_pred.data.max(1)[1].cpu().numpy()\n",
        "        ground_truth = labels_val.data.cpu().numpy()\n",
        "\n",
        "        # replace 100 to change number of images to print. \n",
        "        # 500 % 100 = 5. So, we will get 5 predictions and ground truths\n",
        "        if image_num % 100 == 0:\n",
        "            \n",
        "            # Model Prediction\n",
        "            decoded_pred = val_data.decode_segmap(prediction[0])\n",
        "            plt.imshow(decoded_pred)\n",
        "            plt.show()\n",
        "            plt.clf()\n",
        "            \n",
        "            # Ground Truth\n",
        "            decode_gt = val_data.decode_segmap(ground_truth[0])\n",
        "            plt.imshow(decode_gt)\n",
        "            plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}