3_train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 23332,
     "status": "ok",
     "timestamp": 1578784207643,
     "user": {
      "displayName": "Md Hasib Zunair 1320262643",
      "photoUrl": "",
      "userId": "12069756592370329757"
     },
     "user_tz": 300
    },
    "id": "lDGrrxzTqQqe",
    "outputId": "f46915b0-05ee-41bd-bd4b-a07ef6c9c0b3"
   },
   "outputs": [],
   "source": [
    "# Run this cell if running this notebook for the first time\n",
    "#!pip install -U --pre segmentation-models --user \n",
    "#!conda install -c conda-forge albumentations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Name experiment\n",
    "experiment_name = \"exp-1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 52
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 883,
     "status": "ok",
     "timestamp": 1578784778530,
     "user": {
      "displayName": "Md Hasib Zunair 1320262643",
      "photoUrl": "",
      "userId": "12069756592370329757"
     },
     "user_tz": 300
    },
    "id": "PqrxTSb8pEXX",
    "outputId": "28b38dc5-4bdd-45f8-af78-6f324c136b15"
   },
   "outputs": [],
   "source": [
    "# Import libs\n",
    "import os \n",
    "import time\n",
    "import cv2\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import skimage.draw\n",
    "import random\n",
    "import keras\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.colors\n",
    "import scipy.io as sio\n",
    "from PIL import Image\n",
    "import scipy\n",
    "import scipy.ndimage\n",
    "import keras.backend as K\n",
    "import segmentation_models as sm\n",
    "\n",
    "\n",
    "def create_directory(directory):\n",
    "    '''\n",
    "    Creates a new folder in the specified directory if the folder doesn't exist.\n",
    "    INPUT\n",
    "        directory: Folder to be created, called as \"folder/\".\n",
    "    OUTPUT\n",
    "        New folder in the current directory.\n",
    "    '''\n",
    "    if not os.path.exists(directory):\n",
    "        os.makedirs(directory)\n",
    "\n",
    "        \n",
    "            \n",
    "    \n",
    "# Define paths\n",
    "dataset_name = \"data_processedv5\"\n",
    "base_path = os.path.abspath(\".\")\n",
    "dataset_path = os.path.join(base_path, \"dataset\", dataset_name)\n",
    "log_path = os.path.join(base_path, \"logs\", experiment_name)\n",
    "patchsize = 96\n",
    "\n",
    "\n",
    "# For full data training\n",
    "full_data_path = os.path.join(base_path, \"dataset\")\n",
    "\n",
    "#print(log_path)\n",
    "#print(os.listdir(dataset_path))\n",
    "\n",
    "\n",
    "# Make directory\n",
    "create_directory(log_path)\n",
    "\n",
    "\n",
    "# Get the data\n",
    "\n",
    "# Training data\n",
    "x_train_dir = os.path.join(dataset_path, 'train', 'images')\n",
    "y_train_dir = os.path.join(dataset_path, 'train', 'masks')\n",
    "\n",
    "# Validation data\n",
    "x_valid_dir = os.path.join(dataset_path, 'val', 'images')\n",
    "y_valid_dir = os.path.join(dataset_path, 'val', 'masks')\n",
    "\n",
    "# Test data\n",
    "x_test_dir = x_valid_dir \n",
    "y_test_dir = y_valid_dir \n",
    "\n",
    "print(\"Training and validation samples: \", len(os.listdir(x_train_dir)), len(os.listdir(x_valid_dir))) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Panoptic Quality Metric\n",
    "\n",
    "Panoptic Segmentation\n",
    "Alexander Kirillov, Kaiming He, Ross Girshick, Carsten Rother and Piotr Dollár\n",
    "arXiv:1801.00868, 2018."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute Panoptic quality metric for each image\n",
    "def Panoptic_quality(ground_truth_image,predicted_image):\n",
    "    TP = 0\n",
    "    FP = 0\n",
    "    FN = 0\n",
    "    sum_IOU = 0\n",
    "    matched_instances = {}# Create a dictionary to save ground truth indices in keys and predicted matched instances as velues\n",
    "                        # It will also save IOU of the matched instance in [indx][1]\n",
    "\n",
    "    # Find matched instances and save it in a dictionary\n",
    "    for i in np.unique(ground_truth_image):\n",
    "        if i == 0:\n",
    "            pass\n",
    "        else:\n",
    "            temp_image = np.array(ground_truth_image)\n",
    "            temp_image = temp_image == i\n",
    "            matched_image = temp_image * predicted_image\n",
    "        \n",
    "            for j in np.unique(matched_image):\n",
    "                if j == 0:\n",
    "                    pass\n",
    "                else:\n",
    "                    pred_temp = predicted_image == j\n",
    "                    intersection = sum(sum(temp_image*pred_temp))\n",
    "                    union = sum(sum(temp_image + pred_temp))\n",
    "                    IOU = intersection/union\n",
    "                    if IOU> 0.5:\n",
    "                        matched_instances [i] = j, IOU \n",
    "                        \n",
    "    # Compute TP, FP, FN and sum of IOU of the matched instances to compute Panoptic Quality               \n",
    "                        \n",
    "    pred_indx_list = np.unique(predicted_image)\n",
    "    pred_indx_list = np.array(pred_indx_list[1:])\n",
    "\n",
    "    # Loop on ground truth instances\n",
    "    for indx in np.unique(ground_truth_image):\n",
    "        if indx == 0:\n",
    "            pass\n",
    "        else:\n",
    "            if indx in matched_instances.keys():\n",
    "                pred_indx_list = np.delete(pred_indx_list, np.argwhere(pred_indx_list == [indx][0]))\n",
    "                TP = TP+1\n",
    "                sum_IOU = sum_IOU+matched_instances[indx][1]\n",
    "            else:\n",
    "                FN = FN+1\n",
    "    FP = len(np.unique(pred_indx_list))\n",
    "    PQ = sum_IOU/(TP+0.5*FP+0.5*FN)\n",
    "    \n",
    "    return PQ"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "L6YBZyCW4znX"
   },
   "source": [
    "#### Dataloader and utility functions¶"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "SrfJ9V7QpEcb"
   },
   "outputs": [],
   "source": [
    "# Helper function for data visualization\n",
    "def visualize(**images):\n",
    "    \"\"\"PLot images in one row.\"\"\"\n",
    "    \n",
    "    norm=plt.Normalize(0,4) # 5 classes including BG\n",
    "    map_name = matplotlib.colors.LinearSegmentedColormap.from_list(\"\", [\"black\", \"red\",\"yellow\",\"blue\", \"green\"])\n",
    "\n",
    "    \n",
    "    n = len(images)\n",
    "    plt.figure(figsize=(16, 5))\n",
    "    for i, (name, image) in enumerate(images.items()):\n",
    "        plt.subplot(1, n, i + 1)\n",
    "        plt.xticks([])\n",
    "        plt.yticks([])\n",
    "        plt.title(' '.join(name.split('_')).title())\n",
    "        plt.imshow(image, cmap=map_name, norm=norm)\n",
    "    plt.show()\n",
    "    \n",
    "    \n",
    "# Helper function for data visualization    \n",
    "def denormalize(x):\n",
    "    \"\"\"Scale image to range 0..1 for correct plot\"\"\"\n",
    "    x_max = np.percentile(x, 98)\n",
    "    x_min = np.percentile(x, 2)    \n",
    "    x = (x - x_min) / (x_max - x_min)\n",
    "    x = x.clip(0, 1)\n",
    "    return x\n",
    "    \n",
    "\n",
    "# Classes for data loading and preprocessing\n",
    "class Dataset:\n",
    "    \"\"\"MoNuSAC Dataset. Read images, apply augmentation and preprocessing transformations.\n",
    "    \n",
    "    Args:\n",
    "        images_dir (str): path to images folder\n",
    "        masks_dir (str): path to segmentation masks folder\n",
    "        class_values (list): values of classes to extract from segmentation mask\n",
    "        augmentation (albumentations.Compose): data transfromation pipeline \n",
    "            (e.g. flip, scale, etc.)\n",
    "        preprocessing (albumentations.Compose): data preprocessing \n",
    "            (e.g. noralization, shape manipulation, etc.)\n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    CLASSES = ['Epithelial', 'Lymphocyte', 'Neutrophil', 'Macrophage', 'unlabelled']\n",
    "\n",
    "    \n",
    "#    label_map = {'Epithelial':1, RED\n",
    "#             'Lymphocyte':2, YELLOW\n",
    "#             'Macrophage':4, GREEN\n",
    "#             'Neutrophil':3, BLUE\n",
    "#            }\n",
    "    \n",
    "    \n",
    "    def __init__(\n",
    "            self, \n",
    "            images_dir, \n",
    "            masks_dir, \n",
    "            classes=None, \n",
    "            augmentation=None, \n",
    "            preprocessing=None,\n",
    "    ):\n",
    "        self.ids = os.listdir(images_dir)\n",
    "        \n",
    "        # Unsorted\n",
    "        #self.images_fps = [os.path.join(images_dir, image_id) for image_id in self.ids]\n",
    "        #self.masks_fps = [os.path.join(masks_dir, image_id) for image_id in self.ids]\n",
    "        \n",
    "        # Sorted\n",
    "        self.images_fps = sorted([os.path.join(images_dir, image_id) for image_id in self.ids])\n",
    "        self.masks_fps = sorted([os.path.join(masks_dir, image_id) for image_id in self.ids])\n",
    "\n",
    "        \n",
    "        # convert str names to class values on masks\n",
    "        self.class_values = [self.CLASSES.index(cls) for cls in classes] # cls used instead of cls.lower()\n",
    "\n",
    "        self.augmentation = augmentation\n",
    "        self.preprocessing = preprocessing\n",
    "    \n",
    "    def __getitem__(self, i):\n",
    "        \n",
    "        # Read data\n",
    "        image = cv2.imread(self.images_fps[i])\n",
    "        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "        #image = skimage.io.imread(self.images_fps[i])\n",
    "        \n",
    "        \n",
    "        mask = cv2.imread(self.masks_fps[i], 0)\n",
    "        #mask = skimage.io.imread(self.masks_fps[i])\n",
    "        \n",
    "        \n",
    "        # Extract certain classes from mask\n",
    "        masks = [(mask == v) for v in self.class_values]\n",
    "        mask = np.stack(masks, axis=-1).astype('float')\n",
    "        \n",
    "        # Add background if mask is not binary\n",
    "        if mask.shape[-1] != 1:\n",
    "            background = 1 - mask.sum(axis=-1, keepdims=True)\n",
    "            mask = np.concatenate((mask, background), axis=-1)\n",
    "        \n",
    "        # Apply augmentations\n",
    "        if self.augmentation:\n",
    "            sample = self.augmentation(image=image, mask=mask)\n",
    "            image, mask = sample['image'], sample['mask']\n",
    "        \n",
    "        # Apply preprocessing\n",
    "        if self.preprocessing:\n",
    "            sample = self.preprocessing(image=image, mask=mask)\n",
    "            image, mask = sample['image'], sample['mask']\n",
    "            \n",
    "        return image, mask\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.ids)\n",
    "    \n",
    "    \n",
    "class Dataloder(keras.utils.Sequence):\n",
    "    \"\"\"Load data from dataset and form batches\n",
    "    \n",
    "    Args:\n",
    "        dataset: instance of Dataset class for image loading and preprocessing.\n",
    "        batch_size: Integet number of images in batch.\n",
    "        shuffle: Boolean, if `True` shuffle image indexes each epoch.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, dataset, batch_size=1, shuffle=False):\n",
    "        self.dataset = dataset\n",
    "        self.batch_size = batch_size\n",
    "        self.shuffle = shuffle\n",
    "        self.indexes = np.arange(len(dataset))\n",
    "\n",
    "        self.on_epoch_end()\n",
    "\n",
    "    def __getitem__(self, i):\n",
    "        \n",
    "        # collect batch data\n",
    "        start = i * self.batch_size\n",
    "        stop = (i + 1) * self.batch_size\n",
    "        data = []\n",
    "        for j in range(start, stop):\n",
    "            data.append(self.dataset[j])\n",
    "        \n",
    "        # Transpose list of lists\n",
    "        batch = [np.stack(samples, axis=0) for samples in zip(*data)]\n",
    "        \n",
    "        return batch\n",
    "    \n",
    "    def __len__(self):\n",
    "        \"\"\"Denotes the number of batches per epoch\"\"\"\n",
    "        return len(self.indexes) // self.batch_size\n",
    "    \n",
    "    def on_epoch_end(self):\n",
    "        \"\"\"Callback function to shuffle indexes each epoch\"\"\"\n",
    "        if self.shuffle:\n",
    "            self.indexes = np.random.permutation(self.indexes)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test the dataset class "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load a random image from the images folder\n",
    "image_fns = sorted(next(os.walk(x_train_dir))[2])\n",
    "gt_fns = sorted(next(os.walk(y_train_dir))[2])\n",
    "print(image_fns[:3], gt_fns[:3])\n",
    "\n",
    "# Get random number\n",
    "idx = random.randrange(len(image_fns))\n",
    "print(\"Index: \", idx)\n",
    "\n",
    "\n",
    "# Read image and mask\n",
    "image = skimage.io.imread(os.path.join(x_train_dir, image_fns[idx]))\n",
    "gt = skimage.io.imread(os.path.join(y_train_dir, gt_fns[idx]))\n",
    "\n",
    "# Flags\n",
    "assert image.shape[:2] == gt.shape, \"Wrong image or ground truth!\"\n",
    "assert image.dtype == gt.dtype, \"Wrong data types!\"\n",
    "\n",
    "print(image.shape, gt.shape)\n",
    "\n",
    "val1 = gt.flatten()\n",
    "print(\"Ground truth classes: \", np.unique(val1))\n",
    "\n",
    "#label_map = {'Epithelial':1, RED\n",
    "#             'Lymphocyte':2, YELLOW\n",
    "#             'Macrophage':4, GREEN\n",
    "#             'Neutrophil':3, BLUE\n",
    "#            }\n",
    "\n",
    "import matplotlib.colors\n",
    "# Stolen from https://stackoverflow.com/questions/16834861/create-own-colormap-using-matplotlib-and-plot-color-scale \n",
    "norm=plt.Normalize(0,4) # 5 classes including BG\n",
    "map_name = matplotlib.colors.LinearSegmentedColormap.from_list(\"\", [\"black\", \"red\",\"yellow\",\"blue\", \"green\"])\n",
    "\n",
    "f, axarr = plt.subplots(1,2, figsize=(8,8))\n",
    "\n",
    "# idx = 94 all classes\n",
    "\n",
    "axarr[0].imshow(image, cmap=map_name, norm=norm)\n",
    "axarr[1].imshow(gt, cmap=map_name, norm=norm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# One image and mask\n",
    "\n",
    "#label_map = {'Epithelial':1,\n",
    "#         'Lymphocyte':2,\n",
    "#         'Macrophage':4,\n",
    "#         'Neutrophil':3,\n",
    "#        }\n",
    "\n",
    "med = ['Epithelial', 'Lymphocyte', 'Neutrophil', 'Macrophage']\n",
    "\n",
    "dataset = Dataset(x_train_dir, y_train_dir, classes=med)\n",
    "image, mask = dataset[idx] \n",
    "\n",
    "print(image.shape, mask.shape)\n",
    "visualize(\n",
    "    image=image, \n",
    "    Epithelial_mask = mask[..., 1].squeeze(),\n",
    "    Lymphocyte_mask = mask[..., 2].squeeze(),\n",
    "    Neutrophil_mask = mask[..., 3].squeeze(),\n",
    "    Macrophage_mask = mask[..., 4].squeeze(),\n",
    "    background_mask = mask[..., 0].squeeze(),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 213
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 1255,
     "status": "ok",
     "timestamp": 1578784526143,
     "user": {
      "displayName": "Md Hasib Zunair 1320262643",
      "photoUrl": "",
      "userId": "12069756592370329757"
     },
     "user_tz": 300
    },
    "id": "TthLvAJdpEiP",
    "outputId": "9b0b5108-dc63-4b4a-ee19-ff4acfd292a9"
   },
   "outputs": [],
   "source": [
    "#label_map = {'Epithelial':1,\n",
    "#         'Lymphocyte':2,\n",
    "#         'Macrophage':4,\n",
    "#         'Neutrophil':3,\n",
    "#        }\n",
    "\n",
    "med = ['Epithelial', 'Lymphocyte', 'Neutrophil', 'Macrophage']\n",
    "\n",
    "dataset = Dataset(x_train_dir, y_train_dir, classes=med)\n",
    "\n",
    "for i in range(3):\n",
    "    \n",
    "    image, mask = dataset[random.randrange(len(dataset))] # get some sample\n",
    "    print(image.shape, mask.shape)\n",
    "    \n",
    "    visualize(\n",
    "        image=image, \n",
    "        Epithelial_mask = mask[..., 1].squeeze(),\n",
    "        Lymphocyte_mask = mask[..., 2].squeeze(),\n",
    "        Neutrophil_mask = mask[..., 3].squeeze(),\n",
    "        Macrophage_mask = mask[..., 4].squeeze(),\n",
    "        background_mask = mask[..., 0].squeeze(),\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "b-hv5EQ15Sz_"
   },
   "source": [
    "\n",
    "#### Augmentations\n",
    "\n",
    "    horizontal flip\n",
    "    affine transforms\n",
    "    perspective transforms\n",
    "    brightness/contrast/colors manipulations\n",
    "    image bluring and sharpening\n",
    "    gaussian noise\n",
    "    random crops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "qYa5PpxRqzxV"
   },
   "outputs": [],
   "source": [
    "import albumentations as A\n",
    "\n",
    "def round_clip_0_1(x, **kwargs):\n",
    "    return x.round().clip(0, 1)\n",
    "\n",
    "# define heavy augmentations\n",
    "def get_training_augmentation():\n",
    "    train_transform = [\n",
    "\n",
    "        A.HorizontalFlip(p=0.5),\n",
    "\n",
    "        A.ShiftScaleRotate(scale_limit=0.5, rotate_limit=0, shift_limit=0.1, p=1, border_mode=0),\n",
    "\n",
    "        #A.PadIfNeeded(min_height=patchsize, min_width=patchsize, always_apply=True, border_mode=0), # patchsize instead of 320\n",
    "        A.RandomCrop(height=patchsize, width=patchsize, always_apply=True), # patchsize instead of 320\n",
    "\n",
    "        A.IAAAdditiveGaussianNoise(p=0.2),\n",
    "        #A.IAAPerspective(p=0.5),\n",
    "\n",
    "        #A.OneOf(\n",
    "         #   [\n",
    "          #      A.CLAHE(p=1),\n",
    "           #     A.RandomBrightness(p=1),\n",
    "            #    A.RandomGamma(p=1),\n",
    "           # ],\n",
    "           # p=0.9,\n",
    "       # ),\n",
    "\n",
    "        #A.OneOf(\n",
    "       #     [\n",
    "     #           A.IAASharpen(p=1),\n",
    "       #         A.Blur(blur_limit=3, p=1),\n",
    "       #         A.MotionBlur(blur_limit=3, p=1),\n",
    "      #      ],\n",
    "       #     p=0.9,\n",
    "       # ),\n",
    "\n",
    "        #A.OneOf(\n",
    "        #    [\n",
    "         #       A.RandomContrast(p=1),\n",
    "        #        A.HueSaturationValue(p=1),\n",
    "        #    ],\n",
    "       #     p=0.9,\n",
    "       # ),\n",
    "       # A.Lambda(mask=round_clip_0_1)\n",
    "    ]\n",
    "    return A.Compose(train_transform)\n",
    "\n",
    "\n",
    "def get_validation_augmentation():\n",
    "    \"\"\"Add paddings to make image shape divisible by 96\"\"\"\n",
    "    test_transform = [\n",
    "        A.PadIfNeeded(patchsize, patchsize)\n",
    "    ]\n",
    "    return A.Compose(test_transform)\n",
    "\n",
    "def get_preprocessing(preprocessing_fn):\n",
    "    \"\"\"Construct preprocessing transform\n",
    "    \n",
    "    Args:\n",
    "        preprocessing_fn (callbale): data normalization function \n",
    "            (can be specific for each pretrained neural network)\n",
    "    Return:\n",
    "        transform: albumentations.Compose\n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    _transform = [\n",
    "        A.Lambda(image=preprocessing_fn),\n",
    "    ]\n",
    "    return A.Compose(_transform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 299
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 964,
     "status": "ok",
     "timestamp": 1578784673774,
     "user": {
      "displayName": "Md Hasib Zunair 1320262643",
      "photoUrl": "",
      "userId": "12069756592370329757"
     },
     "user_tz": 300
    },
    "id": "4Tn2jJgvq_71",
    "outputId": "2f5ef237-a674-4b1b-9ae8-1624cdba70d6"
   },
   "outputs": [],
   "source": [
    "# After agument\n",
    "\n",
    "dataset = Dataset(x_train_dir, y_train_dir, classes=med, augmentation=get_training_augmentation())\n",
    "\n",
    "for i in range(3):\n",
    "    image, mask = dataset[random.randrange(len(dataset))] \n",
    "    print(image.shape, mask.shape)\n",
    "    visualize(\n",
    "        image=image, \n",
    "        epi_mask=mask[..., 1].squeeze(),\n",
    "        lym_mask=mask[..., 2].squeeze(),\n",
    "        neu_mask=mask[..., 3].squeeze(),\n",
    "        mac_mask=mask[..., 4].squeeze(),\n",
    "        background_mask=mask[..., 0].squeeze(),\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build and compile network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "mxEbDwzFsEDI"
   },
   "outputs": [],
   "source": [
    "BACKBONE = 'efficientnetb3' \n",
    "BATCH_SIZE = 128\n",
    "CLASSES = ['Epithelial', 'Lymphocyte', 'Neutrophil', 'Macrophage'] \n",
    "LR = 1e-4   # 1e-4(0.0001) \n",
    "EPOCHS = 100\n",
    "\n",
    "preprocess_input = sm.get_preprocessing(BACKBONE)\n",
    "\n",
    "\n",
    "# Define network parameters\n",
    "n_classes = 1 if len(CLASSES) == 1 else (len(CLASSES) + 1)  # case for binary and multiclass segmentation\n",
    "#print(\"Classes: \", n_classes)\n",
    "activation = 'sigmoid' if n_classes == 1 else 'softmax'\n",
    "\n",
    "# Create model\n",
    "model = None\n",
    "model = sm.Unet(BACKBONE, classes=n_classes, activation=activation) # Train all layers\n",
    "\n",
    "# Define optomizer\n",
    "optim = keras.optimizers.Adam(LR)\n",
    "\n",
    "# multiplication factor for computing loss\n",
    "class_weight_array = np.array([0.5, 0.5, 2, 2])\n",
    "\n",
    "# calc loss over the 4 classes\n",
    "class_index_array =  np.array([1,2,3,4]) \n",
    "\n",
    "\n",
    "dice_loss = sm.losses.DiceLoss(class_weights = class_weight_array, class_indexes = class_index_array)\n",
    "jacard_loss = sm.losses.JaccardLoss(class_weights = class_weight_array, class_indexes = class_index_array)\n",
    "focal_loss = sm.losses.CategoricalFocalLoss(alpha=0.25, gamma=2.0, class_indexes = class_index_array)\n",
    "\n",
    "\n",
    "# Initial loss\n",
    "#total_loss = dice_loss + (1 * focal_loss)\n",
    "\n",
    "# Mod 1\n",
    "total_loss = (1 * focal_loss) + (1 * jacard_loss) # best combination\n",
    "\n",
    "# Define metrics\n",
    "metrics = [sm.metrics.IOUScore(threshold=0.5), sm.metrics.FScore(threshold=0.5)]\n",
    "\n",
    "# Compile model with defined optimozer, loss and metrics\n",
    "model.compile(optim, total_loss, metrics)\n",
    "model.summary() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 72
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 1179,
     "status": "ok",
     "timestamp": 1578786916240,
     "user": {
      "displayName": "Md Hasib Zunair 1320262643",
      "photoUrl": "",
      "userId": "12069756592370329757"
     },
     "user_tz": 300
    },
    "id": "PZAJu8P7sq5b",
    "outputId": "701b50ba-4968-4929-f829-bed6d6a223a3"
   },
   "outputs": [],
   "source": [
    "# Dataset for train images\n",
    "train_dataset = Dataset(\n",
    "    x_train_dir, \n",
    "    y_train_dir, \n",
    "    classes=CLASSES,\n",
    "    #augmentation=get_training_augmentation(),\n",
    "    #preprocessing=get_preprocessing(preprocess_input),\n",
    ")\n",
    "\n",
    "# Dataset for validation images\n",
    "valid_dataset = Dataset(\n",
    "    x_valid_dir, \n",
    "    y_valid_dir, \n",
    "    classes=CLASSES ,\n",
    "    #augmentation=get_validation_augmentation(),\n",
    "    #preprocessing=get_preprocessing(preprocess_input),\n",
    ")\n",
    "\n",
    "train_dataloader = Dataloder(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\n",
    "valid_dataloader = Dataloder(valid_dataset, batch_size=BATCH_SIZE, shuffle=False)\n",
    "\n",
    "# Check shapes for errors\n",
    "assert train_dataloader[0][0].shape == (BATCH_SIZE, patchsize, patchsize, 3)\n",
    "assert train_dataloader[0][1].shape == (BATCH_SIZE, patchsize, patchsize, n_classes)\n",
    "\n",
    "\n",
    "# Define callbacks for learning rate scheduling, logging and best checkpoints saving\n",
    "callbacks = [\n",
    "    keras.callbacks.ModelCheckpoint('{}/{}.h5'.format(log_path, experiment_name), monitor='val_loss', save_best_only=True, mode='min'),\n",
    "    keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.1, verbose=1, patience=5, mode='min'), ## new_lr = lr * factor # 5\n",
    "    keras.callbacks.EarlyStopping(monitor='val_loss', min_delta=0, verbose=1, patience=15, mode='min', restore_best_weights=True), # 8\n",
    "    keras.callbacks.CSVLogger('{}/training.csv'.format(log_path))\n",
    "]\n",
    "\n",
    "\n",
    "# Test train loader\n",
    "x , y = next(iter(valid_dataloader))\n",
    "print(x.shape, y.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train model\n",
    "\n",
    "start_time = time.time()\n",
    "\n",
    "\n",
    "history = model.fit_generator(\n",
    "    train_dataloader, \n",
    "    steps_per_epoch=len(train_dataloader), \n",
    "    epochs=EPOCHS, \n",
    "    callbacks=callbacks, \n",
    "    validation_data=valid_dataloader, \n",
    "    validation_steps=len(valid_dataloader),  # val samples = batch size * no of steps\n",
    ")\n",
    "\n",
    "end_time = time.time()\n",
    "print(\"--- Time taken to train : %s hours ---\" % ((end_time - start_time)//3600))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 350
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 3257203,
     "status": "ok",
     "timestamp": 1578790180537,
     "user": {
      "displayName": "Md Hasib Zunair 1320262643",
      "photoUrl": "",
      "userId": "12069756592370329757"
     },
     "user_tz": 300
    },
    "id": "GdqAiyaCtAbA",
    "outputId": "d3d67c42-78da-4c70-f688-099103d31b7e"
   },
   "outputs": [],
   "source": [
    "# Plot training & validation loss values\n",
    "# b, g, r, y, o, -g, -m,\n",
    "\n",
    "plt.figure(figsize=(15, 5))\n",
    "plt.subplot(131)\n",
    "plt.plot(history.history['loss'],linewidth=4)\n",
    "plt.plot(history.history['val_loss'],linewidth=4)\n",
    "plt.title('{} loss'.format(experiment_name))\n",
    "#plt.ylabel('Loss')\n",
    "plt.xlabel('Epoch')\n",
    "plt.legend(['Train', 'Test'], loc='upper left')\n",
    "plt.grid(True)\n",
    "\n",
    "# Plot training & validation iou_score values\n",
    "\n",
    "plt.subplot(132)\n",
    "plt.plot(history.history['iou_score'],linewidth=4)\n",
    "plt.plot(history.history['val_iou_score'],linewidth=4)\n",
    "plt.title('{} IOU score'.format(experiment_name))\n",
    "#plt.ylabel('iou_score')\n",
    "plt.xlabel('Epoch')\n",
    "plt.grid(True)\n",
    "plt.legend(['Train', 'Test'], loc='upper left')\n",
    "\n",
    "# Plot training & validation loss values\n",
    "\n",
    "plt.subplot(133)\n",
    "plt.plot(history.history['f1-score'],linewidth=4)\n",
    "plt.plot(history.history['val_f1-score'],linewidth=4)\n",
    "plt.title('{} F1 score'.format(experiment_name))\n",
    "#plt.ylabel('Loss')\n",
    "plt.xlabel('Epoch')\n",
    "plt.legend(['Train', 'Test'], loc='upper left')\n",
    "plt.grid(True)\n",
    "plt.savefig('{}/{}_graph.png'.format(log_path, experiment_name), dpi=300)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "AWuPhmWl6Vp2"
   },
   "source": [
    "#### Model Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 72
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 1207,
     "status": "ok",
     "timestamp": 1578790181777,
     "user": {
      "displayName": "Md Hasib Zunair 1320262643",
      "photoUrl": "",
      "userId": "12069756592370329757"
     },
     "user_tz": 300
    },
    "id": "EdCicpPztDPI",
    "outputId": "200b0080-8ec5-4992-e462-cf72be01a3d3"
   },
   "outputs": [],
   "source": [
    "# Validation data is test data here\n",
    "\n",
    "test_dataset = Dataset(\n",
    "    x_test_dir, \n",
    "    y_test_dir, \n",
    "    classes=CLASSES,\n",
    "    #augmentation=get_validation_augmentation(),\n",
    "    #preprocessing=get_preprocessing(preprocess_input),\n",
    ")\n",
    "\n",
    "test_dataloader = Dataloder(test_dataset, batch_size=1, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#x, y = next(iter(test_dataloader))\n",
    "#x.shape, y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load best weights\n",
    "#model=None\n",
    "model.load_weights('{}/{}.h5'.format(log_path, experiment_name))\n",
    "\n",
    "#import efficientnet.tfkeras\n",
    "#from tensorflow.keras.models import load_model\n",
    "#model = None\n",
    "#model = load_model('{}/{}.h5'.format(log_path, experiment_name), compile=False)\n",
    "print(\"Model loaded\")\n",
    "\n",
    "\n",
    "# Run on validation set\n",
    "scores = model.evaluate_generator(test_dataloader, verbose=1)\n",
    "\n",
    "\n",
    "mean_scores = []\n",
    "\n",
    "# Print scores\n",
    "print(\"Loss: {:.5}\".format(scores[0]))\n",
    "mean_scores.append(scores[0])\n",
    "\n",
    "for metric, value in zip(metrics, scores[1:]):\n",
    "    print(\"mean {}: {:.5}\".format(metric.__name__, value))\n",
    "    mean_scores.append(value)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Compute mean PQ on validation set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pqs = []\n",
    "pqs_f = []\n",
    "\n",
    "for i in tqdm(range(len(test_dataset))):\n",
    "    \n",
    "    image, gt_mask = test_dataset[i]\n",
    "    image = np.expand_dims(image, axis=0)\n",
    "    pr_mask = model.predict(image)\n",
    "    \n",
    "    gt_mask = np.argmax(gt_mask.squeeze(), axis=-1)\n",
    "    pr_mask = np.argmax(pr_mask.squeeze(), axis=-1)\n",
    "    \n",
    "    # Post processing to refine predictions\n",
    "    pred_filt = cv2.medianBlur(pr_mask.astype(np.uint8), 5)\n",
    "    pred_filt = pred_filt.astype(np.uint8)\n",
    "    \n",
    "    pq = Panoptic_quality(gt_mask, pr_mask)\n",
    "    pqf = Panoptic_quality(gt_mask, pred_filt)\n",
    "    \n",
    "    pqs.append(pq)\n",
    "    pqs_f.append(pqf)\n",
    "    \n",
    "\n",
    "pqs = np.mean(np.array(pqs))\n",
    "pqs_f = np.mean(np.array(pqs_f))\n",
    "\n",
    "mean_scores.append(pqs)\n",
    "mean_scores.append(pqs_f)\n",
    "\n",
    "print(\"Mean PQ without filer: {:.3}\".format( pqs))\n",
    "print(\"Mean PQ without filer: {:.3}\".format( pqs_f))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "| CFG | Loss   |PQ   |IoU   |F1   | Notes |\n",
    "|------|------|------|------|------|------|\n",
    "|   p-unet1 | 0.92| 55.6, 56.0| 83.1| 86.4| weights=[0.5, 0.5, 2, 2], bs=128, lr=10e-4, total_loss=(1 * focal_loss) + (1 * jacard_loss)|\n",
    "|   p-unet2 | 0.89| 55.6, 55.9| 83.1| 86.9| weights=[0.5, 0.5, 2, 2], bs=128, lr=10e-4, total_loss=(1 * focal_loss) + (1 * jacard_loss)|"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Email logs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pynotify import send_email, send_email_with_attachment\n",
    "\n",
    "# loss, iou, f1 and pq\n",
    "vals = [sc for sc in mean_scores]\n",
    "\n",
    "files = ['{}/{}_graph.png'.format(log_path, experiment_name)]\n",
    "\n",
    "subject = \"Experiment results\"\n",
    "message = \"Loss: {:.3}\".format(vals[0]) + \" PQ no MedFil: {:.3}\".format(vals[3]) + \" PQ with MedFil: {:.3}\".format(vals[4]) + \" IOU: {:.3}\".format(vals[1]) + \" F1: {:.3}\".format(vals[2])\n",
    "\n",
    "dest = \"hasibzunair@gmail.com\" \n",
    "\n",
    "# sends an email\n",
    "send_email_with_attachment(dest, files, subject, message)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "6XsiIrWu6SVT"
   },
   "source": [
    "#### Visualization of results on test dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 35413,
     "status": "ok",
     "timestamp": 1578790216040,
     "user": {
      "displayName": "Md Hasib Zunair 1320262643",
      "photoUrl": "",
      "userId": "12069756592370329757"
     },
     "user_tz": 300
    },
    "id": "NqQCLP2vtUyB",
    "outputId": "ceccaaa4-9187-45f2-d587-644b5aa92916"
   },
   "outputs": [],
   "source": [
    "n = 5\n",
    "ids = np.random.choice(np.arange(len(test_dataset)), size=n)\n",
    "\n",
    "for i in ids:\n",
    "    \n",
    "    image, gt_mask = test_dataset[i]\n",
    "    image = np.expand_dims(image, axis=0)\n",
    "    pr_mask = model.predict(image)\n",
    "    \n",
    "    print(image.shape, gt_mask.shape, pr_mask.shape)\n",
    "    \n",
    "    gt_mask = np.argmax(gt_mask.squeeze(), axis=-1)\n",
    "    pr_mask = np.argmax(pr_mask.squeeze(), axis=-1)\n",
    "    print(\"Uniques in label and predicted\", np.unique(gt_mask), np.unique(pr_mask))\n",
    "    \n",
    "    visualize(\n",
    "        image=denormalize(image.squeeze()),\n",
    "        #image= image.squeeze(),\n",
    "        gt_mask = gt_mask,\n",
    "        pr_mask = pr_mask\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Done!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "multiseg_camvid.ipynb",
   "provenance": []
  },
  "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}