2DUNet.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ahDkQz1edLMp"
      },
      "source": [
        "MISA 3D Brain MRI Segmentation using 2D UNet\n",
        "\n",
        "1. Preprocessing - Bias Correction\n",
        "\n",
        "2. Method - Patch Based\n",
        "\n",
        "3. Data Augmentation - Yes"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HMIG5gOwa-Nb"
      },
      "source": [
        "##Importing the Libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "spA2vbr3a0j7",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "1b9aa496-3eb0-473a-89c4-e52ea50eea5d"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: simpleitk in /usr/local/lib/python3.7/dist-packages (2.1.1)\n",
            "Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount(\"/content/drive\", force_remount=True).\n"
          ]
        }
      ],
      "source": [
        "import cv2\n",
        "import glob\n",
        "import warnings\n",
        "import scipy.misc\n",
        "import numpy as np\n",
        "import nibabel as nib\n",
        "!pip install simpleitk\n",
        "import SimpleITK as sitk\n",
        "from scipy import ndimage\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import tensorflow as tf\n",
        "from tensorflow import keras\n",
        "from tensorflow.keras import layers\n",
        "from keras.preprocessing.image import ImageDataGenerator\n",
        "\n",
        "from google.colab import drive\n",
        "drive._mount(\"/content/drive\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IEFUu8FjYvaz"
      },
      "source": [
        "##Defining the Parameters"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eerKWEtxUg_x"
      },
      "outputs": [],
      "source": [
        "# Image Parameters\n",
        "IMAGE_SIZE = (256, 128, 256)\n",
        "\n",
        "# Training, Testing and Validation Parameters\n",
        "TRAINING_VOLUMES = [0, 1, 2, 3, 4, 5, 6, 7, 8]\n",
        "VALIDATION_VOLUMES = [9]\n",
        "\n",
        "# Hyperparameters\n",
        "N_CLASSES = 4\n",
        "N_INPUT_CHANNELS = 1\n",
        "PATCH_SIZE = (32, 32)\n",
        "PATCH_STRIDE = (32, 32)\n",
        "\n",
        "# Data Preparation Parameters\n",
        "CONTENT_THRESHOLD = 0.3 # To Get Rid of Useless Information in the Image\n",
        "\n",
        "# Training Parameters\n",
        "N_EPOCHS = 200\n",
        "BATCH_SIZE = 64\n",
        "PATIENCE = 20\n",
        "MODEL_FNAME_PATTERN = 'model.h5'\n",
        "OPTIMISER = 'Adam'\n",
        "LOSS = 'categorical_crossentropy'\n",
        "dropout_rate = 0.40"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kakQO5X9E0oR"
      },
      "source": [
        "##Define UNet Architecture"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "flPaOqv0bZ3-",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "1e4a7510-bb64-4076-b7f2-74bcf11938bf"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'def get_unet(img_size=PATCH_SIZE, n_classes=N_CLASSES, n_input_channels=N_INPUT_CHANNELS, scale=1):\\n    inputs = keras.Input(shape=img_size + (n_input_channels, ))\\n\\n    # Encoding Path of the UNet\\n    conv1 = layers.Conv2D(32*scale, (3, 3), padding=\"same\", activation=\\'relu\\')(inputs)\\n    drop1 = layers.Dropout(rate=dropout_rate)(conv1, training=True)\\n    max1 = layers.MaxPooling2D((2, 2))(drop1)\\n    # max1 = layers.MaxPooling2D((2, 2))(conv1)\\n\\n    conv2 = layers.Conv2D(64*scale, (3, 3), padding=\"same\", activation=\\'relu\\')(max1)\\n    drop2 = layers.Dropout(rate=dropout_rate)(conv2, training=True)\\n    # max2 = layers.MaxPooling2D((2, 2))(conv2)\\n    max2 = layers.MaxPooling2D((2, 2))(drop2)\\n\\n    conv3 = layers.Conv2D(128*scale, (3, 3), padding=\"same\", activation=\\'relu\\')(max2)\\n    drop3 = layers.Dropout(rate=dropout_rate)(conv3, training=True)\\n    # max3 = layers.MaxPooling2D((2, 2))(conv3)\\n    max3 = layers.MaxPooling2D((2, 2))(drop3)\\n\\n    lat = layers.Conv2D(256*scale, (3, 3), padding=\"same\", activation=\\'relu\\')(max3)\\n    drop4 = layers.Dropout(rate=dropout_rate)(lat, training=True)\\n\\n    # Decoding Path of the UNet\\n    #up1 = layers.UpSampling2D((2, 2))(lat)\\n    up1 = layers.UpSampling2D((2, 2))(drop4)\\n    concat1 = layers.concatenate([conv3, up1], axis=-1)\\n    conv4 = layers.Conv2D(128*scale, (3, 3), padding=\"same\", activation=\\'relu\\')(concat1)\\n    drop5 = layers.Dropout(rate=dropout_rate)(conv4, training=True)\\n    \\n    #up2 = layers.UpSampling2D((2, 2))(conv4)\\n    up2 = layers.UpSampling2D((2, 2))(drop5)\\n    concat2 = layers.concatenate([conv2, up2], axis=-1)\\n    conv5 = layers.Conv2D(64*scale, (3, 3), padding=\"same\", activation=\\'relu\\')(concat2)\\n    drop6 = layers.Dropout(rate=dropout_rate)(conv5, training=True)\\n    \\n    #up3 = layers.UpSampling2D((2, 2))(conv5)\\n    up3 = layers.UpSampling2D((2, 2))(drop6)\\n    concat3 = layers.concatenate([conv1, up3], axis=-1)\\n    conv6 = layers.Conv2D(32*scale, (3, 3), padding=\"same\", activation=\\'relu\\')(concat3)\\n    drop7 = layers.Dropout(rate=dropout_rate)(conv6, training=True)\\n    \\n    #outputs = layers.Conv2D(n_classes, (1, 1), activation=\"softmax\")(conv6)\\n    outputs = layers.Conv2D(n_classes, (1, 1), activation=\"softmax\")(drop7)\\n\\n    model = keras.Model(inputs, outputs)\\n\\n    return model'"
            ]
          },
          "metadata": {},
          "execution_count": 266
        }
      ],
      "source": [
        "\"\"\"def get_unet(img_size=PATCH_SIZE, n_classes=N_CLASSES, n_input_channels=N_INPUT_CHANNELS, scale=1):\n",
        "    inputs = keras.Input(shape=img_size + (n_input_channels, ))\n",
        "\n",
        "    # Encoding Path of the UNet\n",
        "    conv1 = layers.Conv2D(32*scale, (3, 3), padding=\"same\", activation='relu')(inputs)\n",
        "    drop1 = layers.Dropout(rate=dropout_rate)(conv1, training=True)\n",
        "    max1 = layers.MaxPooling2D((2, 2))(drop1)\n",
        "    # max1 = layers.MaxPooling2D((2, 2))(conv1)\n",
        "\n",
        "    conv2 = layers.Conv2D(64*scale, (3, 3), padding=\"same\", activation='relu')(max1)\n",
        "    drop2 = layers.Dropout(rate=dropout_rate)(conv2, training=True)\n",
        "    # max2 = layers.MaxPooling2D((2, 2))(conv2)\n",
        "    max2 = layers.MaxPooling2D((2, 2))(drop2)\n",
        "\n",
        "    conv3 = layers.Conv2D(128*scale, (3, 3), padding=\"same\", activation='relu')(max2)\n",
        "    drop3 = layers.Dropout(rate=dropout_rate)(conv3, training=True)\n",
        "    # max3 = layers.MaxPooling2D((2, 2))(conv3)\n",
        "    max3 = layers.MaxPooling2D((2, 2))(drop3)\n",
        "\n",
        "    lat = layers.Conv2D(256*scale, (3, 3), padding=\"same\", activation='relu')(max3)\n",
        "    drop4 = layers.Dropout(rate=dropout_rate)(lat, training=True)\n",
        "\n",
        "    # Decoding Path of the UNet\n",
        "    #up1 = layers.UpSampling2D((2, 2))(lat)\n",
        "    up1 = layers.UpSampling2D((2, 2))(drop4)\n",
        "    concat1 = layers.concatenate([conv3, up1], axis=-1)\n",
        "    conv4 = layers.Conv2D(128*scale, (3, 3), padding=\"same\", activation='relu')(concat1)\n",
        "    drop5 = layers.Dropout(rate=dropout_rate)(conv4, training=True)\n",
        "    \n",
        "    #up2 = layers.UpSampling2D((2, 2))(conv4)\n",
        "    up2 = layers.UpSampling2D((2, 2))(drop5)\n",
        "    concat2 = layers.concatenate([conv2, up2], axis=-1)\n",
        "    conv5 = layers.Conv2D(64*scale, (3, 3), padding=\"same\", activation='relu')(concat2)\n",
        "    drop6 = layers.Dropout(rate=dropout_rate)(conv5, training=True)\n",
        "    \n",
        "    #up3 = layers.UpSampling2D((2, 2))(conv5)\n",
        "    up3 = layers.UpSampling2D((2, 2))(drop6)\n",
        "    concat3 = layers.concatenate([conv1, up3], axis=-1)\n",
        "    conv6 = layers.Conv2D(32*scale, (3, 3), padding=\"same\", activation='relu')(concat3)\n",
        "    drop7 = layers.Dropout(rate=dropout_rate)(conv6, training=True)\n",
        "    \n",
        "    #outputs = layers.Conv2D(n_classes, (1, 1), activation=\"softmax\")(conv6)\n",
        "    outputs = layers.Conv2D(n_classes, (1, 1), activation=\"softmax\")(drop7)\n",
        "\n",
        "    model = keras.Model(inputs, outputs)\n",
        "\n",
        "    return model\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from keras.layers import BatchNormalization, Activation"
      ],
      "metadata": {
        "id": "KvgLKx5y73i_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def get_unet(img_size=PATCH_SIZE, n_classes=N_CLASSES, n_input_channels=N_INPUT_CHANNELS, scale=1):\n",
        "    inputs = keras.Input(shape=img_size + (n_input_channels, ))\n",
        "\n",
        "    # Encoding Path of the UNet (32-64-128-256-512)\n",
        "    conv1   = Conv2D(32*scale, (3, 3), padding=\"same\", activation='relu')(inputs)\n",
        "    drop1   = Dropout(rate=dropout_rate)(conv1, training=True)\n",
        "    max1    = MaxPooling2D((2, 2))(drop1)\n",
        "\n",
        "    conv2   = Conv2D(64*scale, (3, 3), padding=\"same\", activation='relu')(max1)\n",
        "    drop2   = Dropout(rate=dropout_rate)(conv2, training=True)\n",
        "    max2    = MaxPooling2D((2, 2))(drop2)\n",
        "\n",
        "    conv3   = Conv2D(128*scale, (3, 3), padding=\"same\", activation='relu')(max2)\n",
        "    drop3   = Dropout(rate=dropout_rate)(conv3, training=True)\n",
        "    max3    = MaxPooling2D((2, 2))(drop3)\n",
        "\n",
        "    conv4   = Conv2D(256*scale, (3, 3), padding=\"same\", activation='relu')(max3)\n",
        "    drop4   = Dropout(rate=dropout_rate)(conv4, training=True)\n",
        "    max4    = MaxPooling2D((2, 2))(drop4)\n",
        "\n",
        "    lat     = Conv2D(512*scale, (3, 3), padding=\"same\", activation='relu')(max4)\n",
        "    drop5   = Dropout(rate=dropout_rate)(lat, training=True)\n",
        "\n",
        "    # Decoding Path of the UNet\n",
        "    up1     = UpSampling2D((2, 2))(drop5)\n",
        "    concat1 = concatenate([conv4, up1], axis=-1)\n",
        "    conv5   = Conv2D(256*scale, (3, 3), padding=\"same\", activation='relu')(concat1)\n",
        "    drop6   = Dropout(rate=dropout_rate)(conv5, training=True)\n",
        "    \n",
        "    up2     = UpSampling2D((2, 2))(drop6)\n",
        "    concat2 = concatenate([conv3, up2], axis=-1)\n",
        "    conv6   = Conv2D(128*scale, (3, 3), padding=\"same\", activation='relu')(concat2)\n",
        "    drop7   = Dropout(rate=dropout_rate)(conv6, training=True)\n",
        "    \n",
        "    up3     = UpSampling2D((2, 2))(drop7)\n",
        "    concat3 = concatenate([conv2, up3], axis=-1)\n",
        "    conv7   = Conv2D(64*scale, (3, 3), padding=\"same\", activation='relu')(concat3)\n",
        "    drop8   = Dropout(rate=dropout_rate)(conv7, training=True)\n",
        "\n",
        "    up4     = UpSampling2D((2, 2))(drop8)\n",
        "    concat4 = concatenate([conv1, up4], axis=-1)\n",
        "    conv8   = Conv2D(32*scale, (3, 3), padding=\"same\", activation='relu')(concat4)\n",
        "    drop9   = Dropout(rate=dropout_rate)(conv8, training=True)\n",
        "    \n",
        "    outputs = Conv2D(n_classes, (1, 1), activation=\"softmax\")(drop9)\n",
        "\n",
        "    model   = Model(inputs, outputs)\n",
        "\n",
        "    return model"
      ],
      "metadata": {
        "id": "wcAmfYyQoM_q"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ByBoNkPGWKrN"
      },
      "source": [
        "##Generate Bias Corrected Images"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "\"\"\"\n",
        "def N4(inputImagePath, maskImagePath, outputPath):\n",
        "  # inputImagePath = input('Enter the path of the image : ')\n",
        "  inputImage = sitk.ReadImage(inputImagePath)\n",
        "\n",
        "  print(\"N4 bias correction runs.\")\n",
        "\n",
        "  # maskImage = sitk.ReadImage(\"06-t1c_mask.nii.gz\")\n",
        "  maskImage = sitk.OtsuThreshold(inputImage,0,1,200)\n",
        "  maskImagePath = input('Enter the name of the mask image to be saved : ')\n",
        "  sitk.WriteImage(maskImage, maskImagePath)\n",
        "  print(\"Mask image is saved.\")\n",
        "\n",
        "  inputImage = sitk.Cast(inputImage,sitk.sitkFloat32)\n",
        "\n",
        "  corrector = sitk.N4BiasFieldCorrectionImageFilter();\n",
        "\n",
        "  output = corrector.Execute(inputImage,maskImage)\n",
        "\n",
        "  outputPath = input(\"Enter the name of the Bias Field Corrected Image :\")\n",
        "  sitk.WriteImage(output,outputPath)\n",
        "  print(\"Finished N4 Bias Field Correction.....\")\n",
        "\"\"\""
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6qJA60k2n2Cl",
        "outputId": "908e22ec-82c5-4768-cbc7-36f0e2494a93"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'\\ndef N4(inputImagePath, maskImagePath, outputPath):\\n  # inputImagePath = input(\\'Enter the path of the image : \\')\\n  inputImage = sitk.ReadImage(inputImagePath)\\n\\n  print(\"N4 bias correction runs.\")\\n\\n  # maskImage = sitk.ReadImage(\"06-t1c_mask.nii.gz\")\\n  maskImage = sitk.OtsuThreshold(inputImage,0,1,200)\\n  maskImagePath = input(\\'Enter the name of the mask image to be saved : \\')\\n  sitk.WriteImage(maskImage, maskImagePath)\\n  print(\"Mask image is saved.\")\\n\\n  inputImage = sitk.Cast(inputImage,sitk.sitkFloat32)\\n\\n  corrector = sitk.N4BiasFieldCorrectionImageFilter();\\n\\n  output = corrector.Execute(inputImage,maskImage)\\n\\n  outputPath = input(\"Enter the name of the Bias Field Corrected Image :\")\\n  sitk.WriteImage(output,outputPath)\\n  print(\"Finished N4 Bias Field Correction.....\")\\n'"
            ]
          },
          "metadata": {},
          "execution_count": 268
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7e5tn28UYaO6"
      },
      "outputs": [],
      "source": [
        "def data_bias_correction(setName) :\n",
        "  \n",
        "  data_file = '/content/drive/My Drive/MISA/Normal Segmentations/data/{}/*'.format(setName)\n",
        "\n",
        "  for filename in glob.glob(data_file):\n",
        "    \n",
        "    #print(filename)\n",
        "    name = filename[-7:]\n",
        "    #print(name)\n",
        "    print(\"Working on image {0}\".format(name))\n",
        "\n",
        "    img_path = '/content/drive/My Drive/MISA/Normal Segmentations/data/{}/{}/{}.nii.gz'.format(setName, name, name)\n",
        "    inputImage = sitk.ReadImage(img_path)\n",
        "    \n",
        "    mask_path = '/content/drive/My Drive/MISA/Normal Segmentations/data/{}/{}/{}_mask.nii.gz'.format(setName, name, name)\n",
        "    maskImage = sitk.OtsuThreshold(inputImage,0,1,200)\n",
        "    sitk.WriteImage(maskImage, mask_path)\n",
        "    print(\"Mask image is saved.\")\n",
        "\n",
        "    inputImage = sitk.Cast(inputImage,sitk.sitkFloat32)\n",
        "    corrector = sitk.N4BiasFieldCorrectionImageFilter();\n",
        "    output = corrector.Execute(inputImage,maskImage)\n",
        "\n",
        "    bias_corrected_path = '/content/drive/My Drive/MISA/Normal Segmentations/data/{}/{}/{}_bias_corrected.nii.gz'.format(setName, name, name)\n",
        "    sitk.WriteImage(output,bias_corrected_path)\n",
        "    print(\"Finished N4 Bias Field Correction.....\")\n",
        "\n",
        "    #bias_corrected_path = '/content/drive/My Drive/MISA/Normal Segmentations/data/{}/{}/{}_bias_corrected.nii.gz'.format(setName, name, name)\n",
        "    #N4(img_path, mask_path, bias_corrected_path)\n",
        "    #N4('/content/drive/MyDrive/MISA/Normal Segmentations/data/Training_Set/IBSR_01/IBSR_01.nii.gz', '/content/hello/', '/content/hello/')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "O2InxCFyek-_",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "7b05ccc9-3f45-4ebb-af94-17b4dd95d2e5"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "\"# Calling the Bias Removal Function (N4)\\ndata_bias_correction('Training_Set')\\ndata_bias_correction('Validation_Set')\\ndata_bias_correction('Test_Set')\""
            ]
          },
          "metadata": {},
          "execution_count": 270
        }
      ],
      "source": [
        "\"\"\"# Calling the Bias Removal Function (N4)\n",
        "data_bias_correction('Training_Set')\n",
        "data_bias_correction('Validation_Set')\n",
        "data_bias_correction('Test_Set')\"\"\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4Z1E1IglE-0w"
      },
      "source": [
        "##Loading the training and validation data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cgujGm9tX_Rj"
      },
      "outputs": [],
      "source": [
        "def load_data_bias(image_size, setName) :\n",
        "  \n",
        "  data_file = '/content/drive/My Drive/MISA/Normal Segmentations/data/{}/*'.format(setName)\n",
        "\n",
        "  folders = glob.glob(data_file)\n",
        "  n_volumes = len(folders)\n",
        "  \n",
        "  volumes = np.zeros((n_volumes, *image_size, 1))\n",
        "  labels = np.zeros((n_volumes, *image_size, 1))\n",
        "\n",
        "  i = 0\n",
        "\n",
        "  for filename in glob.glob(data_file):\n",
        "    \n",
        "    #print(filename)\n",
        "    name = filename[-7:]\n",
        "    #print(name)\n",
        "\n",
        "    img_data = nib.load('/content/drive/My Drive/MISA/Normal Segmentations/data/{}/{}/{}_bias_corrected.nii.gz'.format(setName, name, name))\n",
        "    img_data_temp = img_data.get_fdata()\n",
        "    img_data_temp = img_data_temp.reshape((*image_size, 1))\n",
        "    #print(img_data_temp.shape)\n",
        "    volumes[i] = img_data_temp\n",
        "\n",
        "    seg_data = nib.load('/content/drive/My Drive/MISA/Normal Segmentations/data/{}/{}/{}_seg.nii.gz'.format(setName, name, name))\n",
        "    labels[i] = seg_data.get_fdata()\n",
        "    \n",
        "    print(\"Working on image {0}\".format(name))\n",
        "    i = i+1\n",
        "\n",
        "  return (volumes, labels)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IoSMUWMWgKQj",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "5747ea6b-feac-4dda-a0f9-282335a932eb"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Working on image IBSR_05\n",
            "Working on image IBSR_03\n",
            "Working on image IBSR_08\n",
            "Working on image IBSR_04\n",
            "Working on image IBSR_01\n",
            "Working on image IBSR_16\n",
            "Working on image IBSR_18\n",
            "Working on image IBSR_07\n",
            "Working on image IBSR_09\n",
            "Working on image IBSR_06\n",
            "Working on image IBSR_14\n",
            "Working on image IBSR_17\n",
            "Working on image IBSR_12\n",
            "Working on image IBSR_13\n"
          ]
        }
      ],
      "source": [
        "(t_volumes, t_labels) = load_data_bias(IMAGE_SIZE, 'Training_Set')\n",
        "(v_volumes, v_labels) = load_data_bias(IMAGE_SIZE, 'Validation_Set')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gQ3ZCG2Wc92_"
      },
      "source": [
        "Visualising the Training Images"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NuRrxxLM9TCp",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "d3ff5a99-cd92-4976-d9be-287ff79e880d"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "\"check_vol = t_volumes[1,:,:,:,:]\\ncheck_vol = check_vol.reshape((256, 128, 256))\\nrotated_vol = ndimage.rotate(check_vol, 90)\\nplt.axis('off')\\nplt.imshow(rotated_vol[:, :, 150], cmap='gray')\\nplt.show()\""
            ]
          },
          "metadata": {},
          "execution_count": 273
        }
      ],
      "source": [
        "\"\"\"check_vol = t_volumes[1,:,:,:,:]\n",
        "check_vol = check_vol.reshape((256, 128, 256))\n",
        "rotated_vol = ndimage.rotate(check_vol, 90)\n",
        "plt.axis('off')\n",
        "plt.imshow(rotated_vol[:, :, 150], cmap='gray')\n",
        "plt.show()\"\"\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uZ2BevUuIkGs"
      },
      "source": [
        "## Denoising the Volumes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9wQy2mw7ETMt",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "6e22708b-d28c-48d3-f0cc-cf07dabf7cbb"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'def anisodiff3(stack,niter=1,kappa=50,gamma=0.1,step=(1.,1.,1.),option=1,ploton=False):\\n\\n    # ...you could always diffuse each color channel independently if you\\n    # really want\\n    if stack.ndim == 4:\\n        warnings.warn(\"Only grayscale stacks allowed, converting to 3D matrix\")\\n        stack = stack.mean(3)\\n\\n    # initialize output array\\n    stack = stack.astype(\\'float32\\')\\n    stackout = stack.copy()\\n\\n    # initialize some internal variables\\n    deltaS = np.zeros_like(stackout)\\n    deltaE = deltaS.copy()\\n    deltaD = deltaS.copy()\\n    NS = deltaS.copy()\\n    EW = deltaS.copy()\\n    UD = deltaS.copy()\\n    gS = np.ones_like(stackout)\\n    gE = gS.copy()\\n    gD = gS.copy()\\n\\n    # create the plot figure, if requested\\n    if ploton:\\n        import pylab as pl\\n        from time import sleep\\n\\n        showplane = stack.shape[0]//2\\n\\n        fig = pl.figure(figsize=(20,5.5),num=\"Anisotropic diffusion\")\\n        ax1,ax2 = fig.add_subplot(1,2,1),fig.add_subplot(1,2,2)\\n\\n        ax1.imshow(stack[showplane,...].squeeze(),interpolation=\\'nearest\\')\\n        ih = ax2.imshow(stackout[showplane,...].squeeze(),interpolation=\\'nearest\\',animated=True)\\n        ax1.set_title(\"Original stack (Z = %i)\" %showplane)\\n        ax2.set_title(\"Iteration 0\")\\n\\n        fig.canvas.draw()\\n\\n    for ii in range(niter):\\n\\n        # calculate the diffs\\n        deltaD[:-1,: ,:  ] = np.diff(stackout,axis=0)\\n        deltaS[:  ,:-1,: ] = np.diff(stackout,axis=1)\\n        deltaE[:  ,: ,:-1] = np.diff(stackout,axis=2)\\n\\n        # conduction gradients (only need to compute one per dim!)\\n        if option == 1:\\n            gD = np.exp(-(deltaD/kappa)**2.)/step[0]\\n            gS = np.exp(-(deltaS/kappa)**2.)/step[1]\\n            gE = np.exp(-(deltaE/kappa)**2.)/step[2]\\n        elif option == 2:\\n            gD = 1./(1.+(deltaD/kappa)**2.)/step[0]\\n            gS = 1./(1.+(deltaS/kappa)**2.)/step[1]\\n            gE = 1./(1.+(deltaE/kappa)**2.)/step[2]\\n\\n        # update matrices\\n        D = gD*deltaD\\n        E = gE*deltaE\\n        S = gS*deltaS\\n\\n        # subtract a copy that has been shifted \\'Up/North/West\\' by one\\n        # pixel. don\\'t as questions. just do it. trust me.\\n        UD[:] = D\\n        NS[:] = S\\n        EW[:] = E\\n        UD[1:,: ,: ] -= D[:-1,:  ,:  ]\\n        NS[: ,1:,: ] -= S[:  ,:-1,:  ]\\n        EW[: ,: ,1:] -= E[:  ,:  ,:-1]\\n\\n        # update the image\\n        stackout += gamma*(UD+NS+EW)\\n\\n        if ploton:\\n            iterstring = \"Iteration %i\" %(ii+1)\\n            ih.set_data(stackout[showplane,...].squeeze())\\n            ax2.set_title(iterstring)\\n            fig.canvas.draw()\\n            # sleep(0.01)\\n\\n    return stackout'"
            ]
          },
          "metadata": {},
          "execution_count": 274
        }
      ],
      "source": [
        "\"\"\"def anisodiff3(stack,niter=1,kappa=50,gamma=0.1,step=(1.,1.,1.),option=1,ploton=False):\n",
        "\n",
        "    # ...you could always diffuse each color channel independently if you\n",
        "    # really want\n",
        "    if stack.ndim == 4:\n",
        "        warnings.warn(\"Only grayscale stacks allowed, converting to 3D matrix\")\n",
        "        stack = stack.mean(3)\n",
        "\n",
        "    # initialize output array\n",
        "    stack = stack.astype('float32')\n",
        "    stackout = stack.copy()\n",
        "\n",
        "    # initialize some internal variables\n",
        "    deltaS = np.zeros_like(stackout)\n",
        "    deltaE = deltaS.copy()\n",
        "    deltaD = deltaS.copy()\n",
        "    NS = deltaS.copy()\n",
        "    EW = deltaS.copy()\n",
        "    UD = deltaS.copy()\n",
        "    gS = np.ones_like(stackout)\n",
        "    gE = gS.copy()\n",
        "    gD = gS.copy()\n",
        "\n",
        "    # create the plot figure, if requested\n",
        "    if ploton:\n",
        "        import pylab as pl\n",
        "        from time import sleep\n",
        "\n",
        "        showplane = stack.shape[0]//2\n",
        "\n",
        "        fig = pl.figure(figsize=(20,5.5),num=\"Anisotropic diffusion\")\n",
        "        ax1,ax2 = fig.add_subplot(1,2,1),fig.add_subplot(1,2,2)\n",
        "\n",
        "        ax1.imshow(stack[showplane,...].squeeze(),interpolation='nearest')\n",
        "        ih = ax2.imshow(stackout[showplane,...].squeeze(),interpolation='nearest',animated=True)\n",
        "        ax1.set_title(\"Original stack (Z = %i)\" %showplane)\n",
        "        ax2.set_title(\"Iteration 0\")\n",
        "\n",
        "        fig.canvas.draw()\n",
        "\n",
        "    for ii in range(niter):\n",
        "\n",
        "        # calculate the diffs\n",
        "        deltaD[:-1,: ,:  ] = np.diff(stackout,axis=0)\n",
        "        deltaS[:  ,:-1,: ] = np.diff(stackout,axis=1)\n",
        "        deltaE[:  ,: ,:-1] = np.diff(stackout,axis=2)\n",
        "\n",
        "        # conduction gradients (only need to compute one per dim!)\n",
        "        if option == 1:\n",
        "            gD = np.exp(-(deltaD/kappa)**2.)/step[0]\n",
        "            gS = np.exp(-(deltaS/kappa)**2.)/step[1]\n",
        "            gE = np.exp(-(deltaE/kappa)**2.)/step[2]\n",
        "        elif option == 2:\n",
        "            gD = 1./(1.+(deltaD/kappa)**2.)/step[0]\n",
        "            gS = 1./(1.+(deltaS/kappa)**2.)/step[1]\n",
        "            gE = 1./(1.+(deltaE/kappa)**2.)/step[2]\n",
        "\n",
        "        # update matrices\n",
        "        D = gD*deltaD\n",
        "        E = gE*deltaE\n",
        "        S = gS*deltaS\n",
        "\n",
        "        # subtract a copy that has been shifted 'Up/North/West' by one\n",
        "        # pixel. don't as questions. just do it. trust me.\n",
        "        UD[:] = D\n",
        "        NS[:] = S\n",
        "        EW[:] = E\n",
        "        UD[1:,: ,: ] -= D[:-1,:  ,:  ]\n",
        "        NS[: ,1:,: ] -= S[:  ,:-1,:  ]\n",
        "        EW[: ,: ,1:] -= E[:  ,:  ,:-1]\n",
        "\n",
        "        # update the image\n",
        "        stackout += gamma*(UD+NS+EW)\n",
        "\n",
        "        if ploton:\n",
        "            iterstring = \"Iteration %i\" %(ii+1)\n",
        "            ih.set_data(stackout[showplane,...].squeeze())\n",
        "            ax2.set_title(iterstring)\n",
        "            fig.canvas.draw()\n",
        "            # sleep(0.01)\n",
        "\n",
        "    return stackout\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CjVgNZuKJhtH",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "0cc50496-6a19-4c93-a8c6-e1739952bea8"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'def denoise_volumes(in_volumes) :\\n\\n  n_loop = in_volumes.shape[0]\\n\\n  out_volumes = np.zeros(in_volumes.shape)\\n  #print(out_volumes.shape)\\n\\n  for i in range(0,n_loop,1):\\n    temp = in_volumes[i,:,:,:,:]\\n    temp = anisodiff3(temp,niter=10)\\n    temp = temp.reshape((*temp.shape, 1))\\n    out_volumes[i] = temp\\n\\n    #print(temp.shape)\\n\\n  return out_volumes'"
            ]
          },
          "metadata": {},
          "execution_count": 275
        }
      ],
      "source": [
        "\"\"\"def denoise_volumes(in_volumes) :\n",
        "\n",
        "  n_loop = in_volumes.shape[0]\n",
        "\n",
        "  out_volumes = np.zeros(in_volumes.shape)\n",
        "  #print(out_volumes.shape)\n",
        "\n",
        "  for i in range(0,n_loop,1):\n",
        "    temp = in_volumes[i,:,:,:,:]\n",
        "    temp = anisodiff3(temp,niter=10)\n",
        "    temp = temp.reshape((*temp.shape, 1))\n",
        "    out_volumes[i] = temp\n",
        "\n",
        "    #print(temp.shape)\n",
        "\n",
        "  return out_volumes\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sfo3B9gWKRrY",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "ca87a2f6-96e2-44c7-df4a-6079908c5706"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'t_volumes = denoise_volumes(t_volumes)\\nv_volumes = denoise_volumes(v_volumes)\\n\\nt_volumes = t_volumes_clean\\nv_volumes = v_volumes_clean\\n\\nprint(t_volumes.shape)\\nprint(v_volumes.shape)'"
            ]
          },
          "metadata": {},
          "execution_count": 276
        }
      ],
      "source": [
        "\"\"\"t_volumes = denoise_volumes(t_volumes)\n",
        "v_volumes = denoise_volumes(v_volumes)\n",
        "\n",
        "t_volumes = t_volumes_clean\n",
        "v_volumes = v_volumes_clean\n",
        "\n",
        "print(t_volumes.shape)\n",
        "print(v_volumes.shape)\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9Omu--2PGX6P"
      },
      "outputs": [],
      "source": [
        "# check_vol_clean = anisodiff3(check_vol)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "r-9Q9zKvGsdB",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "0cb2a0b2-95e6-480d-dae2-d205947f0e27"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "\"rotated_vol_clean = ndimage.rotate(check_vol_clean, 90)\\nplt.axis('off')\\nplt.imshow(rotated_vol_clean[:, :, 150], cmap='gray')\\nplt.show()\""
            ]
          },
          "metadata": {},
          "execution_count": 278
        }
      ],
      "source": [
        "\"\"\"rotated_vol_clean = ndimage.rotate(check_vol_clean, 90)\n",
        "plt.axis('off')\n",
        "plt.imshow(rotated_vol_clean[:, :, 150], cmap='gray')\n",
        "plt.show()\"\"\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OAok6QlCeBm0"
      },
      "source": [
        "Splitting the Dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "x_aGLgc60ceM",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "29bc6d18-5a03-4caa-e90c-d467f55a416e"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "(9, 256, 128, 256, 1)\n",
            "(1, 256, 128, 256, 1)\n"
          ]
        }
      ],
      "source": [
        "# Split the training data into training and validation\n",
        "training_volumes = t_volumes[TRAINING_VOLUMES]\n",
        "training_labels = t_labels[TRAINING_VOLUMES]\n",
        "\n",
        "validation_volumes = t_volumes[VALIDATION_VOLUMES]\n",
        "validation_labels = t_labels[VALIDATION_VOLUMES]\n",
        "\n",
        "print(training_volumes.shape)\n",
        "#print(training_labels.shape)\n",
        "\n",
        "print(validation_volumes.shape)\n",
        "#print(validation_labels.shape)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LpKWGJ-Fh1B_"
      },
      "source": [
        "##Extracting Patches"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8GPSNW8HhzNb"
      },
      "outputs": [],
      "source": [
        "# def z_score_standardisation(x, avg, std):\n",
        "#   return (x-avg)/std"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BFDtzyW6FH7x"
      },
      "source": [
        "**Extract *useful* patches**\n",
        "\n",
        "This step is fundamental, we want to provide the network with useful information"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vMvPbAZ4gb8C"
      },
      "outputs": [],
      "source": [
        "def extract_patches(x, patch_size, patch_stride) :\n",
        "  return tf.image.extract_patches(\n",
        "    x,\n",
        "    sizes=[1, *patch_size, 1],\n",
        "    strides=[1, *patch_stride, 1],\n",
        "    rates=[1, 1, 1, 1],\n",
        "    padding='SAME', name=None)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nsKNTm-Lf-sb"
      },
      "outputs": [],
      "source": [
        "def extract_useful_patches(\n",
        "    volumes, labels,\n",
        "    image_size=IMAGE_SIZE,\n",
        "    patch_size=PATCH_SIZE,\n",
        "    stride=PATCH_STRIDE,\n",
        "    threshold=CONTENT_THRESHOLD,\n",
        "    num_classes=N_CLASSES) :\n",
        "\n",
        "  volumes = volumes.reshape([-1, image_size[1], image_size[2], 1])\n",
        "  labels = labels.reshape([-1, image_size[1], image_size[2], 1])\n",
        "\n",
        "  vol_patches = extract_patches(volumes, patch_size, stride).numpy()\n",
        "  seg_patches = extract_patches(labels, patch_size, stride).numpy()\n",
        "\n",
        "  vol_patches = vol_patches.reshape([-1, *patch_size, 1])\n",
        "  seg_patches = seg_patches.reshape([-1, *patch_size, ])\n",
        "\n",
        "  # this will get rid of the background and only take foreground\n",
        "  foreground_mask = seg_patches != 0 \n",
        "\n",
        "  # we only keep the useful forground patches\n",
        "  # threshold too small - takes even the useless patches\n",
        "  # threshold too high - might leave out useful patches\n",
        "  useful_patches = foreground_mask.sum(axis=(1, 2)) > threshold * np.prod(patch_size)\n",
        "\n",
        "  vol_patches = vol_patches[useful_patches]\n",
        "  seg_patches = seg_patches[useful_patches]\n",
        "\n",
        "  seg_patches = tf.keras.utils.to_categorical(\n",
        "    seg_patches, num_classes=N_CLASSES, dtype='float32')\n",
        "  \n",
        "  return (vol_patches, seg_patches)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tTcoEi426bfg",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "3c7b24b8-963a-4fc5-af1b-6ea04d68cabd"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "(11546, 32, 32, 1)\n"
          ]
        }
      ],
      "source": [
        "# extract patches from training set\n",
        "(training_patches, training_patches_seg) = extract_useful_patches(training_volumes, training_labels)\n",
        "\n",
        "# extract patches from validation set\n",
        "(validation_patches, validation_patches_seg) = extract_useful_patches(validation_volumes, validation_labels)\n",
        "\n",
        "print(training_patches.shape)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mx9fsloJgxrm"
      },
      "source": [
        "##Data Augmentation"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Degree of Augmentation\n",
        "deg     = 0.2\n",
        "\n",
        "datagen = ImageDataGenerator(\n",
        "        rotation_range=40, #40\n",
        "        width_shift_range=deg,\n",
        "        height_shift_range=deg,\n",
        "        # rescale=1./255,\n",
        "        shear_range=deg,\n",
        "        zoom_range=deg,\n",
        "        horizontal_flip=True,\n",
        "        vertical_flip=True,\n",
        "        fill_mode='nearest') #reflect, wrap, constant(black)"
      ],
      "metadata": {
        "id": "CwyB93VgyuIn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6o21TRkXedUO"
      },
      "outputs": [],
      "source": [
        "train_generator = datagen.flow(training_patches, batch_size=int(training_patches.shape[0]/BATCH_SIZE), seed=1)\n",
        "train_label_generator = datagen.flow(training_patches_seg, batch_size=int(training_patches.shape[0]/BATCH_SIZE), seed=1)\n",
        "\n",
        "val_generator = datagen.flow(validation_patches, batch_size=int(validation_patches.shape[0]/BATCH_SIZE), seed=1)\n",
        "val_label_generator = datagen.flow(validation_patches_seg, batch_size=int(validation_patches.shape[0]/BATCH_SIZE), seed=1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sOw8_p4Oedag"
      },
      "outputs": [],
      "source": [
        "X_train = train_generator.next()\n",
        "y_train = train_label_generator.next()\n",
        "\n",
        "X_val = val_generator.next()\n",
        "y_val = val_label_generator.next()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IELRGzM3r2mz",
        "outputId": "9350d930-b553-40c4-9501-854a4ddeef33"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "(11546, 32, 32, 1)\n",
            "(11546, 32, 32, 4)\n",
            "----------------\n",
            "(1201, 32, 32, 1)\n",
            "(1201, 32, 32, 4)\n"
          ]
        }
      ],
      "source": [
        "print(training_patches.shape)\n",
        "print(training_patches_seg.shape)\n",
        "print(\"----------------\")\n",
        "print(validation_patches.shape)\n",
        "print(validation_patches_seg.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ce8mcx6Nr3QG",
        "outputId": "0b4adb3a-b7f8-4e0e-9c67-e47f7b34ad6d"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "(11726, 32, 32, 1)\n",
            "(11726, 32, 32, 4)\n",
            "(1219, 32, 32, 1)\n",
            "(1219, 32, 32, 4)\n"
          ]
        }
      ],
      "source": [
        "full_train = np.concatenate((training_patches, X_train))\n",
        "print(full_train.shape)\n",
        "full_train_label = np.concatenate((training_patches_seg, y_train))\n",
        "print(full_train_label.shape)\n",
        "\n",
        "full_val = np.concatenate((validation_patches, X_val))\n",
        "print(full_val.shape)\n",
        "full_val_label = np.concatenate((validation_patches_seg, y_val))\n",
        "print(full_val_label.shape)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3fdqvKEK1UES"
      },
      "source": [
        "##Train the Model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bJYxf3HiFopr"
      },
      "source": [
        "\n",
        "\n",
        "----\n",
        "\n",
        "---\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "**Instantiate UNet model and train it**\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EkBp70aWGB_b"
      },
      "source": [
        "Using callbacks to stop training and avoid overfitting\n",
        "\n",
        "\n",
        "*   Early stopping with a certain patience\n",
        "*   Save (and load!) best model\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "unet = get_unet()\n",
        "# unet.summary()"
      ],
      "metadata": {
        "id": "nl0moG5se2IT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "j7H9hIPwEwJZ",
        "outputId": "5a6caebe-acf8-4ffb-9170-c0cd721b9756"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 1/200\n",
            "184/184 [==============================] - 12s 57ms/step - loss: 0.8213 - val_loss: 0.6550\n",
            "Epoch 2/200\n",
            "184/184 [==============================] - 10s 54ms/step - loss: 0.5001 - val_loss: 0.5260\n",
            "Epoch 3/200\n",
            "184/184 [==============================] - 10s 54ms/step - loss: 0.4190 - val_loss: 0.4734\n",
            "Epoch 4/200\n",
            "184/184 [==============================] - 10s 55ms/step - loss: 0.3658 - val_loss: 0.4588\n",
            "Epoch 5/200\n",
            "184/184 [==============================] - 10s 55ms/step - loss: 0.3326 - val_loss: 0.4419\n",
            "Epoch 6/200\n",
            "184/184 [==============================] - 10s 53ms/step - loss: 0.3050 - val_loss: 0.4513\n",
            "Epoch 7/200\n",
            "184/184 [==============================] - 10s 55ms/step - loss: 0.2774 - val_loss: 0.4387\n",
            "Epoch 8/200\n",
            "184/184 [==============================] - 10s 56ms/step - loss: 0.2646 - val_loss: 0.3885\n",
            "Epoch 9/200\n",
            "184/184 [==============================] - 10s 54ms/step - loss: 0.2475 - val_loss: 0.4258\n",
            "Epoch 10/200\n",
            "184/184 [==============================] - 10s 54ms/step - loss: 0.2343 - val_loss: 0.4442\n",
            "Epoch 11/200\n",
            "184/184 [==============================] - 10s 55ms/step - loss: 0.2217 - val_loss: 0.3760\n",
            "Epoch 12/200\n",
            "184/184 [==============================] - 10s 53ms/step - loss: 0.2148 - val_loss: 0.4416\n",
            "Epoch 13/200\n",
            "184/184 [==============================] - 10s 54ms/step - loss: 0.2059 - val_loss: 0.5007\n",
            "Epoch 14/200\n",
            "184/184 [==============================] - 10s 53ms/step - loss: 0.2014 - val_loss: 0.4063\n",
            "Epoch 15/200\n",
            "184/184 [==============================] - 10s 53ms/step - loss: 0.1942 - val_loss: 0.4094\n",
            "Epoch 16/200\n",
            "184/184 [==============================] - 10s 54ms/step - loss: 0.1890 - val_loss: 0.4257\n",
            "Epoch 17/200\n",
            "184/184 [==============================] - 10s 53ms/step - loss: 0.1848 - val_loss: 0.4012\n",
            "Epoch 18/200\n",
            "184/184 [==============================] - 10s 54ms/step - loss: 0.1821 - val_loss: 0.4097\n",
            "Epoch 19/200\n",
            "184/184 [==============================] - 10s 54ms/step - loss: 0.1788 - val_loss: 0.3952\n",
            "Epoch 20/200\n",
            "184/184 [==============================] - 10s 55ms/step - loss: 0.1779 - val_loss: 0.4697\n",
            "Epoch 21/200\n",
            "184/184 [==============================] - 10s 54ms/step - loss: 0.1718 - val_loss: 0.4042\n",
            "Epoch 22/200\n",
            "184/184 [==============================] - 10s 54ms/step - loss: 0.1680 - val_loss: 0.4180\n",
            "Epoch 23/200\n",
            "184/184 [==============================] - 10s 55ms/step - loss: 0.1700 - val_loss: 0.4178\n",
            "Epoch 24/200\n",
            "184/184 [==============================] - 10s 54ms/step - loss: 0.1657 - val_loss: 0.3933\n",
            "Epoch 25/200\n",
            "184/184 [==============================] - 10s 54ms/step - loss: 0.1621 - val_loss: 0.4273\n",
            "Epoch 26/200\n",
            "184/184 [==============================] - 10s 54ms/step - loss: 0.1620 - val_loss: 0.5142\n",
            "Epoch 27/200\n",
            "184/184 [==============================] - 10s 54ms/step - loss: 0.1582 - val_loss: 0.4007\n",
            "Epoch 28/200\n",
            "184/184 [==============================] - 10s 54ms/step - loss: 0.1548 - val_loss: 0.4588\n",
            "Epoch 29/200\n",
            "184/184 [==============================] - 10s 54ms/step - loss: 0.1560 - val_loss: 0.4444\n",
            "Epoch 30/200\n",
            "184/184 [==============================] - 10s 54ms/step - loss: 0.1525 - val_loss: 0.4136\n",
            "Epoch 31/200\n",
            "184/184 [==============================] - 10s 55ms/step - loss: 0.1527 - val_loss: 0.4143\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<keras.callbacks.History at 0x7f8f2d36d450>"
            ]
          },
          "metadata": {},
          "execution_count": 290
        }
      ],
      "source": [
        "my_callbacks = [\n",
        "    tf.keras.callbacks.EarlyStopping(patience=PATIENCE), # early stopping\n",
        "    tf.keras.callbacks.ModelCheckpoint(filepath=MODEL_FNAME_PATTERN, save_best_only=True) # save the best based on validation\n",
        "]\n",
        "\n",
        "unet = get_unet()\n",
        "unet.compile(optimizer=OPTIMISER, loss=LOSS)\n",
        "unet.fit(\n",
        "    x=full_train, \n",
        "    y=full_train_label,\n",
        "    validation_data=(full_val, full_val_label),\n",
        "    batch_size=BATCH_SIZE,\n",
        "    epochs=N_EPOCHS,\n",
        "    callbacks=my_callbacks,\n",
        "    verbose=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XGcfrZCMjz4M"
      },
      "source": [
        "##Load the best model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5fNsYJL7KnON"
      },
      "outputs": [],
      "source": [
        "unet = get_unet(\n",
        "    img_size=(IMAGE_SIZE[1], IMAGE_SIZE[2]),\n",
        "    n_classes=N_CLASSES,\n",
        "    n_input_channels=N_INPUT_CHANNELS)\n",
        "unet.compile(optimizer=OPTIMISER, loss=LOSS)\n",
        "unet.load_weights('model.h5')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CwPzN0jrSm8r"
      },
      "source": [
        "##Prepare test data using the validation volumes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vuRHLNYGkDFp"
      },
      "outputs": [],
      "source": [
        "def prepare_val_data(the_volumes, the_labels):\n",
        "  testing_volumes_processed = the_volumes.reshape([-1, IMAGE_SIZE[1], IMAGE_SIZE[2], 1])\n",
        "  testing_labels_processed = the_labels.reshape([-1, IMAGE_SIZE[1], IMAGE_SIZE[2], 1])\n",
        "\n",
        "  testing_labels_processed = tf.keras.utils.to_categorical(testing_labels_processed, num_classes=4, dtype='float32')\n",
        "\n",
        "  #print(testing_volumes_processed.shape)\n",
        "  #print(testing_labels_processed.shape)\n",
        "\n",
        "  return (testing_volumes_processed, testing_labels_processed)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xbJZUQyLqF4p"
      },
      "source": [
        "###Predict labels for test data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ItoY31x0K3r8"
      },
      "outputs": [],
      "source": [
        "def pred_val_data(testing_volumes_processed)  :\n",
        "  # creates probability map of each label for all volumes\n",
        "  prediction = unet.predict(x=testing_volumes_processed)\n",
        "\n",
        "  prediction = np.argmax(prediction, axis=3)\n",
        "\n",
        "  #plt.axis('off')\n",
        "  #plt.imshow(prediction[:, :, 150])\n",
        "\n",
        "  return prediction"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VU2XAj_wxgJv",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "outputId": "b60eb3cb-b5c6-427e-c993-d65841e6e865"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'\\nprint(prediction.shape)\\nprint(testing_labels_processed.shape)\\nprint(testing_volumes_T1_processed.shape)\\n'"
            ]
          },
          "metadata": {},
          "execution_count": 294
        }
      ],
      "source": [
        "\"\"\"\n",
        "print(prediction.shape)\n",
        "print(testing_labels_processed.shape)\n",
        "print(testing_volumes_T1_processed.shape)\n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "##Computing Dice, AVD and HD (Final)\n",
        "\n"
      ],
      "metadata": {
        "id": "5RHaF_LhpWXr"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def compute_hausdorff_distance(in1, in2, label = 'all'):\n",
        "    in1=in1.squeeze()\n",
        "    in2=in2.squeeze()\n",
        "    hausdorff_distance_filter = sitk.HausdorffDistanceImageFilter()\n",
        "    if label == 'all':\n",
        "        # Hausdorff distance\n",
        "        hausdorff_distance_filter.Execute(in1, in2)\n",
        "    else:\n",
        "        in1_array  = in1 #sitk.GetArrayFromImage(in1)\n",
        "        in1_array = (in1_array == label) *1 \n",
        "        in1_array = in1_array.astype('uint16')  \n",
        "        img1 = sitk.GetImageFromArray(in1_array)\n",
        "        \n",
        "        in2_array  = in2 #sitk.GetArrayFromImage(in2)\n",
        "        in2_array = (in2_array == label) *1 \n",
        "        in2_array = in2_array.astype('uint16')  \n",
        "        img2 = sitk.GetImageFromArray(in2_array)\n",
        "        # Hausdorff distance\n",
        "        hausdorff_distance_filter.Execute(img1, img2)\n",
        "    return hausdorff_distance_filter.GetHausdorffDistance()\n",
        "\n",
        "def compute_dice_coefficient(in1, in2, label  = 'all'):\n",
        "    in1=in1.squeeze()\n",
        "    in2=in2.squeeze()\n",
        "    if label=='all': \n",
        "        return 2 * np.sum( (in1>0) &  (in2>0) & (in1 == in2)) / (np.sum(in1 > 0) + np.sum(in2 > 0))\n",
        "    else:\n",
        "        return 2 * np.sum((in1 == label) & (in2 == label)) / (np.sum(in1 == label) + np.sum(in2 == label))\n",
        "\n",
        "def compute_volumentric_difference(in1, in2, label  = 'all'):\n",
        "    in1=in1.squeeze()\n",
        "    in2=in2.squeeze()\n",
        "    if label  == 'all':\n",
        "      #  vol_dif  = np.sum((in1 != in2) & (in1 !=0) & (in2 !=0))\n",
        "        return np.sum((in1 != in2)) / ((np.sum(in1 > 0) + np.sum(in2 > 0)))\n",
        "    else:\n",
        "        in1  = (in1 == label) * 1\n",
        "        in2  = (in2 == label) * 1\n",
        "        return np.sum((in1 != in2)) / ((np.sum(in1 > 0) + np.sum(in2 > 0)))"
      ],
      "metadata": {
        "id": "iFBnx17Bpa_b"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "for cl in range(0,4,1): \n",
        "  overallDSC = np.zeros(N_CLASSES)\n",
        "  overall_Hausdorff = np.zeros(N_CLASSES)\n",
        "  overall_vol = np.zeros(N_CLASSES)\n",
        "\n",
        "  for i in range(0,validation_volumes.shape[0], 1):\n",
        "      \n",
        "      testing_volumes_processed, testing_labels_processed = prepare_val_data(v_volumes[i], v_labels[i])\n",
        "      prediction = pred_val_data(testing_volumes_processed)\n",
        "      \n",
        "      #cl = 3\n",
        "\n",
        "      cur_DSC = compute_dice_coefficient(prediction, v_labels[i], label=cl)\n",
        "      overallDSC = overallDSC + cur_DSC\n",
        "\n",
        "      cur_Hausdorff = compute_hausdorff_distance(prediction, v_labels[i], label=cl) \n",
        "      overall_Hausdorff = overall_Hausdorff + cur_Hausdorff\n",
        "\n",
        "      cur_vol = compute_volumentric_difference(prediction, v_labels[i], label=cl)\n",
        "      overall_vol = overall_vol + cur_vol\n",
        "      \n",
        "      #print(prediction.shape)\n",
        "      #print(v_labels[i].shape)\n",
        "      \n",
        "  #print(overall_Hausdorff)\n",
        "  overallDSC = overallDSC/validation_volumes.shape[0]\n",
        "  overall_Hausdorff = overall_Hausdorff/validation_volumes.shape[0]\n",
        "  overall_vol = overall_vol/validation_volumes.shape[0]\n",
        "\n",
        "  # for i in range(0,cl,1):\n",
        "  #print(\"Class {} - Dice Coefficient = {:.4f}\".format(cl, overallDSC[i]))\n",
        "  #print(\"Class {} - HD = {:.4f}\".format(cl, overall_Hausdorff[i]))\n",
        "  #print(\"Class {} - AVD = {:.4f}\".format(cl, overall_vol[i]))\n",
        "  print(\"Class {}\".format(cl))\n",
        "  print(\"\\tDice Coefficient = {:.4f}\".format(overallDSC[i]))\n",
        "  print(\"\\tHD = {:.4f}\".format(overall_Hausdorff[i]))\n",
        "  print(\"\\tAVD = {:.4f}\".format(overall_vol[i]))"
      ],
      "metadata": {
        "id": "8PrStjpepejb",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "0bf23c95-0865-4261-c045-704e879d3994"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Class 0\n",
            "\tDice Coefficient = 0.9972\n",
            "\tHD = 20.3224\n",
            "\tAVD = 0.0028\n",
            "Class 1\n",
            "\tDice Coefficient = 0.7940\n",
            "\tHD = 145.8869\n",
            "\tAVD = 0.2060\n",
            "Class 2\n",
            "\tDice Coefficient = 0.9106\n",
            "\tHD = 112.4189\n",
            "\tAVD = 0.0894\n",
            "Class 3\n",
            "\tDice Coefficient = 0.8646\n",
            "\tHD = 113.6574\n",
            "\tAVD = 0.1354\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j3zTHEeiN4kE",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "outputId": "578911dc-e02c-4edf-dab7-f3269045e8d3"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'\\nClass 0 - Dice Coefficient 0.9975\\nClass 1 - Dice Coefficient 0.8342\\nClass 2 - Dice Coefficient 0.9209\\nClass 3 - Dice Coefficient 0.8825\\n'"
            ]
          },
          "metadata": {},
          "execution_count": 297
        }
      ],
      "source": [
        "# batch size = 32, patient = 5, dropout=0.15, epoch = 50\n",
        "\"\"\"\n",
        "Class 0 - Dice Coefficient 0.9976\n",
        "Class 1 - Dice Coefficient 0.8288\n",
        "Class 2 - Dice Coefficient 0.9186\n",
        "Class 3 - Dice Coefficient 0.8765\n",
        "\"\"\"\n",
        "\n",
        "# batch size = 40, patient = 5, dropout=0.15, epoch = 50\n",
        "\"\"\"\n",
        "Class 0 - Dice Coefficient 0.9977\n",
        "Class 1 - Dice Coefficient 0.8261\n",
        "Class 2 - Dice Coefficient 0.9202\n",
        "Class 3 - Dice Coefficient 0.8790\n",
        "\"\"\"\n",
        "\n",
        "# batch size = 50, patient = 20, dropout=0.15, epoch = 200\n",
        "\"\"\"\n",
        "Class 0 - Dice Coefficient 0.9977\n",
        "Class 1 - Dice Coefficient 0.8261\n",
        "Class 2 - Dice Coefficient 0.9202\n",
        "Class 3 - Dice Coefficient 0.8790\n",
        "\"\"\"\n",
        "\n",
        "# batch size = 64, patient = 20, dropout=0.40, epoch = 200\n",
        "\"\"\"\n",
        "Class 0 - Dice Coefficient 0.9975\n",
        "Class 1 - Dice Coefficient 0.8342\n",
        "Class 2 - Dice Coefficient 0.9209\n",
        "Class 3 - Dice Coefficient 0.8825\n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "source": [
        ""
      ],
      "metadata": {
        "id": "CipR5ZtUqSDC"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "ByBoNkPGWKrN",
        "uZ2BevUuIkGs"
      ],
      "name": "2DUNet.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}