Transformers_Unsupervised_Anomaly_Segmentation.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Transformers_Unsupervised_Anomaly_Segmentation.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "He5bnMxF4eUO",
        "pUt7KJYD4jH_",
        "BODSO6Mr4ymn",
        "ONbEgrjv55Xq",
        "A_KV8CdA6MU7",
        "UO40zWJ96_SY",
        "JBL3NyI957XU",
        "Guj7hFmq7Pjy",
        "9ytUmxAK7Rp4",
        "NeeKAOcK6Ahf",
        "nuEgxobC7T_N",
        "ob9yVvQg7V__",
        "PrHDynS67XoN",
        "QucRsBiN7b5F",
        "rx-OXyQG7eDv"
      ]
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "He5bnMxF4eUO"
      },
      "source": [
        "# Drive mount"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "J0TlXj0x4WlK"
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive', force_remount=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pUt7KJYD4jH_"
      },
      "source": [
        "# Installs "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zqXin51A4qN_"
      },
      "source": [
        "!pip install -r requirements.txt"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BODSO6Mr4ymn"
      },
      "source": [
        "# Imports"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wiBWiMst40Zz"
      },
      "source": [
        "import json\n",
        "import os\n",
        "\n",
        "import cv2\n",
        "from skimage import filters\n",
        "from einops import rearrange\n",
        "import statistics\n",
        "import seaborn as sns; sns.set_theme()\n",
        "\n",
        "import random\n",
        "import traceback\n",
        "import nibabel as nib\n",
        "import scipy \n",
        "\n",
        "import numpy as np\n",
        "from numpy import save\n",
        "import matplotlib.pyplot as plt\n",
        "import time\n",
        "from datetime import datetime\n",
        "\n",
        "import tensorflow as tf\n",
        "import tensorflow_addons as tfa\n",
        "\n",
        "from tensorflow.keras import layers\n",
        "from plot_keras_history import plot_history\n",
        "\n",
        "from sklearn.model_selection import ParameterGrid\n",
        "from sklearn import metrics\n",
        "\n",
        "from scripts.evalresults import *\n",
        "from scripts.utils import *"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ONbEgrjv55Xq"
      },
      "source": [
        "# Autoencoders"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A_KV8CdA6MU7"
      },
      "source": [
        "## Dense Convolutional Autoencoder"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8DICbwKf54j3"
      },
      "source": [
        "#-- Model implementation : Dense Convolutional Autoencoder\n",
        "\n",
        "def DCAE():\n",
        "  \n",
        "    input_img = tf.keras.Input(shape=(image_size, image_size, num_channels)) \n",
        "\n",
        "    x = layers.Conv2D(32 , 5, activation=layers.LeakyReLU(), strides=2, padding=\"same\")(input_img)\n",
        "    x = layers.Conv2D(64 , 5, activation=layers.LeakyReLU(), strides=2, padding=\"same\")(x)\n",
        "    x = layers.Conv2D(128, 5, activation=layers.LeakyReLU(), strides=2, padding=\"same\")(x)    \n",
        "    x = layers.Conv2D(128, 5, activation=layers.LeakyReLU(), strides=2, padding=\"same\")(x)   \n",
        "    x = layers.Conv2D(16 , 1, activation=layers.LeakyReLU(), strides=1, padding=\"same\")(x)\n",
        "    \n",
        "    x = layers.Flatten()(x)\n",
        "    encoded = layers.Dense(intermediate_dim, activation=layers.LeakyReLU())(x)\n",
        "\n",
        "    #-- BOTTELNECK SIZE : 512 \n",
        "    \n",
        "    x = layers.Dense(16 * 16 * 16, activation=layers.LeakyReLU())(encoded)\n",
        "    x = layers.Reshape((16, 16, 16))(x)\n",
        "\n",
        "    x = layers.Conv2D(128, 1, strides=1, activation=layers.LeakyReLU(), padding=\"same\")(x)    \n",
        "    x = layers.Conv2DTranspose(128, 5, strides=2, activation=layers.LeakyReLU(), padding=\"same\")(x) \n",
        "    x = layers.Conv2DTranspose(64 , 5, strides=2, activation=layers.LeakyReLU(), padding=\"same\")(x)\n",
        "    x = layers.Conv2DTranspose(32 , 5, strides=2, activation=layers.LeakyReLU(), padding=\"same\")(x)\n",
        "    x = layers.Conv2DTranspose(32 , 5, strides=2, activation=layers.LeakyReLU(), padding=\"same\")(x)\n",
        "    \n",
        "    decoded = layers.Conv2D(num_channels, 1, activation=layers.LeakyReLU(), padding='same')(x)\n",
        "    \n",
        "    model = tf.keras.Model(input_img, decoded)\n",
        "\n",
        "    return model\n",
        "\n",
        "    \n",
        "#-- Configure the hyperparameters\n",
        "\n",
        "model_name = 'Dense Convolutional Autoencoder'\n",
        "numEpochs = 50\n",
        "learning_rate = 0.00001\n",
        "image_size = 256\n",
        "num_channels = 1\n",
        "batch_size = 1\n",
        "intermediate_dim = 512\n",
        "\n",
        "\n",
        "#-- Configure training and testing Datasets \n",
        "\n",
        "saved_dir = './saved/'\n",
        "data_dir  = './data/OASIS/'\n",
        "\n",
        "test_healthy_path = './data/OASIS_all/OASIS_Flair_Test.npy'\n",
        "\n",
        "test_path = './data/BraTS/s0/BraTS_Flair.npy'\n",
        "brainmask_path = './data/BraTS/s0/BraTS_Brainmask.npy'\n",
        "x_prior_path = './data/BraTS/s0/BraTS_prior_52.npy.npy'\n",
        "label_path = './data/BraTS/s0/BraTS_GT.npy'\n",
        "\n",
        "'''\n",
        "#-- If using MSLUB as test-set\n",
        "test_path = './data/MSLUB/MSLUB_Flair.npy'\n",
        "brainmask_path = './data/MSLUB/MSLUB_Brainmask.npy'\n",
        "x_prior_path = './data/MSLUB/MSLUB_prior_57.npy'\n",
        "label_path = './data/MSLUB/MSLUB_GT.npy'\n",
        "'''\n",
        "\n",
        "list_len = [4805, 3078]    #-- BraTS and MSLUB test-set sizes, respectively.\n",
        "len_testset = list_len[0]  #-- 0 for BraTS and 1 for MSLUB\n",
        "\n",
        "train_paths = list_of_paths(data_dir)\n",
        "\n",
        "nb_train_files = 66\n",
        "data_gen = data_generator(train_paths[:nb_train_files], batch_size)\n",
        "training_steps = (256 / batch_size) * nb_train_files\n",
        "\n",
        "nb_val_files = 5\n",
        "val_gen = data_generator(train_paths[-nb_val_files:], batch_size)\n",
        "validation_steps = (256 / batch_size) * nb_val_files\n",
        "\n",
        "\n",
        "#-- Checkpoints dir\n",
        "\n",
        "date = datetime. now(). strftime(\"%Y_%m_%d-%I:%M:%S_%p\")\n",
        "ckpts_dir = os.path.join(saved_dir, f'Ckpts_{date}')\n",
        "os.makedirs(ckpts_dir)\n",
        "      \n",
        "ckpts_path = os.path.join(ckpts_dir, 'Model_Ckpts.h5')\n",
        "params_path = os.path.join(ckpts_dir, 'Parameters.txt')\n",
        "results_path = os.path.join(ckpts_dir, 'Results.txt')\n",
        "fig_path = os.path.join(ckpts_dir, 'History_plot.png')\n",
        "dice_plot_path = os.path.join(ckpts_dir, 'Dice_plot.png')\n",
        "predicted_path = os.path.join(ckpts_dir, 'Predicted.npy')\n",
        "residual_path = os.path.join(ckpts_dir, 'Residuals.npy')\n",
        "residual_BP_path = os.path.join(ckpts_dir, 'Residuals_BP.npy')\n",
        " \n",
        "\n",
        "#-- Configure the training\n",
        "\n",
        "opt = tf.keras.optimizers.Adam(learning_rate = learning_rate)\n",
        "          \n",
        "calbks = tf.keras.callbacks.ModelCheckpoint(filepath=ckpts_path, monitor='val_loss', save_best_only=True, save_weights_only=True, verbose=2)\n",
        "tqdm_callback = tfa.callbacks.TQDMProgressBar() \n",
        "\n",
        "model = DCAE()\n",
        "model.summary()\n",
        "model.compile(optimizer=opt, loss='mae', metrics=['mse', SSIMLoss, MS_SSIMLoss])\n",
        "\n",
        "\n",
        "#-- Print & Write model Parameters\n",
        "\n",
        "parameters = (f'\\nSelected model \"{model_name}\" with :\\n - {batch_size}: Batche(s)\\n - {numEpochs}: Epochs\\n - {intermediate_dim}: Bottelneck size\\n')\n",
        "print(parameters)\n",
        "\n",
        "           \n",
        "#-- TRAIN \n",
        "\n",
        "print('\\nTrain =>\\n')\n",
        "history = model.fit(x = data_gen,\n",
        "                    steps_per_epoch = training_steps,\n",
        "                    validation_data = val_gen,\n",
        "                    validation_steps = validation_steps,\n",
        "                    verbose = 0,\n",
        "                    epochs = numEpochs,\n",
        "                    callbacks = [calbks, tqdm_callback]\n",
        "                    )\n",
        "\n",
        "                          \n",
        "#-- Get training and test loss histories \n",
        "\n",
        "plot_history(history, path=fig_path)\n",
        "plt.close()\n",
        "time.sleep(2)\n",
        "\n",
        "\n",
        "#-- Test  \n",
        "\n",
        "print('\\nTest ===>\\n')\n",
        "my_test = np.load(test_path)\n",
        "brainmask = np.load(brainmask_path)\n",
        "x_prior = np.load(x_prior_path)\n",
        "my_labels = np.load(label_path)\n",
        "\n",
        "healthy_test = np.load(test_healthy_path)\n",
        "steps = healthy_test.shape[0]\n",
        "\n",
        "score = model.evaluate(x=healthy_test, y=healthy_test, verbose=0, steps=steps, callbacks = [tqdm_callback])    \n",
        "score_out = (f'\\nTest on healthy unseen data :\\n - Test loss (MAE): {score[0]},\\n - Test MSE : {score[1]},\\n - Test SSIM : {score[2]}\\n - Test MS_SSIM : {score[3]}\\n')\n",
        "print(score_out)\n",
        "\n",
        "\n",
        "#-- Predict\n",
        "\n",
        "print('\\nPredict =====>\\n')\n",
        "predicted = model.predict(x=my_test, verbose=1, steps=len_testset)\n",
        "np.save(predicted_path, predicted)\n",
        "time.sleep(4)\n",
        "\n",
        "\n",
        "#-- Calculate, Post-process and Save Residuals\n",
        "\n",
        "print('\\nCalculate, Post-process and Save Residuals =====>\\n')     \n",
        "residual_BP = calculate_residual_BP(my_test, predicted, brainmask)  #-- You can use either brainmask or x_prior\n",
        "np.save(residual_BP_path, residual_BP)\n",
        "        \n",
        "residual = calculate_residual(my_test, predicted, brainmask)  #-- You can use either brainmask or x_prior\n",
        "np.save(residual_path, residual)\n",
        "        \n",
        "\n",
        "#-- Evaluation\n",
        "\n",
        "print('\\nEvaluate =========>\\n')        \n",
        "[AUROC, AUPRC, AVG_DICE, MAD, STD], DICE = eval_residuals(my_labels, residual)     \n",
        "results = (f'\\nResults after median_filter :\\n - AUROC = {AUROC}\\n - AUPRC = {AUPRC}\\n - AVG_DICE = {AVG_DICE}\\n - MEDIAN_ABSOLUTE_DEVIATION = {MAD}\\n - STANDARD_DEVIATION = {STD}')\n",
        "print(results)\n",
        "                      \n",
        "plt.figure()\n",
        "hor_axis = [x for x in range(len_testset)]\n",
        "plt.scatter(hor_axis, DICE, s = 5, marker = '.', c = 'blue')\n",
        "plt.ylabel('Dice Score')\n",
        "plt.xlabel('N° Samples')\n",
        "plt.title('Dice scores')\n",
        "plt.savefig(dice_plot_path)\n",
        "time.sleep(2)\n",
        "\n",
        "\n",
        "#-- Save\n",
        "\n",
        "print('\\nSave Results and Parameters =============>\\n')\n",
        "f = open(results_path, \"w\")\n",
        "f.write(results)       \n",
        "f.close()   \n",
        "                         \n",
        "f = open(params_path, \"w\")\n",
        "f.write(parameters)\n",
        "f.write(score_out)\n",
        "f.close()\n",
        "\n",
        "\n",
        "#-- End\n",
        "\n",
        "print('\\nEnd !\\n')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UO40zWJ96_SY"
      },
      "source": [
        "## Spatial Convolutional Autoencoder"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SX_mLKHp7DKI"
      },
      "source": [
        "#-- Model implementation : Spatial Convolutional Autoencoder\n",
        "\n",
        "def SCAE():\n",
        "  \n",
        "    input_img = tf.keras.Input(shape=(image_size, image_size, num_channels)) \n",
        "\n",
        "    x = layers.Conv2D(32 , 5, activation=layers.LeakyReLU(), strides=2, padding=\"same\")(input_img)\n",
        "    x = layers.Conv2D(64 , 5, activation=layers.LeakyReLU(), strides=2, padding=\"same\")(x)\n",
        "    x = layers.Conv2D(128, 5, activation=layers.LeakyReLU(), strides=2, padding=\"same\")(x)    \n",
        "    x = layers.Conv2D(128, 5, activation=layers.LeakyReLU(), strides=2, padding=\"same\")(x)   \n",
        "    encoded = layers.Conv2D(16 , 1, activation=layers.LeakyReLU(), strides=1, padding=\"same\")(x)\n",
        "   \n",
        "    #-- BOTTELNECK SIZE : (16, 16, 16)\n",
        "\n",
        "    x = layers.Conv2D(128, 1, strides=1, activation=layers.LeakyReLU(), padding=\"same\")(encoded)    \n",
        "    x = layers.Conv2DTranspose(128, 5, strides=2, activation=layers.LeakyReLU(), padding=\"same\")(x) \n",
        "    x = layers.Conv2DTranspose(64 , 5, strides=2, activation=layers.LeakyReLU(), padding=\"same\")(x)\n",
        "    x = layers.Conv2DTranspose(32 , 5, strides=2, activation=layers.LeakyReLU(), padding=\"same\")(x)\n",
        "    x = layers.Conv2DTranspose(32 , 5, strides=2, activation=layers.LeakyReLU(), padding=\"same\")(x)\n",
        "    \n",
        "    decoded = layers.Conv2D(num_channels, 1, activation=layers.LeakyReLU(), padding='same')(x)\n",
        "    \n",
        "    model = tf.keras.Model(input_img, decoded)\n",
        "\n",
        "    return model\n",
        "\n",
        "    \n",
        "#-- Configure the hyperparameters\n",
        "\n",
        "model_name = 'Spatial Convolutional Autoencoder'\n",
        "numEpochs = 50\n",
        "learning_rate = 0.00001\n",
        "image_size = 256\n",
        "num_channels = 1\n",
        "batch_size = 1\n",
        "intermediate_dim = (16, 16, 16)\n",
        "\n",
        "\n",
        "#-- Configure training and testing Datasets \n",
        "\n",
        "saved_dir = './saved/'\n",
        "data_dir  = './data/OASIS/'\n",
        "\n",
        "test_healthy_path = './data/OASIS_all/OASIS_Flair_Test.npy'\n",
        "\n",
        "test_path = './data/BraTS/s0/BraTS_Flair.npy'\n",
        "brainmask_path = './data/BraTS/s0/BraTS_Brainmask.npy'\n",
        "x_prior_path = './data/BraTS/s0/BraTS_prior_52.npy.npy'    \n",
        "label_path = './data/BraTS/s0/BraTS_GT.npy'\n",
        "\n",
        "'''\n",
        "#-- If using MSLUB as test-set\n",
        "test_path = './data/MSLUB/MSLUB_Flair.npy'\n",
        "brainmask_path = './data/MSLUB/MSLUB_Brainmask.npy'\n",
        "x_prior_path = './data/MSLUB/MSLUB_prior_57.npy'\n",
        "label_path = './data/MSLUB/MSLUB_GT.npy'\n",
        "'''\n",
        "\n",
        "list_len = [4805, 3078]    #-- BraTS and MSLUB test-set sizes, respectively.\n",
        "len_testset = list_len[0]  #-- 0 for BraTS and 1 for MSLUB\n",
        "\n",
        "train_paths = list_of_paths(data_dir)\n",
        "\n",
        "nb_train_files = 66\n",
        "data_gen = data_generator(train_paths[:nb_train_files], batch_size)\n",
        "training_steps = (256 / batch_size) * nb_train_files\n",
        "\n",
        "nb_val_files = 5\n",
        "val_gen = data_generator(train_paths[-nb_val_files:], batch_size)\n",
        "validation_steps = (256 / batch_size) * nb_val_files\n",
        "\n",
        "\n",
        "#-- Checkpoints dir\n",
        "\n",
        "date = datetime. now(). strftime(\"%Y_%m_%d-%I:%M:%S_%p\")\n",
        "ckpts_dir = os.path.join(saved_dir, f'Ckpts_{date}')\n",
        "os.makedirs(ckpts_dir)\n",
        "     \n",
        "ckpts_path = os.path.join(ckpts_dir, 'Model_Ckpts.h5')\n",
        "params_path = os.path.join(ckpts_dir, 'Parameters.txt')\n",
        "results_path = os.path.join(ckpts_dir, 'Results.txt')\n",
        "fig_path = os.path.join(ckpts_dir, 'History_plot.png')\n",
        "dice_plot_path = os.path.join(ckpts_dir, 'Dice_plot.png')\n",
        "predicted_path = os.path.join(ckpts_dir, 'Predicted.npy')\n",
        "residual_path = os.path.join(ckpts_dir, 'Residuals.npy')\n",
        "residual_BP_path = os.path.join(ckpts_dir, 'Residuals_BP.npy')\n",
        "      \n",
        "\n",
        "#-- Configure the training\n",
        "\n",
        "opt = tf.keras.optimizers.Adam(learning_rate = learning_rate)\n",
        "          \n",
        "calbks = tf.keras.callbacks.ModelCheckpoint(filepath=ckpts_path, monitor='val_loss', save_best_only=True, save_weights_only=True, verbose=2)\n",
        "tqdm_callback = tfa.callbacks.TQDMProgressBar() \n",
        "\n",
        "model = SCAE()\n",
        "model.summary()\n",
        "model.compile(optimizer=opt, loss='mae', metrics=['mse', SSIMLoss, MS_SSIMLoss])\n",
        "\n",
        "\n",
        "#-- Print & Write model Parameters\n",
        "\n",
        "parameters = (f'\\nSelected model \"{model_name}\" with :\\n - {batch_size}: Batche(s)\\n - {numEpochs}: Epochs\\n - {intermediate_dim}: Bottelneck size\\n')\n",
        "print(parameters)\n",
        "\n",
        "         \n",
        "#-- TRAIN \n",
        "\n",
        "print('\\nTrain =>\\n')\n",
        "history = model.fit(x = data_gen,\n",
        "                    steps_per_epoch = training_steps,\n",
        "                    validation_data = val_gen,\n",
        "                    validation_steps = validation_steps,\n",
        "                    verbose = 0,\n",
        "                    epochs = numEpochs,\n",
        "                    callbacks = [calbks, tqdm_callback]\n",
        "                    )\n",
        "\n",
        "                          \n",
        "#-- Get training and test loss histories   \n",
        "\n",
        "plot_history(history, path=fig_path)\n",
        "plt.close()\n",
        "time.sleep(2)\n",
        "\n",
        "\n",
        "#-- Test  \n",
        "    \n",
        "print('\\nTest ===>\\n')\n",
        "my_test = np.load(test_path)\n",
        "brainmask = np.load(brainmask_path)\n",
        "x_prior = np.load(x_prior_path)\n",
        "my_labels = np.load(label_path)\n",
        "\n",
        "healthy_test = np.load(test_healthy_path)\n",
        "steps = healthy_test.shape[0]\n",
        "\n",
        "score = model.evaluate(x=healthy_test, y=healthy_test, verbose=0, steps=steps, callbacks = [tqdm_callback])    \n",
        "score_out = (f'\\nTest on healthy unseen data :\\n - Test loss (MAE): {score[0]},\\n - Test MSE : {score[1]},\\n - Test SSIM : {score[2]}\\n - Test MS_SSIM : {score[3]}\\n')\n",
        "print(score_out)\n",
        "\n",
        "\n",
        "#-- Predict\n",
        "\n",
        "print('\\nPredict =====>\\n')\n",
        "predicted = model.predict(x=my_test, verbose=1, steps=len_testset)\n",
        "np.save(predicted_path, predicted)\n",
        "time.sleep(4)\n",
        "\n",
        "\n",
        "#-- Calculate, Post-process and Save Residuals\n",
        "\n",
        "print('\\nCalculate, Post-process and Save Residuals =====>\\n')     \n",
        "residual_BP = calculate_residual_BP(my_test, predicted, brainmask)  #-- You can use either brainmask or x_prior\n",
        "np.save(residual_BP_path, residual_BP)\n",
        "        \n",
        "residual = calculate_residual(my_test, predicted, brainmask)  #-- You can use either brainmask or x_prior\n",
        "np.save(residual_path, residual)\n",
        "        \n",
        "\n",
        "#-- Evaluation\n",
        "\n",
        "print('\\nEvaluate =========>\\n')        \n",
        "[AUROC, AUPRC, AVG_DICE, MAD, STD], DICE = eval_residuals(my_labels, residual)     \n",
        "results = (f'\\nResults after median_filter :\\n - AUROC = {AUROC}\\n - AUPRC = {AUPRC}\\n - AVG_DICE = {AVG_DICE}\\n - MEDIAN_ABSOLUTE_DEVIATION = {MAD}\\n - STANDARD_DEVIATION = {STD}')\n",
        "print(results)\n",
        "                    \n",
        "plt.figure()\n",
        "hor_axis = [x for x in range(len_testset)]\n",
        "plt.scatter(hor_axis, DICE, s = 5, marker = '.', c = 'blue')\n",
        "plt.ylabel('Dice Score')\n",
        "plt.xlabel('N° Samples')\n",
        "plt.title('Dice scores')\n",
        "plt.savefig(dice_plot_path)\n",
        "time.sleep(2)\n",
        "\n",
        "\n",
        "#-- Save\n",
        "\n",
        "print('\\nSave Results and Parameters =============>\\n')\n",
        "f = open(results_path, \"w\")\n",
        "f.write(results)       \n",
        "f.close()   \n",
        "                    \n",
        "f = open(params_path, \"w\")\n",
        "f.write(parameters)\n",
        "f.write(score_out)\n",
        "f.close()\n",
        "\n",
        "      \n",
        "#-- End\n",
        "\n",
        "print('\\nEnd !\\n')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JBL3NyI957XU"
      },
      "source": [
        "# Latent Variable Models"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Guj7hFmq7Pjy"
      },
      "source": [
        "## Variational Autoencoder"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ipxZQ5kC5-cr"
      },
      "source": [
        "#-- Model implementation : Variational Autoencoder\n",
        "\n",
        "class Sampling(layers.Layer):\n",
        "    def call(self, inputs):\n",
        "        z_mean, z_log_var = inputs\n",
        "        batch = tf.shape(z_mean)[0]\n",
        "        dim = tf.shape(z_mean)[1]\n",
        "        epsilon = tf.keras.backend.random_normal(shape=(batch, dim))\n",
        "        return z_mean + tf.exp(0.5 * z_log_var) * epsilon\n",
        "        \n",
        "\n",
        "def VAE():\n",
        "\n",
        "  inputs = tf.keras.Input(shape=(image_size, image_size, num_channels))\n",
        "  \n",
        "  x = layers.Conv2D(32 , 5, activation=layers.LeakyReLU(), strides=2, padding=\"same\")(inputs)\n",
        "  x = layers.Conv2D(64 , 5, activation=layers.LeakyReLU(), strides=2, padding=\"same\")(x)\n",
        "  x = layers.Conv2D(128, 5, activation=layers.LeakyReLU(), strides=2, padding=\"same\")(x)    \n",
        "  x = layers.Conv2D(128, 5, activation=layers.LeakyReLU(), strides=2, padding=\"same\")(x)   \n",
        "  x = layers.Conv2D(16 , 1, activation=layers.LeakyReLU(), strides=1, padding=\"same\")(x)   \n",
        "  x = layers.Flatten()(x)\n",
        "  encoded = layers.Dense(intermediate_dim, activation=layers.LeakyReLU())(x)   \n",
        "  \n",
        "  z_mean = layers.Dense(latent_dim, name=\"z_mean\")(encoded)\n",
        "  z_log_var = layers.Dense(latent_dim, name=\"z_log_var\")(encoded)\n",
        "  z = Sampling()([z_mean, z_log_var])\n",
        "  \n",
        "  x = layers.Dense(16 * 16 * 16, activation=layers.LeakyReLU())(z)\n",
        "  x = layers.Reshape((16, 16, 16))(x)\n",
        "  x = layers.Conv2D(128, 1, strides=1, activation=layers.LeakyReLU(), padding=\"same\")(x)    \n",
        "  x = layers.Conv2DTranspose(128, 5, strides=2, activation=layers.LeakyReLU(), padding=\"same\")(x) \n",
        "  x = layers.Conv2DTranspose(64 , 5, strides=2, activation=layers.LeakyReLU(), padding=\"same\")(x)\n",
        "  x = layers.Conv2DTranspose(32 , 5, strides=2, activation=layers.LeakyReLU(), padding=\"same\")(x)\n",
        "  x = layers.Conv2DTranspose(32 , 5, strides=2, activation=layers.LeakyReLU(), padding=\"same\")(x)\n",
        "  decoder_outputs = layers.Conv2D(num_channels, 1, activation=layers.LeakyReLU(), padding='same')(x)\n",
        "  \n",
        "  model = tf.keras.Model(inputs, decoder_outputs)\n",
        "  \n",
        "  return model\n",
        "\n",
        "\n",
        "#-- Configure the hyperparameters\n",
        "\n",
        "model_name = 'Variational Autoencoder'\n",
        "numEpochs = 50\n",
        "learning_rate = 0.00001\n",
        "rate = 0.  \n",
        "image_size = 256\n",
        "num_channels = 1\n",
        "batch_size = 1\n",
        "latent_dim = 512\n",
        "intermediate_dim = 2048\n",
        "\n",
        "\n",
        "#-- Configure training and testing Datasets \n",
        "\n",
        "saved_dir = './saved/'\n",
        "data_dir  = './data/OASIS/'\n",
        "\n",
        "test_healthy_path = './data/OASIS_all/OASIS_Flair_Test.npy'\n",
        "\n",
        "test_path = './data/BraTS/s0/BraTS_Flair.npy'\n",
        "brainmask_path = './data/BraTS/s0/BraTS_Brainmask.npy'\n",
        "x_prior_path = './data/BraTS/s0/BraTS_prior_52.npy.npy' \n",
        "label_path = './data/BraTS/s0/BraTS_GT.npy'\n",
        "\n",
        "'''\n",
        "#-- If using MSLUB as test-set\n",
        "test_path = './data/MSLUB/MSLUB_Flair.npy'\n",
        "brainmask_path = './data/MSLUB/MSLUB_Brainmask.npy'\n",
        "x_prior_path = './data/MSLUB/MSLUB_prior_57.npy'    \n",
        "label_path = './data/MSLUB/MSLUB_GT.npy'\n",
        "'''\n",
        "\n",
        "list_len = [4805, 3078]    #-- BraTS and MSLUB test-set sizes, respectively.\n",
        "len_testset = list_len[0]  #-- 0 for BraTS and 1 for MSLUB\n",
        "\n",
        "train_paths = list_of_paths(data_dir)\n",
        "\n",
        "nb_train_files = 66\n",
        "data_gen = data_generator(train_paths[:nb_train_files], batch_size)\n",
        "training_steps = (256 / batch_size) * nb_train_files\n",
        "\n",
        "nb_val_files = 5\n",
        "val_gen = data_generator(train_paths[-nb_val_files:], batch_size)\n",
        "validation_steps = (256 / batch_size) * nb_val_files\n",
        "\n",
        "\n",
        "#-- Checkpoints dir\n",
        "\n",
        "date = datetime. now(). strftime(\"%Y_%m_%d-%I:%M:%S_%p\")\n",
        "ckpts_dir = os.path.join(saved_dir, f'Ckpts_{date}')\n",
        "os.makedirs(ckpts_dir)\n",
        "     \n",
        "ckpts_path = os.path.join(ckpts_dir, 'Model_Ckpts.h5')\n",
        "params_path = os.path.join(ckpts_dir, 'Parameters.txt')\n",
        "results_path = os.path.join(ckpts_dir, 'Results.txt')\n",
        "fig_path = os.path.join(ckpts_dir, 'History_plot.png')\n",
        "dice_plot_path = os.path.join(ckpts_dir, 'Dice_plot.png')\n",
        "predicted_path = os.path.join(ckpts_dir, 'Predicted.npy')\n",
        "residual_path = os.path.join(ckpts_dir, 'Residuals.npy')\n",
        "residual_BP_path = os.path.join(ckpts_dir, 'Residuals_BP.npy')\n",
        "\n",
        "      \n",
        "#-- Configure the training\n",
        "\n",
        "opt = tf.keras.optimizers.Adam(learning_rate = learning_rate)\n",
        "           \n",
        "calbks = tf.keras.callbacks.ModelCheckpoint(filepath=ckpts_path, monitor='loss', save_best_only=True, save_weights_only=True, verbose=2)\n",
        "tqdm_callback = tfa.callbacks.TQDMProgressBar() \n",
        "\n",
        "model = VAE()\n",
        "model.summary()\n",
        "model.compile(optimizer=opt, loss='mae', metrics=['mse', SSIMLoss, MS_SSIMLoss])\n",
        "\n",
        "\n",
        "#-- Print & Write model Parameters\n",
        "\n",
        "parameters = (f'\\nSelected model \"{model_name}\" with :\\n - {batch_size}: Batche(s)\\n - {numEpochs}: Epochs\\n - {intermediate_dim}: Intermediate dim\\n - {latent_dim}: Bottelneck size\\n')\n",
        "print(parameters)\n",
        "\n",
        "  \n",
        "#-- TRAIN \n",
        "\n",
        "print('\\nTrain =>\\n')\n",
        "history = model.fit(x = data_gen,\n",
        "                    steps_per_epoch = training_steps,\n",
        "                    validation_data = val_gen,\n",
        "                    validation_steps = validation_steps,\n",
        "                    verbose = 0,\n",
        "                    epochs = numEpochs,\n",
        "                    callbacks = [calbks, tqdm_callback]\n",
        "                    )\n",
        "\n",
        "                          \n",
        "#-- Get training and test loss histories \n",
        "\n",
        "plot_history(history, path=fig_path)\n",
        "plt.close()\n",
        "time.sleep(2)\n",
        "\n",
        "\n",
        "#-- Test\n",
        "\n",
        "print('\\nTest ===>\\n')\n",
        "my_test = np.load(test_path)\n",
        "brainmask = np.load(brainmask_path)\n",
        "x_prior = np.load(x_prior_path)\n",
        "my_labels = np.load(label_path)\n",
        "\n",
        "healthy_test = np.load(test_healthy_path)\n",
        "steps = healthy_test.shape[0]\n",
        "\n",
        "score = model.evaluate(x=healthy_test, y=healthy_test, verbose=0, steps=steps, callbacks = [tqdm_callback])    \n",
        "score_out = (f'\\nTest on healthy unseen data :\\n - Test loss (MAE): {score[0]},\\n - Test MSE : {score[1]},\\n - Test SSIM : {score[2]}\\n - Test MS_SSIM : {score[3]}\\n')\n",
        "print(score_out)\n",
        "\n",
        "\n",
        "#-- Predict\n",
        "\n",
        "print('\\nPredict =====>\\n')\n",
        "predicted = model.predict(x=my_test, verbose=1, steps=len_testset)\n",
        "np.save(predicted_path, predicted)\n",
        "time.sleep(4)\n",
        "\n",
        "\n",
        "#-- Calculate, Post-process and Save Residuals\n",
        "\n",
        "print('\\nCalculate, Post-process and Save Residuals =====>\\n')     \n",
        "residual_BP = calculate_residual_BP(my_test, predicted, brainmask)  #-- You can use either brainmask or x_prior\n",
        "np.save(residual_BP_path, residual_BP)\n",
        "        \n",
        "residual = calculate_residual(my_test, predicted, brainmask)  #-- You can use either brainmask or x_prior\n",
        "np.save(residual_path, residual)\n",
        "        \n",
        "\n",
        "#-- Evaluation\n",
        "\n",
        "print('\\nEvaluate =========>\\n')        \n",
        "[AUROC, AUPRC, AVG_DICE, MAD, STD], DICE = eval_residuals(my_labels, residual)     \n",
        "results = (f'\\nResults after median_filter :\\n - AUROC = {AUROC}\\n - AUPRC = {AUPRC}\\n - AVG_DICE = {AVG_DICE}\\n - MEDIAN_ABSOLUTE_DEVIATION = {MAD}\\n - STANDARD_DEVIATION = {STD}')\n",
        "print(results)\n",
        "                      \n",
        "plt.figure()\n",
        "hor_axis = [x for x in range(len_testset)]\n",
        "plt.scatter(hor_axis, DICE, s = 5, marker = '.', c = 'blue')\n",
        "plt.ylabel('Dice Score')\n",
        "plt.xlabel('N° Samples')\n",
        "plt.title('Dice scores')\n",
        "plt.savefig(dice_plot_path)\n",
        "time.sleep(2)\n",
        "\n",
        "\n",
        "#-- Save\n",
        "\n",
        "print('\\nSave Results and Parameters =============>\\n')\n",
        "f = open(results_path, \"w\")\n",
        "f.write(results)       \n",
        "f.close()   \n",
        "                       \n",
        "f = open(params_path, \"w\")\n",
        "f.write(parameters)\n",
        "f.write(score_out)\n",
        "f.close()\n",
        "\n",
        "      \n",
        "#-- End\n",
        "\n",
        "print('\\nEnd !\\n')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9ytUmxAK7Rp4"
      },
      "source": [
        "## VQ-VAE"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ayqdhjM27TX6"
      },
      "source": [
        "#-- Model implementation : Vector Quantizer Variational Autoencoder\n",
        "\n",
        "class VectorQuantizer(layers.Layer):\n",
        "    def __init__(self, num_embeddings, embedding_dim, beta=0.25, **kwargs):\n",
        "        super().__init__(**kwargs)\n",
        "        self.embedding_dim = embedding_dim\n",
        "        self.num_embeddings = num_embeddings\n",
        "        self.beta = (\n",
        "            beta  # This parameter is best kept between [0.25, 2] as per the paper.\n",
        "        )\n",
        "\n",
        "        # Initialize the embeddings which we will quantize.\n",
        "        w_init = tf.random_uniform_initializer()\n",
        "        self.embeddings = tf.Variable(\n",
        "            initial_value=w_init(\n",
        "                shape=(self.embedding_dim, self.num_embeddings), dtype=\"float32\"\n",
        "            ),\n",
        "            trainable=True,\n",
        "            name=\"embeddings_vqvae\",\n",
        "        )\n",
        "\n",
        "    def call(self, x):\n",
        "        # Calculate the input shape of the inputs and\n",
        "        # then flatten the inputs keeping `embedding_dim` intact.\n",
        "        input_shape = tf.shape(x)\n",
        "        flattened = tf.reshape(x, [-1, self.embedding_dim])\n",
        "\n",
        "        # Quantization.\n",
        "        encoding_indices = self.get_code_indices(flattened)\n",
        "        encodings = tf.one_hot(encoding_indices, self.num_embeddings)\n",
        "        quantized = tf.matmul(encodings, self.embeddings, transpose_b=True)\n",
        "        quantized = tf.reshape(quantized, input_shape)\n",
        "\n",
        "        commitment_loss = self.beta * tf.reduce_mean(\n",
        "            (tf.stop_gradient(quantized) - x) ** 2\n",
        "        )\n",
        "        codebook_loss = tf.reduce_mean((quantized - tf.stop_gradient(x)) ** 2)\n",
        "        self.add_loss(commitment_loss + codebook_loss)\n",
        "\n",
        "        # Straight-through estimator.\n",
        "        quantized = x + tf.stop_gradient(quantized - x)\n",
        "        return quantized\n",
        "\n",
        "    def get_code_indices(self, flattened_inputs):\n",
        "        # Calculate L2-normalized distance between the inputs and the codes.\n",
        "        similarity = tf.matmul(flattened_inputs, self.embeddings)\n",
        "        distances = (\n",
        "            tf.reduce_sum(flattened_inputs ** 2, axis=1, keepdims=True)\n",
        "            + tf.reduce_sum(self.embeddings ** 2, axis=0)\n",
        "            - 2 * similarity\n",
        "        )\n",
        "\n",
        "        # Derive the indices for minimum distances.\n",
        "        encoding_indices = tf.argmin(distances, axis=1)\n",
        "        return encoding_indices\n",
        "      \n",
        "              \n",
        "def VQVAE():\n",
        "\n",
        "    inputs = tf.keras.Input(shape=(image_size, image_size, num_channels))\n",
        "    \n",
        "    x = layers.Conv2D(32 , 5, activation=layers.LeakyReLU(), strides=2, padding=\"same\")(inputs)\n",
        "    x = layers.Conv2D(64 , 5, activation=layers.LeakyReLU(), strides=2, padding=\"same\")(x)\n",
        "    x = layers.Conv2D(128, 5, activation=layers.LeakyReLU(), strides=2, padding=\"same\")(x)    \n",
        "    x = layers.Conv2D(128, 5, activation=layers.LeakyReLU(), strides=2, padding=\"same\")(x)   \n",
        "    encoder_outputs = layers.Conv2D(latent_dim , 1, activation=layers.LeakyReLU(), strides=1, padding=\"same\")(x)\n",
        "\n",
        "    quantized_latents = VectorQuantizer(num_embeddings, latent_dim, name=\"vector_quantizer\")(encoder_outputs)\n",
        "\n",
        "    x = layers.Conv2D(128, 1, strides=1, activation=layers.LeakyReLU(), padding=\"same\")(quantized_latents)    \n",
        "    x = layers.Conv2DTranspose(128, 5, strides=2, activation=layers.LeakyReLU(), padding=\"same\")(x) \n",
        "    x = layers.Conv2DTranspose(64 , 5, strides=2, activation=layers.LeakyReLU(), padding=\"same\")(x)\n",
        "    x = layers.Conv2DTranspose(32 , 5, strides=2, activation=layers.LeakyReLU(), padding=\"same\")(x)\n",
        "    x = layers.Conv2DTranspose(32 , 5, strides=2, activation=layers.LeakyReLU(), padding=\"same\")(x)\n",
        "    decoder_outputs = layers.Conv2D(num_channels, 1, activation=layers.LeakyReLU(), padding='same')(x)\n",
        "\n",
        "    model = tf.keras.Model(inputs, decoder_outputs)\n",
        "\n",
        "    return model\n",
        "\n",
        "\n",
        "#-- Configure the hyperparameters\n",
        "\n",
        "model_name = 'Vector Quantizer Variational Autoencoder'\n",
        "numEpochs = 50\n",
        "learning_rate = 0.00001\n",
        "rate = 0.  \n",
        "image_size = 256\n",
        "num_channels = 1\n",
        "batch_size = 1\n",
        "latent_dim = 16\n",
        "num_embeddings = 512\n",
        "intermediate_dim = 4096\n",
        "\n",
        "\n",
        "#-- Configure training and testing Datasets\n",
        "\n",
        "saved_dir = './saved/'\n",
        "data_dir  = './data/OASIS/'\n",
        "\n",
        "test_healthy_path = './data/OASIS_all/OASIS_Flair_Test.npy'\n",
        "\n",
        "test_path = './data/BraTS/s0/BraTS_Flair.npy'\n",
        "brainmask_path = './data/BraTS/s0/BraTS_Brainmask.npy'\n",
        "x_prior_path = './data/BraTS/s0/BraTS_prior_52.npy.npy'\n",
        "label_path = './data/BraTS/s0/BraTS_GT.npy'\n",
        "\n",
        "'''\n",
        "#-- If using MSLUB as test-set\n",
        "test_path = './data/MSLUB/MSLUB_Flair.npy'\n",
        "brainmask_path = './data/MSLUB/MSLUB_Brainmask.npy'\n",
        "x_prior_path = './data/MSLUB/MSLUB_prior_57.npy'\n",
        "label_path = './data/MSLUB/MSLUB_GT.npy'\n",
        "'''\n",
        "\n",
        "list_len = [4805, 3078]    #-- BraTS and MSLUB test-set sizes, respectively.\n",
        "len_testset = list_len[0]  #-- 0 for BraTS and 1 for MSLUB\n",
        "\n",
        "train_paths = list_of_paths(data_dir)\n",
        "\n",
        "nb_train_files = 66\n",
        "data_gen = data_generator(train_paths[:nb_train_files], batch_size)\n",
        "training_steps = (256 / batch_size) * nb_train_files\n",
        "\n",
        "nb_val_files = 5\n",
        "val_gen = data_generator(train_paths[-nb_val_files:], batch_size)\n",
        "validation_steps = (256 / batch_size) * nb_val_files\n",
        "\n",
        "\n",
        "#-- Checkpoints dir\n",
        "\n",
        "date = datetime. now(). strftime(\"%Y_%m_%d-%I:%M:%S_%p\")\n",
        "ckpts_dir = os.path.join(saved_dir, f'Ckpts_{date}')\n",
        "os.makedirs(ckpts_dir)\n",
        "    \n",
        "ckpts_path = os.path.join(ckpts_dir, 'Model_Ckpts.h5')\n",
        "params_path = os.path.join(ckpts_dir, 'Parameters.txt')\n",
        "results_path = os.path.join(ckpts_dir, 'Results.txt')\n",
        "fig_path = os.path.join(ckpts_dir, 'History_plot.png')\n",
        "dice_plot_path = os.path.join(ckpts_dir, 'Dice_plot.png')\n",
        "predicted_path = os.path.join(ckpts_dir, 'Predicted.npy')\n",
        "residual_path = os.path.join(ckpts_dir, 'Residuals.npy')\n",
        "residual_BP_path = os.path.join(ckpts_dir, 'Residuals_BP.npy')\n",
        "\n",
        "      \n",
        "#-- Configure the training\n",
        "\n",
        "opt = tf.keras.optimizers.Adam(learning_rate = learning_rate)    \n",
        "\n",
        "calbks = tf.keras.callbacks.ModelCheckpoint(filepath=ckpts_path, monitor='loss', save_best_only=True, save_weights_only=True, verbose=2)\n",
        "tqdm_callback = tfa.callbacks.TQDMProgressBar()  \n",
        "\n",
        "model = VQVAE()\n",
        "model.summary()\n",
        "model.compile(optimizer=opt, loss='mae', metrics=['mse', SSIMLoss, MS_SSIMLoss])\n",
        "\n",
        "\n",
        "#-- Print & Write model Parameters\n",
        "\n",
        "parameters = (f'\\nSelected model \"{model_name}\" with :\\n - {batch_size}: Batche(s)\\n - {numEpochs}: Epochs\\n - {num_embeddings}: Bottelneck size\\n - {intermediate_dim}: Intermediate dim\\n')\n",
        "print(parameters)\n",
        "\n",
        "\n",
        "#-- TRAIN \n",
        "\n",
        "print('\\nTrain =>\\n')\n",
        "history = model.fit(x = data_gen,\n",
        "                    steps_per_epoch = training_steps,\n",
        "                    validation_data = val_gen,\n",
        "                    validation_steps = validation_steps,\n",
        "                    verbose = 0,\n",
        "                    epochs = numEpochs,\n",
        "                    callbacks = [calbks, tqdm_callback]\n",
        "                    )\n",
        "\n",
        "                          \n",
        "#-- Get training and test loss histories \n",
        "\n",
        "plot_history(history, path=fig_path)\n",
        "plt.close()\n",
        "time.sleep(2)\n",
        "\n",
        "\n",
        "#-- Test  \n",
        "\n",
        "print('\\nTest ===>\\n')\n",
        "my_test = np.load(test_path)\n",
        "brainmask = np.load(brainmask_path)\n",
        "x_prior = np.load(x_prior_path)\n",
        "my_labels = np.load(label_path)\n",
        "\n",
        "healthy_test = np.load(test_healthy_path)\n",
        "steps = healthy_test.shape[0]\n",
        "\n",
        "score = model.evaluate(x=healthy_test, y=healthy_test, verbose=0, steps=steps, callbacks = [tqdm_callback])    \n",
        "score_out = (f'\\nTest on healthy unseen data :\\n - Test loss (MAE): {score[0]},\\n - Test MSE : {score[1]},\\n - Test SSIM : {score[2]}\\n - Test MS_SSIM : {score[3]}\\n')\n",
        "print(score_out)\n",
        "\n",
        "\n",
        "#-- Predict\n",
        "\n",
        "print('\\nPredict =====>\\n')\n",
        "predicted = model.predict(x=my_test, verbose=1, steps=len_testset)\n",
        "np.save(predicted_path, predicted)\n",
        "time.sleep(4)\n",
        "\n",
        "\n",
        "#-- Calculate, Post-process and Save Residuals\n",
        "\n",
        "print('\\nCalculate, Post-process and Save Residuals =====>\\n')     \n",
        "residual_BP = calculate_residual_BP(my_test, predicted, brainmask)  #-- You can use either brainmask or x_prior\n",
        "np.save(residual_BP_path, residual_BP)\n",
        "        \n",
        "residual = calculate_residual(my_test, predicted, brainmask)  #-- You can use either brainmask or x_prior\n",
        "np.save(residual_path, residual)\n",
        "        \n",
        "\n",
        "#-- Evaluation\n",
        "\n",
        "print('\\nEvaluate =========>\\n')        \n",
        "[AUROC, AUPRC, AVG_DICE, MAD, STD], DICE = eval_residuals(my_labels, residual)     \n",
        "results = (f'\\nResults after median_filter :\\n - AUROC = {AUROC}\\n - AUPRC = {AUPRC}\\n - AVG_DICE = {AVG_DICE}\\n - MEDIAN_ABSOLUTE_DEVIATION = {MAD}\\n - STANDARD_DEVIATION = {STD}')\n",
        "print(results)\n",
        "                     \n",
        "plt.figure()\n",
        "hor_axis = [x for x in range(len_testset)]\n",
        "plt.scatter(hor_axis, DICE, s = 5, marker = '.', c = 'blue')\n",
        "plt.ylabel('Dice Score')\n",
        "plt.xlabel('N° Samples')\n",
        "plt.title('Dice scores')\n",
        "plt.savefig(dice_plot_path)\n",
        "time.sleep(2)\n",
        "\n",
        "\n",
        "#-- Save\n",
        "\n",
        "print('\\nSave Results and Parameters =============>\\n')\n",
        "f = open(results_path, \"w\")\n",
        "f.write(results)       \n",
        "f.close()   \n",
        "                   \n",
        "f = open(params_path, \"w\")\n",
        "f.write(parameters)\n",
        "f.write(score_out)\n",
        "f.close()\n",
        "  \n",
        "   \n",
        "#-- End\n",
        "\n",
        "print('\\nEnd !\\n')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NeeKAOcK6Ahf"
      },
      "source": [
        "# Transformer based models"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nuEgxobC7T_N"
      },
      "source": [
        "## B_TAE"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LOZbeiQa6D4W"
      },
      "source": [
        "#-- Implement patch creation as a layer\n",
        "\n",
        "class Patches(layers.Layer):\n",
        "\n",
        "    def __init__(self, patch_size):\n",
        "        super(Patches, self).__init__()\n",
        "        self.patch_size = patch_size\n",
        "    \n",
        "    def get_config(self):\n",
        "\n",
        "        config = super().get_config().copy()\n",
        "        config.update({\n",
        "            'patch_size': self.patch_size\n",
        "        })\n",
        "        return config\n",
        "\n",
        "    def call(self, images):\n",
        "        batch_size = tf.shape(images)[0]\n",
        "        patches = tf.image.extract_patches(\n",
        "            images=images,\n",
        "            sizes=[1, self.patch_size, self.patch_size, 1],\n",
        "            strides=[1, self.patch_size, self.patch_size, 1],\n",
        "            rates=[1, 1, 1, 1],\n",
        "            padding=\"VALID\")\n",
        "        patch_dims = patches.shape[-1]\n",
        "        patches = tf.reshape(patches, [batch_size, -1, patch_dims])\n",
        "        return patches\n",
        "\n",
        "\n",
        "#-- Implement the patch encoding layer\n",
        "\n",
        "class PatchEncoder(layers.Layer):\n",
        "\n",
        "    def __init__(self, embed_shape):\n",
        "        super(PatchEncoder, self).__init__()\n",
        "        self.embed_shape = embed_shape\n",
        "        self.position_embedding = layers.Embedding(\n",
        "            input_dim = embed_shape[0],\n",
        "            output_dim = embed_shape[1]\n",
        "            )\n",
        "        self.projection = layers.Dense(units=embed_shape[1])\n",
        "        \n",
        "    def get_config(self):\n",
        "      config = super().get_config().copy()\n",
        "      config.update({\n",
        "          'embed_shape': self.embed_shape,\n",
        "          'position_embedding': self.position_embedding,\n",
        "          'projection': self.projection\n",
        "        })\n",
        "      return config\n",
        "\n",
        "    def call(self, patches):\n",
        "        positions = tf.range(start=0, limit=self.embed_shape[0], delta=1)\n",
        "        encoded = self.projection(patches) + self.position_embedding(positions)\n",
        "        return encoded\n",
        "\n",
        "\n",
        "#-- Implement image reconstruction from patches as a layer\n",
        "\n",
        "class Images(layers.Layer):\n",
        "  \n",
        "    def __init__(self, image_size, patch_size, num_channels):\n",
        "        super(Images, self).__init__()\n",
        "        self.image_size = image_size\n",
        "        self.patch_size = patch_size\n",
        "        self.num_channels = num_channels\n",
        "\n",
        "    def get_config(self):\n",
        "        config = super().get_config().copy()\n",
        "        config.update({\n",
        "            'image_size': self.image_size,\n",
        "            'patch_size': self.patch_size,\n",
        "            'num_channels': self.num_channels\n",
        "          })\n",
        "        return config\n",
        "\n",
        "    def call(self, patches):\n",
        "        batch_size = tf.shape(patches)[0]\n",
        "        reconstructed = tf.reshape(patches, [batch_size, self.image_size, self.image_size, num_channels])\n",
        "        rec_new = tf.nn.space_to_depth(reconstructed, self.patch_size) \n",
        "        image = tf.reshape(rec_new, [batch_size, self.image_size, self.image_size, self.num_channels])   \n",
        "        return image\n",
        "\n",
        "\n",
        "#-- Dense Layer    \n",
        "\n",
        "class TruncatedDense(layers.Dense):\n",
        "    def __init__(self, units, use_bias=True, initializer = tf.keras.initializers.TruncatedNormal(mean=0., stddev=.02)):\n",
        "        super().__init__(units, use_bias=use_bias, kernel_initializer=initializer)\n",
        "\n",
        "        \n",
        "#-- Mlp Head\n",
        "\n",
        "class Mlp(layers.Layer):\n",
        "\n",
        "    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=layers.Activation(tf.nn.gelu), drop=0.):\n",
        "        super().__init__()\n",
        "        out_features = out_features or in_features\n",
        "        hidden_features = hidden_features or in_features\n",
        "        self.fc1 = TruncatedDense(hidden_features)\n",
        "        self.act = act_layer\n",
        "        self.fc2 = TruncatedDense(out_features)\n",
        "        self.drop = layers.Dropout(drop)\n",
        "    \n",
        "    def get_config(self):\n",
        "        config = super().get_config().copy()\n",
        "        config.update({\n",
        "            'fc1': self.fc1,\n",
        "            'act': self.act,\n",
        "            'fc2': self.fc2,\n",
        "            'drop': self.drop\n",
        "          })\n",
        "        return config\n",
        "\n",
        "    def call(self, x):\n",
        "        x = self.fc1(x)\n",
        "        x = self.act(x)\n",
        "        x = self.drop(x)\n",
        "        x = self.fc2(x)\n",
        "        x = self.drop(x)\n",
        "        return x\n",
        "\n",
        "\n",
        "#-- Model implementation : Basic Vision Transformer Autoencoder \n",
        "\n",
        "def transformer_autoencoder(encoded_patches):\n",
        "\n",
        "    for _ in range(transformer_layers):\n",
        "        query = key = encoded_patches\n",
        "        attn_encoded_patches = layers.MultiHeadAttention(\n",
        "            num_heads = num_heads,\n",
        "            key_dim = embed_shape[1],\n",
        "            dropout = rate)(query=query, value=encoded_patches, key=key)   \n",
        "        attn_encoded_patches = layers.Dropout(rate)(attn_encoded_patches) \n",
        "        encoded_patches += attn_encoded_patches  \n",
        "        encoded_patches = layers.LayerNormalization(epsilon=1e-6)(encoded_patches)\n",
        "        ffn_out = Mlp(\n",
        "            in_features=embed_shape[1],\n",
        "            hidden_features=4*embed_shape[1],\n",
        "            drop=rate)(encoded_patches)\n",
        "        encoded_patches += ffn_out   \n",
        "        encoded_patches = layers.LayerNormalization(epsilon=1e-6)(encoded_patches)\n",
        "\n",
        "    positions = tf.range(start=0, limit = embed_shape[0], delta=1)\n",
        "    pos_embed = layers.Embedding(\n",
        "        input_dim = embed_shape[0],\n",
        "        output_dim = embed_shape[1])(positions)\n",
        "    target = encoded_patches\n",
        "\n",
        "    for _ in range(transformer_layers):\n",
        "        query_tgt = key_tgt = target + pos_embed    \n",
        "        attn_target1 = layers.MultiHeadAttention(\n",
        "            num_heads = num_heads,\n",
        "            key_dim = embed_shape[1],\n",
        "            dropout=rate)(query=query_tgt, value=target, key=key_tgt)   \n",
        "        attn_target1 = layers.Dropout(rate)(attn_target1)\n",
        "        target += attn_target1\n",
        "        target = layers.LayerNormalization(epsilon=1e-6)(target)\n",
        "        query_tgt = target + pos_embed\n",
        "        attn_target2 = layers.MultiHeadAttention(\n",
        "            num_heads = num_heads,\n",
        "            key_dim = embed_shape[1],\n",
        "            dropout=rate)(query=query_tgt, value=encoded_patches + pos_embed, key=encoded_patches)\n",
        "        attn_target2 = layers.Dropout(rate)(attn_target2)\n",
        "        target += attn_target2\n",
        "        target = layers.LayerNormalization(epsilon=1e-6)(target)\n",
        "        ffn_out = Mlp(\n",
        "            in_features=embed_shape[1],\n",
        "            hidden_features=4*embed_shape[1],\n",
        "            drop=rate)(encoded_patches)\n",
        "        target +=ffn_out\n",
        "        target = layers.LayerNormalization(epsilon=1e-6)(target)\n",
        "    return target\n",
        "\n",
        "\n",
        "#-- Build the B_TAE model\n",
        "\n",
        "def TAE():\n",
        "\n",
        "    input_img = tf.keras.Input(shape=(image_size, image_size, num_channels))\n",
        "    \n",
        "    patches = Patches(patch_size)(input_img)\n",
        "    enc_patches = PatchEncoder(embed_shape)(patches)\n",
        "    decoded_patches = transformer_autoencoder(enc_patches)\n",
        "    decoded_patches = layers.Dense(units=patch_size*patch_size*num_channels)(decoded_patches)  \n",
        "    reconstructed = Images(image_size, patch_size, num_channels)(decoded_patches)\n",
        "\n",
        "    model = tf.keras.Model(input_img, reconstructed)\n",
        "\n",
        "    return model\n",
        " \n",
        " \n",
        "#-- Configure the hyperparameters\n",
        "\n",
        "model_name = 'BASIC Transformer'\n",
        "numEpochs = 50\n",
        "learning_rate = 0.00001\n",
        "rate = 0.1  \n",
        "image_size = 256\n",
        "num_channels = 1\n",
        "batch_size = 4\n",
        "filters = 96\n",
        "\n",
        "\n",
        "#-- Transformer parameters variation\n",
        "\n",
        "param_grid = {\n",
        "              'transformer_layers': [8],\n",
        "              'patch_size' : [16],\n",
        "              'num_heads' : [4]\n",
        "             }\n",
        "             \n",
        "            \n",
        "PARAMS = ParameterGrid(param_grid)\n",
        "list_PARAMS = list(PARAMS)\n",
        "\n",
        "\n",
        "#-- Configure training and testing Datasets \n",
        "\n",
        "saved_dir = './saved/'\n",
        "data_dir  = './data/OASIS/'\n",
        "\n",
        "test_healthy_path = './data/OASIS_all/OASIS_Flair_Test.npy'\n",
        "\n",
        "test_path = './data/BraTS/s0/BraTS_Flair.npy'\n",
        "brainmask_path = './data/BraTS/s0/BraTS_Brainmask.npy'\n",
        "x_prior_path = './data/BraTS/s0/BraTS_prior_52.npy.npy'\n",
        "label_path = './data/BraTS/s0/BraTS_GT.npy'\n",
        "\n",
        "'''\n",
        "#-- If using MSLUB as test-set\n",
        "test_path = './data/MSLUB/MSLUB_Flair.npy'\n",
        "brainmask_path = './data/MSLUB/MSLUB_Brainmask.npy'\n",
        "x_prior_path = './data/MSLUB/MSLUB_prior_57.npy'\n",
        "label_path = './data/MSLUB/MSLUB_GT.npy'\n",
        "'''\n",
        "\n",
        "list_len = [4805, 3078]    #-- BraTS and MSLUB test-set sizes, respectively.\n",
        "len_testset = list_len[0]  #-- 0 for BraTS and 1 for MSLUB\n",
        "\n",
        "train_paths = list_of_paths(data_dir)\n",
        "\n",
        "nb_train_files = 66\n",
        "data_gen = data_generator(train_paths[:nb_train_files], batch_size)\n",
        "training_steps = (256 / batch_size) * nb_train_files\n",
        "\n",
        "nb_val_files = 5\n",
        "val_gen = data_generator(train_paths[-nb_val_files:], batch_size)\n",
        "validation_steps = (256 / batch_size) * nb_val_files\n",
        "\n",
        "\n",
        "#-- Train, Test and Evaluate\n",
        "\n",
        "for param in list_PARAMS:\n",
        "   try:\n",
        "   \n",
        "        #-- Checkpoints dir\n",
        "\n",
        "        date = datetime. now(). strftime(\"%Y_%m_%d-%I:%M:%S_%p\")\n",
        "        ckpts_dir = os.path.join(saved_dir, f'Ckpts_{date}')\n",
        "        os.makedirs(ckpts_dir)\n",
        "        \n",
        "        ckpts_path = os.path.join(ckpts_dir, 'Model_Ckpts.h5')\n",
        "        params_path = os.path.join(ckpts_dir, 'Parameters.txt')\n",
        "        results_path = os.path.join(ckpts_dir, 'Results.txt')\n",
        "        fig_path = os.path.join(ckpts_dir, 'History_plot.png')\n",
        "        dice_plot_path = os.path.join(ckpts_dir, 'Dice_plot.png')\n",
        "        predicted_path = os.path.join(ckpts_dir, 'Predicted.npy')\n",
        "        residual_path = os.path.join(ckpts_dir, 'Residuals.npy')\n",
        "        residual_BP_path = os.path.join(ckpts_dir, 'Residuals_BP.npy')\n",
        "      \n",
        "\n",
        "        #-- Configure the parameters\n",
        "\n",
        "        transformer_layers = param['transformer_layers']\n",
        "        patch_size = param['patch_size']\n",
        "        num_heads = param['num_heads']\n",
        "        input_resolution = image_size // patch_size\n",
        "        num_patches = input_resolution ** 2\n",
        "        embed_shape = (num_patches, filters)\n",
        "\n",
        "\n",
        "        #-- Configure the training\n",
        "\n",
        "        opt = tf.keras.optimizers.Adam(learning_rate = learning_rate)\n",
        "            \n",
        "        calbks = tf.keras.callbacks.ModelCheckpoint(filepath=ckpts_path, monitor='val_loss', save_best_only=True, save_weights_only=True, verbose=2)\n",
        "        tqdm_callback = tfa.callbacks.TQDMProgressBar() \n",
        "\n",
        "        model = TAE()\n",
        "        model.summary()        \n",
        "        model.compile(optimizer=opt, loss='mae', metrics=['mse', SSIMLoss, MS_SSIMLoss])\n",
        "\n",
        "\n",
        "        #-- Print & Write model Parameters\n",
        "\n",
        "        parameters = (f'\\nSelected model \"{model_name}\" with :\\n - {num_heads}: Heads,\\n - {patch_size}: Patch size,\\n - {transformer_layers}: Transformer Layer,\\n - ({embed_shape[0]}, {embed_shape[1]}): Embedding Shape,\\n - {batch_size}: Batche(s)\\n - {numEpochs}: Epochs\\n')\n",
        "        print(parameters)\n",
        "     \n",
        "        \n",
        "        #-- Train \n",
        "\n",
        "        print('\\nTrain =>\\n')\n",
        "        history = model.fit(x = data_gen,\n",
        "                            steps_per_epoch = training_steps,\n",
        "                            validation_data = val_gen,\n",
        "                            validation_steps = validation_steps,\n",
        "                            verbose = 0,\n",
        "                            epochs = numEpochs,\n",
        "                            callbacks = [calbks, tqdm_callback]\n",
        "                            )\n",
        "        \n",
        "                        \n",
        "        #-- Get training and test loss histories \n",
        "\n",
        "        plot_history(history, path=fig_path)\n",
        "        plt.close()\n",
        "        time.sleep(2)\n",
        "        \n",
        "                \n",
        "        #-- Test \n",
        "\n",
        "        print('\\nTest ===>\\n')\n",
        "        my_test = np.load(test_path)\n",
        "        brainmask = np.load(brainmask_path)\n",
        "        x_prior = np.load(x_prior_path)\n",
        "        my_labels = np.load(label_path)    \n",
        "                \n",
        "        healthy_test = np.load(test_healthy_path)\n",
        "        steps = healthy_test.shape[0]\n",
        "                \n",
        "        score = model.evaluate(x=healthy_test, y=healthy_test, verbose=0, steps=steps, callbacks = [tqdm_callback])\n",
        "        score_out = (f'\\nTest on healthy unseen data :\\n - Test loss (MAE): {score[0]},\\n - Test MSE : {score[1]},\\n - Test SSIM: {score[2]},\\n - Test MS_SSIMLoss: {score[3]},\\n')\n",
        "        print(score_out)\n",
        "        \n",
        "                \n",
        "        #-- Predict\n",
        "\n",
        "        print('\\nPredict =====>\\n')\n",
        "        predicted = model.predict(x=my_test, verbose=1, steps=len_testset)\n",
        "        np.save(predicted_path, predicted)\n",
        "        time.sleep(4)\n",
        "        \n",
        "        \n",
        "        #-- Calculate, Post-process and Save Residuals\n",
        "\n",
        "        print('\\nCalculate, Post-process and Save Residuals =====>\\n')     \n",
        "        residual_BP = calculate_residual_BP(my_test, predicted, brainmask)  #-- You can use either brainmask or x_prior\n",
        "        np.save(residual_BP_path, residual_BP)\n",
        "        \n",
        "        residual = calculate_residual(my_test, predicted, brainmask)  #-- You can use either brainmask or x_prior\n",
        "        np.save(residual_path, residual)\n",
        "        \n",
        "\n",
        "        #-- Evaluation\n",
        "\n",
        "        print('\\nEvaluate =========>\\n')        \n",
        "        [AUROC, AUPRC, AVG_DICE, MAD, STD], DICE = eval_residuals(my_labels, residual)     \n",
        "        results = (f'\\nResults after median_filter :\\n - AUROC = {AUROC}\\n - AUPRC = {AUPRC}\\n - AVG_DICE = {AVG_DICE}\\n - MEDIAN_ABSOLUTE_DEVIATION = {MAD}\\n - STANDARD_DEVIATION = {STD}')\n",
        "        print(results)\n",
        "      \n",
        "                 \n",
        "        plt.figure()\n",
        "        hor_axis = [x for x in range(len_testset)]\n",
        "        plt.scatter(hor_axis, DICE, s = 5, marker = '.', c = 'blue')\n",
        "        plt.ylabel('Dice Score')\n",
        "        plt.xlabel('N° Samples')\n",
        "        plt.title('Dice scores')\n",
        "        plt.savefig(dice_plot_path)\n",
        "        time.sleep(2)\n",
        "\n",
        "\n",
        "        #-- Save\n",
        "\n",
        "        print('\\nSave Results and Parameters =============>\\n')\n",
        "        f = open(results_path, \"w\")\n",
        "        f.write(results)       \n",
        "        f.close()   \n",
        "         \n",
        "                        \n",
        "        f = open(params_path, \"w\")\n",
        "        f.write(parameters)\n",
        "        f.write(score_out)\n",
        "        f.close()\n",
        "        \n",
        "         \n",
        "        #-- End\n",
        "\n",
        "        print('\\nEnd !\\n')\n",
        "\n",
        "   except:\n",
        "        print('Error encountered !')\n",
        "        tf.keras.backend.clear_session()\n",
        "        continue"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ob9yVvQg7V__"
      },
      "source": [
        "## DC_TAE"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uenJD7Lb7Vi-"
      },
      "source": [
        "#-- Implement patch creation as a layer\n",
        "\n",
        "class Patches(layers.Layer):\n",
        "\n",
        "    def __init__(self, patch_size):\n",
        "        super(Patches, self).__init__()\n",
        "        self.patch_size = patch_size\n",
        "    \n",
        "    def get_config(self):\n",
        "\n",
        "        config = super().get_config().copy()\n",
        "        config.update({\n",
        "            'patch_size': self.patch_size\n",
        "        })\n",
        "        return config\n",
        "\n",
        "    def call(self, images):\n",
        "        batch_size = tf.shape(images)[0]\n",
        "        patches = tf.image.extract_patches(\n",
        "            images=images,\n",
        "            sizes=[1, self.patch_size, self.patch_size, 1],\n",
        "            strides=[1, self.patch_size, self.patch_size, 1],\n",
        "            rates=[1, 1, 1, 1],\n",
        "            padding=\"VALID\")\n",
        "        patch_dims = patches.shape[-1]\n",
        "        patches = tf.reshape(patches, [batch_size, -1, patch_dims])\n",
        "        return patches\n",
        "\n",
        "\n",
        "#-- Implement the patch encoding layer\n",
        "\n",
        "class PatchEncoder(layers.Layer):\n",
        "\n",
        "    def __init__(self, embed_shape):\n",
        "        super(PatchEncoder, self).__init__()\n",
        "        self.embed_shape = embed_shape\n",
        "        self.position_embedding = layers.Embedding(\n",
        "            input_dim = embed_shape[0],\n",
        "            output_dim = embed_shape[1]\n",
        "            )\n",
        "        self.projection = layers.Dense(units=embed_shape[1])\n",
        "        \n",
        "    def get_config(self):\n",
        "      config = super().get_config().copy()\n",
        "      config.update({\n",
        "          'embed_shape': self.embed_shape,\n",
        "          'position_embedding': self.position_embedding,\n",
        "          'projection': self.projection\n",
        "        })\n",
        "      return config\n",
        "\n",
        "    def call(self, patches):\n",
        "        positions = tf.range(start=0, limit=self.embed_shape[0], delta=1)\n",
        "        encoded = self.projection(patches) + self.position_embedding(positions)\n",
        "        return encoded\n",
        "\n",
        "\n",
        "#-- Implement image reconstruction from patches as a layer\n",
        "\n",
        "class Images(layers.Layer):\n",
        "  \n",
        "    def __init__(self, image_size, patch_size, num_channels):\n",
        "        super(Images, self).__init__()\n",
        "        self.image_size = image_size\n",
        "        self.patch_size = patch_size\n",
        "        self.num_channels = num_channels\n",
        "\n",
        "    def get_config(self):\n",
        "        config = super().get_config().copy()\n",
        "        config.update({\n",
        "            'image_size': self.image_size,\n",
        "            'patch_size': self.patch_size,\n",
        "            'num_channels': self.num_channels\n",
        "          })\n",
        "        return config\n",
        "\n",
        "    def call(self, patches):\n",
        "        batch_size = tf.shape(patches)[0]\n",
        "        reconstructed = tf.reshape(patches, [batch_size, self.image_size, self.image_size, num_channels])\n",
        "        rec_new = tf.nn.space_to_depth(reconstructed, self.patch_size) \n",
        "        image = tf.reshape(rec_new, [batch_size, self.image_size, self.image_size, self.num_channels])   \n",
        "        return image\n",
        "\n",
        "\n",
        "#-- Dense Layer    \n",
        "\n",
        "class TruncatedDense(layers.Dense):\n",
        "    def __init__(self, units, use_bias=True, initializer = tf.keras.initializers.TruncatedNormal(mean=0., stddev=.02)):\n",
        "        super().__init__(units, use_bias=use_bias, kernel_initializer=initializer)\n",
        "\n",
        "\n",
        "#-- Mlp Head\n",
        "\n",
        "class Mlp(layers.Layer):\n",
        "\n",
        "    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=layers.Activation(tf.nn.gelu), drop=0.):\n",
        "        super().__init__()\n",
        "        out_features = out_features or in_features\n",
        "        hidden_features = hidden_features or in_features\n",
        "        self.fc1 = TruncatedDense(hidden_features)\n",
        "        self.act = act_layer\n",
        "        self.fc2 = TruncatedDense(out_features)\n",
        "        self.drop = layers.Dropout(drop)\n",
        "    \n",
        "    def get_config(self):\n",
        "        config = super().get_config().copy()\n",
        "        config.update({\n",
        "            'fc1': self.fc1,\n",
        "            'act': self.act,\n",
        "            'fc2': self.fc2,\n",
        "            'drop': self.drop\n",
        "          })\n",
        "        return config\n",
        "\n",
        "    def call(self, x):\n",
        "        x = self.fc1(x)\n",
        "        x = self.act(x)\n",
        "        x = self.drop(x)\n",
        "        x = self.fc2(x)\n",
        "        x = self.drop(x)\n",
        "        return x\n",
        "\n",
        "\n",
        "#-- Dense Convolutional Autoencoder\n",
        "\n",
        "def DCAE(input_image):\n",
        "\n",
        "    x = layers.Conv2D(32 , 3, activation=layers.LeakyReLU(), strides=2, padding=\"same\")(input_image)\n",
        "    x = layers.Conv2D(64 , 3, activation=layers.LeakyReLU(), strides=2, padding=\"same\")(x)\n",
        "    x = layers.Conv2D(128, 3, activation=layers.LeakyReLU(), strides=2, padding=\"same\")(x)    \n",
        "    x = layers.Conv2D(128, 3, activation=layers.LeakyReLU(), strides=2, padding=\"same\")(x)   \n",
        "    x = layers.Conv2D(16 , 1, activation=layers.LeakyReLU(), strides=1, padding=\"same\")(x)\n",
        "    \n",
        "    x = layers.Flatten()(x)\n",
        "    encoded = layers.Dense(intermediate_dim, activation=layers.LeakyReLU())(x)\n",
        "\n",
        "    #-- BOTTELNECK SIZE : 512 \n",
        "    \n",
        "    x = layers.Dense(16 * 16 * 16, activation=layers.LeakyReLU())(encoded)\n",
        "    x = layers.Reshape((16, 16, 16))(x)\n",
        "\n",
        "    x = layers.Conv2D(128, 1, strides=1, activation=layers.LeakyReLU(), padding=\"same\")(x)    \n",
        "    x = layers.Conv2DTranspose(128, 3, strides=2, activation=layers.LeakyReLU(), padding=\"same\")(x) \n",
        "    x = layers.Conv2DTranspose(64 , 3, strides=2, activation=layers.LeakyReLU(), padding=\"same\")(x)\n",
        "    x = layers.Conv2DTranspose(32 , 3, strides=2, activation=layers.LeakyReLU(), padding=\"same\")(x)\n",
        "    x = layers.Conv2DTranspose(32 , 3, strides=2, activation=layers.LeakyReLU(), padding=\"same\")(x)\n",
        "    \n",
        "    decoded = layers.Conv2D(num_channels, 1, activation=layers.LeakyReLU(), padding='same')(x)\n",
        "\n",
        "    return decoded\n",
        "\n",
        "\n",
        "#-- Model implementation : Dense Convolutional Autoencoder inside Transformer\n",
        "\n",
        "def transformer_autoencoder(encoded_patches):\n",
        "\n",
        "    for _ in range(transformer_layers):\n",
        "        query = key = encoded_patches\n",
        "        attn_encoded_patches = layers.MultiHeadAttention(\n",
        "            num_heads = num_heads,\n",
        "            key_dim = embed_shape[1],\n",
        "            dropout = rate)(query=query, value=encoded_patches, key=key)   \n",
        "        attn_encoded_patches = layers.Dropout(rate)(attn_encoded_patches) \n",
        "        encoded_patches += attn_encoded_patches  \n",
        "        encoded_patches = layers.LayerNormalization(epsilon=1e-6)(encoded_patches)\n",
        "        ffn_out = Mlp(\n",
        "            in_features=embed_shape[1],\n",
        "            hidden_features=4*embed_shape[1],\n",
        "            drop=rate)(encoded_patches)\n",
        "        encoded_patches += ffn_out   \n",
        "        encoded_patches = layers.LayerNormalization(epsilon=1e-6)(encoded_patches)\n",
        "    \n",
        "    #-- Convolutional AutoEncoder\n",
        "    encoded_img = Images(image_size, patch_size, num_channels)(encoded_patches)\n",
        "    decoded_patches = DCAE(encoded_img)  \n",
        "    encoded_patches = Patches(patch_size)(decoded_patches)\n",
        "    encoded_patches = PatchEncoder(embed_shape)(encoded_patches)\n",
        "\n",
        "    positions = tf.range(start=0, limit = embed_shape[0], delta=1)\n",
        "    pos_embed = layers.Embedding(\n",
        "        input_dim = embed_shape[0],\n",
        "        output_dim = embed_shape[1])(positions)\n",
        "    target = encoded_patches\n",
        "\n",
        "    for _ in range(transformer_layers):\n",
        "        query_tgt = key_tgt = target + pos_embed    \n",
        "        attn_target1 = layers.MultiHeadAttention(\n",
        "            num_heads = num_heads,\n",
        "            key_dim = embed_shape[1],\n",
        "            dropout=rate)(query=query_tgt, value=target, key=key_tgt)   \n",
        "        attn_target1 = layers.Dropout(rate)(attn_target1)\n",
        "        target += attn_target1\n",
        "        target = layers.LayerNormalization(epsilon=1e-6)(target)\n",
        "        query_tgt = target + pos_embed\n",
        "        attn_target2 = layers.MultiHeadAttention(\n",
        "            num_heads = num_heads,\n",
        "            key_dim = embed_shape[1],\n",
        "            dropout=rate)(query=query_tgt, value=encoded_patches + pos_embed, key=encoded_patches)\n",
        "        attn_target2 = layers.Dropout(rate)(attn_target2)\n",
        "        target += attn_target2\n",
        "        target = layers.LayerNormalization(epsilon=1e-6)(target)\n",
        "        ffn_out = Mlp(\n",
        "            in_features=embed_shape[1],\n",
        "            hidden_features=4*embed_shape[1],\n",
        "            drop=rate)(encoded_patches)\n",
        "        target +=ffn_out\n",
        "        target = layers.LayerNormalization(epsilon=1e-6)(target)\n",
        "    return target\n",
        "\n",
        "\n",
        "#-- Build the DC_TAE Model\n",
        "\n",
        "def TAE():\n",
        "  \n",
        "    input_img = tf.keras.Input(shape=(image_size, image_size, num_channels))\n",
        "\n",
        "    patches = Patches(patch_size)(input_img)    \n",
        "    enc_patches = PatchEncoder(embed_shape)(patches)\n",
        "\n",
        "    decoded_patches = transformer_autoencoder(enc_patches)\n",
        "  \n",
        "    reconstructed = Images(image_size, patch_size, num_channels)(decoded_patches)\n",
        "\n",
        "    model = tf.keras.Model(input_img, reconstructed)\n",
        "    \n",
        "    return model\n",
        " \n",
        "    \n",
        "#-- Configure the hyperparameters\n",
        "\n",
        "model_name = 'DCAE Inside Transformer'\n",
        "numEpochs = 50\n",
        "learning_rate = 0.00001\n",
        "rate = 0.1  \n",
        "image_size = 256\n",
        "num_channels = 1\n",
        "batch_size = 1\n",
        "intermediate_dim = 512\n",
        "\n",
        "\n",
        "#-- Transformer parameters variation\n",
        "\n",
        "param_grid = {\n",
        "              'transformer_layers': [8],      \n",
        "              'patch_size' : [16],\n",
        "              'num_heads' : [4]\n",
        "             }\n",
        "             \n",
        "            \n",
        "PARAMS = ParameterGrid(param_grid)\n",
        "list_PARAMS = list(PARAMS)\n",
        "\n",
        "\n",
        "#-- Configure training and testing Datasets \n",
        "\n",
        "saved_dir = './saved/'\n",
        "data_dir  = './data/OASIS/'\n",
        "\n",
        "test_healthy_path = './data/OASIS_all/OASIS_Flair_Test.npy'\n",
        "\n",
        "test_path = './data/BraTS/s0/BraTS_Flair.npy'\n",
        "brainmask_path = './data/BraTS/s0/BraTS_Brainmask.npy'\n",
        "x_prior_path = './data/BraTS/s0/BraTS_prior_52.npy.npy'\n",
        "label_path = './data/BraTS/s0/BraTS_GT.npy'\n",
        "\n",
        "'''\n",
        "#-- If using MSLUB as test-set\n",
        "test_path = './data/MSLUB/MSLUB_Flair.npy'\n",
        "brainmask_path = './data/MSLUB/MSLUB_Brainmask.npy'\n",
        "x_prior_path = './data/MSLUB/MSLUB_prior_57.npy'\n",
        "label_path = './data/MSLUB/MSLUB_GT.npy'\n",
        "'''\n",
        "\n",
        "list_len = [4805, 3078]    #-- BraTS and MSLUB test-set sizes, respectively.\n",
        "len_testset = list_len[0]  #-- 0 for BraTS and 1 for MSLUB\n",
        "\n",
        "train_paths = list_of_paths(data_dir)\n",
        "\n",
        "nb_train_files = 66\n",
        "data_gen = data_generator(train_paths[:nb_train_files], batch_size)\n",
        "training_steps = (256 / batch_size) * nb_train_files\n",
        "\n",
        "nb_val_files = 5\n",
        "val_gen = data_generator(train_paths[-nb_val_files:], batch_size)\n",
        "validation_steps = (256 / batch_size) * nb_val_files\n",
        "\n",
        "\n",
        "#-- Train, Test and Evaluate\n",
        "\n",
        "for param in list_PARAMS:\n",
        "  try:\n",
        "\n",
        "        #-- Checkpoints dir\n",
        "\n",
        "        date = datetime. now(). strftime(\"%Y_%m_%d-%I:%M:%S_%p\")\n",
        "        ckpts_dir = os.path.join(saved_dir, f'Ckpts_{date}')\n",
        "        os.makedirs(ckpts_dir)\n",
        "        \n",
        "        ckpts_path = os.path.join(ckpts_dir, 'Model_Ckpts.h5')\n",
        "        params_path = os.path.join(ckpts_dir, 'Parameters.txt')\n",
        "        results_path = os.path.join(ckpts_dir, 'Results.txt')\n",
        "        fig_path = os.path.join(ckpts_dir, 'History_plot.png')\n",
        "        dice_plot_path = os.path.join(ckpts_dir, 'Dice_plot.png')\n",
        "        predicted_path = os.path.join(ckpts_dir, 'Predicted.npy')\n",
        "        residual_path = os.path.join(ckpts_dir, 'Residuals.npy')\n",
        "        residual_BP_path = os.path.join(ckpts_dir, 'Residuals_BP.npy')\n",
        "      \n",
        "\n",
        "        #-- Configure the parameters\n",
        "\n",
        "        transformer_layers = param['transformer_layers']\n",
        "        patch_size = param['patch_size']\n",
        "        num_heads = param['num_heads']\n",
        "        input_resolution = image_size // patch_size\n",
        "        num_patches = input_resolution ** 2\n",
        "        embed_shape = (num_patches, patch_size*patch_size*num_channels)\n",
        "\n",
        "\n",
        "        #-- Configure the training\n",
        "\n",
        "        opt = tf.keras.optimizers.Adam(learning_rate = learning_rate)\n",
        "            \n",
        "        calbks = tf.keras.callbacks.ModelCheckpoint(filepath=ckpts_path, monitor='val_loss', save_best_only=True, save_weights_only=True, verbose=2)\n",
        "        tqdm_callback = tfa.callbacks.TQDMProgressBar() \n",
        "\n",
        "        model = TAE()\n",
        "        model.summary()        \n",
        "        model.compile(optimizer=opt, loss='mae', metrics=['mse', SSIMLoss, MS_SSIMLoss])       \n",
        "    \n",
        "\n",
        "        #-- Print & Write model Parameters\n",
        "\n",
        "        parameters = (f'\\nSelected model \"{model_name}\" with :\\n - {num_heads}: Heads,\\n - {patch_size}: Patch size,\\n - {transformer_layers}: Transformer Layer,\\n - ({embed_shape[0]}, {embed_shape[1]}): Embedding Shape,\\n - {batch_size}: Batche(s)\\n - {numEpochs}: Epochs\\n')\n",
        "        print(parameters)\n",
        "        \n",
        "        \n",
        "        #-- Train\n",
        "\n",
        "        print('\\nTrain =>\\n')\n",
        "        history = model.fit(x = data_gen,\n",
        "                            steps_per_epoch = training_steps,\n",
        "                            validation_data = val_gen,\n",
        "                            validation_steps = validation_steps,\n",
        "                            verbose = 0,\n",
        "                            epochs = numEpochs,\n",
        "                            callbacks = [calbks, tqdm_callback]\n",
        "                            )\n",
        "        \n",
        "                        \n",
        "        #-- Get training and test loss histories \n",
        "\n",
        "        plot_history(history, path=fig_path)\n",
        "        plt.close()\n",
        "        time.sleep(2)              \n",
        "        \n",
        "        \n",
        "        #-- Test\n",
        "\n",
        "        print('\\nTest ===>\\n')\n",
        "        my_test = np.load(test_path)\n",
        "        brainmask = np.load(brainmask_path)\n",
        "        x_prior = np.load(x_prior_path)\n",
        "        my_labels = np.load(label_path)              \n",
        "               \n",
        "        healthy_test = np.load(test_healthy_path)\n",
        "        steps = healthy_test.shape[0]\n",
        "                \n",
        "        score = model.evaluate(x=healthy_test, y=healthy_test, verbose=0, steps=steps, callbacks = [tqdm_callback])\n",
        "        score_out = (f'\\nTest on healthy unseen data :\\n - Test loss (MSE): {score[0]},\\n - Test MAE : {score[1]},\\n - Test SSIM: {score[2]},\\n - Test MS_SSIMLoss: {score[3]},\\n')\n",
        "        print(score_out)\n",
        "        \n",
        "                \n",
        "        #-- Predict\n",
        "\n",
        "        print('\\nPredict =====>\\n')\n",
        "        predicted = model.predict(x=my_test, verbose=1, steps=len_testset)\n",
        "        np.save(predicted_path, predicted)\n",
        "        time.sleep(4)\n",
        "        \n",
        "        \n",
        "        #-- Calculate, Post-process and Save Residuals\n",
        "\n",
        "        print('\\nCalculate, Post-process and Save Residuals =====>\\n')     \n",
        "        residual_BP = calculate_residual_BP(my_test, predicted, brainmask)  #-- You can use either brainmask or x_prior\n",
        "        np.save(residual_BP_path, residual_BP)\n",
        "        \n",
        "        residual = calculate_residual(my_test, predicted, brainmask)  #-- You can use either brainmask or x_prior\n",
        "        np.save(residual_path, residual)\n",
        "        \n",
        "\n",
        "        #-- Evaluation\n",
        "\n",
        "        print('\\nEvaluate =========>\\n')        \n",
        "        [AUROC, AUPRC, AVG_DICE, MAD, STD], DICE = eval_residuals(my_labels, residual)     \n",
        "        results = (f'\\nResults after median_filter :\\n - AUROC = {AUROC}\\n - AUPRC = {AUPRC}\\n - AVG_DICE = {AVG_DICE}\\n - MEDIAN_ABSOLUTE_DEVIATION = {MAD}\\n - STANDARD_DEVIATION = {STD}')\n",
        "        print(results)\n",
        "                       \n",
        "        plt.figure()\n",
        "        hor_axis = [x for x in range(len_testset)]\n",
        "        plt.scatter(hor_axis, DICE, s = 5, marker = '.', c = 'blue')\n",
        "        plt.ylabel('Dice Score')\n",
        "        plt.xlabel('N° Samples')\n",
        "        plt.title('Dice scores')\n",
        "        plt.savefig(dice_plot_path)\n",
        "        time.sleep(2)\n",
        "\n",
        "\n",
        "        #-- Save\n",
        "\n",
        "        print('\\nSave Results and Parameters =============>\\n')\n",
        "        f = open(results_path, \"w\")\n",
        "        f.write(results)       \n",
        "        f.close()   \n",
        "                    \n",
        "        f = open(params_path, \"w\")\n",
        "        f.write(parameters)\n",
        "        f.write(score_out)\n",
        "        f.close()       \n",
        "         \n",
        "\n",
        "        #-- End\n",
        "\n",
        "        print('\\nEnd of step !\\n')\n",
        "   \n",
        "\n",
        "  except:\n",
        "        print('Error encountered !')\n",
        "        tf.keras.backend.clear_session()\n",
        "        continue"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PrHDynS67XoN"
      },
      "source": [
        "## SC_TAE"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0J09PmPi7ZCF"
      },
      "source": [
        "#-- Implement patch creation as a layer\n",
        "\n",
        "class Patches(layers.Layer):\n",
        "\n",
        "    def __init__(self, patch_size):\n",
        "        super(Patches, self).__init__()\n",
        "        self.patch_size = patch_size\n",
        "    \n",
        "    def get_config(self):\n",
        "\n",
        "        config = super().get_config().copy()\n",
        "        config.update({\n",
        "            'patch_size': self.patch_size\n",
        "        })\n",
        "        return config\n",
        "\n",
        "    def call(self, images):\n",
        "        batch_size = tf.shape(images)[0]\n",
        "        patches = tf.image.extract_patches(\n",
        "            images=images,\n",
        "            sizes=[1, self.patch_size, self.patch_size, 1],\n",
        "            strides=[1, self.patch_size, self.patch_size, 1],\n",
        "            rates=[1, 1, 1, 1],\n",
        "            padding=\"VALID\")\n",
        "        patch_dims = patches.shape[-1]\n",
        "        patches = tf.reshape(patches, [batch_size, -1, patch_dims])\n",
        "        return patches\n",
        "\n",
        "\n",
        "#-- Implement the patch encoding layer\n",
        "\n",
        "class PatchEncoder(layers.Layer):\n",
        "\n",
        "    def __init__(self, embed_shape):\n",
        "        super(PatchEncoder, self).__init__()\n",
        "        self.embed_shape = embed_shape\n",
        "        self.position_embedding = layers.Embedding(\n",
        "            input_dim = embed_shape[0],\n",
        "            output_dim = embed_shape[1]\n",
        "            )\n",
        "        self.projection = layers.Dense(units=embed_shape[1])\n",
        "        \n",
        "    def get_config(self):\n",
        "      config = super().get_config().copy()\n",
        "      config.update({\n",
        "          'embed_shape': self.embed_shape,\n",
        "          'position_embedding': self.position_embedding,\n",
        "          'projection': self.projection\n",
        "        })\n",
        "      return config\n",
        "\n",
        "    def call(self, patches):\n",
        "        positions = tf.range(start=0, limit=self.embed_shape[0], delta=1)\n",
        "        encoded = self.projection(patches) + self.position_embedding(positions)\n",
        "        return encoded\n",
        "\n",
        "\n",
        "#-- Implement image reconstruction from patches as a layer\n",
        "\n",
        "class Images(layers.Layer):\n",
        "  \n",
        "    def __init__(self, image_size, patch_size, num_channels):\n",
        "        super(Images, self).__init__()\n",
        "        self.image_size = image_size\n",
        "        self.patch_size = patch_size\n",
        "        self.num_channels = num_channels\n",
        "\n",
        "    def get_config(self):\n",
        "        config = super().get_config().copy()\n",
        "        config.update({\n",
        "            'image_size': self.image_size,\n",
        "            'patch_size': self.patch_size,\n",
        "            'num_channels': self.num_channels\n",
        "          })\n",
        "        return config\n",
        "\n",
        "    def call(self, patches):\n",
        "        batch_size = tf.shape(patches)[0]\n",
        "        reconstructed = tf.reshape(patches, [batch_size, self.image_size, self.image_size, num_channels])\n",
        "        rec_new = tf.nn.space_to_depth(reconstructed, self.patch_size) \n",
        "        image = tf.reshape(rec_new, [batch_size, self.image_size, self.image_size, self.num_channels])   \n",
        "        return image\n",
        "\n",
        "\n",
        "#-- Dense Layer  \n",
        "\n",
        "class TruncatedDense(layers.Dense):\n",
        "    def __init__(self, units, use_bias=True, initializer = tf.keras.initializers.TruncatedNormal(mean=0., stddev=.02)):\n",
        "        super().__init__(units, use_bias=use_bias, kernel_initializer=initializer)\n",
        "\n",
        "\n",
        "#-- Mlp Head\n",
        "\n",
        "class Mlp(layers.Layer):\n",
        "\n",
        "    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=layers.Activation(tf.nn.gelu), drop=0.):\n",
        "        super().__init__()\n",
        "        out_features = out_features or in_features\n",
        "        hidden_features = hidden_features or in_features\n",
        "        self.fc1 = TruncatedDense(hidden_features)\n",
        "        self.act = act_layer\n",
        "        self.fc2 = TruncatedDense(out_features)\n",
        "        self.drop = layers.Dropout(drop)\n",
        "    \n",
        "    def get_config(self):\n",
        "        config = super().get_config().copy()\n",
        "        config.update({\n",
        "            'fc1': self.fc1,\n",
        "            'act': self.act,\n",
        "            'fc2': self.fc2,\n",
        "            'drop': self.drop\n",
        "          })\n",
        "        return config\n",
        "\n",
        "    def call(self, x):\n",
        "        x = self.fc1(x)\n",
        "        x = self.act(x)\n",
        "        x = self.drop(x)\n",
        "        x = self.fc2(x)\n",
        "        x = self.drop(x)\n",
        "        return x\n",
        "\n",
        "\n",
        "#-- Spatial Convolutional Autoencoder\n",
        "\n",
        "def SCAE(input_img):\n",
        "  \n",
        "    x = layers.Conv2D(32 , 3, activation=layers.LeakyReLU(), strides=2, padding=\"same\")(input_img)\n",
        "    x = layers.Conv2D(64 , 3, activation=layers.LeakyReLU(), strides=2, padding=\"same\")(x)\n",
        "    x = layers.Conv2D(128, 3, activation=layers.LeakyReLU(), strides=2, padding=\"same\")(x)    \n",
        "    x = layers.Conv2D(128, 3, activation=layers.LeakyReLU(), strides=2, padding=\"same\")(x)   \n",
        "    encoded = layers.Conv2D(16 , 1, activation=layers.LeakyReLU(), strides=1, padding=\"same\")(x)\n",
        "   \n",
        "    #-- BOTTELNECK SIZE : (16, 16, 16)\n",
        "\n",
        "    x = layers.Conv2D(128, 1, strides=1, activation=layers.LeakyReLU(), padding=\"same\")(encoded)    \n",
        "    x = layers.Conv2DTranspose(128, 3, strides=2, activation=layers.LeakyReLU(), padding=\"same\")(x) \n",
        "    x = layers.Conv2DTranspose(64 , 3, strides=2, activation=layers.LeakyReLU(), padding=\"same\")(x)\n",
        "    x = layers.Conv2DTranspose(32 , 3, strides=2, activation=layers.LeakyReLU(), padding=\"same\")(x)\n",
        "    x = layers.Conv2DTranspose(32 , 3, strides=2, activation=layers.LeakyReLU(), padding=\"same\")(x)\n",
        "    \n",
        "    decoded = layers.Conv2D(num_channels, 1, activation=layers.LeakyReLU(), padding='same')(x)\n",
        "\n",
        "    return decoded\n",
        "  \n",
        "\n",
        "#-- Model implementation : Spatial Convolutional Autoencoder inside Transformer\n",
        "\n",
        "def transformer_autoencoder(encoded_patches):\n",
        "\n",
        "    for _ in range(transformer_layers):\n",
        "        query = key = encoded_patches\n",
        "        attn_encoded_patches = layers.MultiHeadAttention(\n",
        "            num_heads = num_heads,\n",
        "            key_dim = embed_shape[1],\n",
        "            dropout = rate)(query=query, value=encoded_patches, key=key)   \n",
        "        attn_encoded_patches = layers.Dropout(rate)(attn_encoded_patches) \n",
        "        encoded_patches += attn_encoded_patches  \n",
        "        encoded_patches = layers.LayerNormalization(epsilon=1e-6)(encoded_patches)\n",
        "        ffn_out = Mlp(\n",
        "            in_features=embed_shape[1],\n",
        "            hidden_features=4*embed_shape[1],\n",
        "            drop=rate)(encoded_patches)\n",
        "        encoded_patches += ffn_out   \n",
        "        encoded_patches = layers.LayerNormalization(epsilon=1e-6)(encoded_patches)\n",
        "    \n",
        "    #-- Convolutional AutoEncoder\n",
        "    encoded_img = Images(image_size, patch_size, num_channels)(encoded_patches)\n",
        "    decoded_patches = SCAE(encoded_img)  \n",
        "    encoded_patches = Patches(patch_size)(decoded_patches)\n",
        "    encoded_patches = PatchEncoder(embed_shape)(encoded_patches)\n",
        "\n",
        "    positions = tf.range(start=0, limit = embed_shape[0], delta=1)\n",
        "    pos_embed = layers.Embedding(\n",
        "        input_dim = embed_shape[0],\n",
        "        output_dim = embed_shape[1])(positions)\n",
        "    target = encoded_patches\n",
        "\n",
        "    for _ in range(transformer_layers):\n",
        "        query_tgt = key_tgt = target + pos_embed    \n",
        "        attn_target1 = layers.MultiHeadAttention(\n",
        "            num_heads = num_heads,\n",
        "            key_dim = embed_shape[1],\n",
        "            dropout=rate)(query=query_tgt, value=target, key=key_tgt)   \n",
        "        attn_target1 = layers.Dropout(rate)(attn_target1)\n",
        "        target += attn_target1\n",
        "        target = layers.LayerNormalization(epsilon=1e-6)(target)\n",
        "        query_tgt = target + pos_embed\n",
        "        attn_target2 = layers.MultiHeadAttention(\n",
        "            num_heads = num_heads,\n",
        "            key_dim = embed_shape[1],\n",
        "            dropout=rate)(query=query_tgt, value=encoded_patches + pos_embed, key=encoded_patches)\n",
        "        attn_target2 = layers.Dropout(rate)(attn_target2)\n",
        "        target += attn_target2\n",
        "        target = layers.LayerNormalization(epsilon=1e-6)(target)\n",
        "        ffn_out = Mlp(\n",
        "            in_features=embed_shape[1],\n",
        "            hidden_features=4*embed_shape[1],\n",
        "            drop=rate)(encoded_patches)\n",
        "        target +=ffn_out\n",
        "        target = layers.LayerNormalization(epsilon=1e-6)(target)\n",
        "    return target\n",
        "\n",
        "\n",
        "#-- Build the SC_TAE Model\n",
        "\n",
        "def TAE():\n",
        "  \n",
        "    input_img = tf.keras.Input(shape=(image_size, image_size, num_channels))\n",
        "\n",
        "    patches = Patches(patch_size)(input_img)    \n",
        "    enc_patches = PatchEncoder(embed_shape)(patches)\n",
        "\n",
        "    decoded_patches = transformer_autoencoder(enc_patches)\n",
        "  \n",
        "    reconstructed = Images(image_size, patch_size, num_channels)(decoded_patches)\n",
        "\n",
        "    model = tf.keras.Model(input_img, reconstructed)\n",
        "    \n",
        "    return model\n",
        " \n",
        "    \n",
        "#-- Configure the hyperparameters\n",
        "\n",
        "model_name = 'SCAE Inside Transformer'\n",
        "numEpochs = 50\n",
        "learning_rate = 0.00001\n",
        "rate = 0.1  \n",
        "image_size = 256\n",
        "num_channels = 1\n",
        "batch_size = 1 \n",
        "\n",
        "\n",
        "#-- Transformer parameters variation\n",
        "\n",
        "param_grid = {\n",
        "              'transformer_layers': [8],\n",
        "              'patch_size' : [16],\n",
        "              'num_heads' : [4]\n",
        "             }\n",
        "             \n",
        "            \n",
        "PARAMS = ParameterGrid(param_grid)\n",
        "list_PARAMS = list(PARAMS)\n",
        "\n",
        "\n",
        "#-- Configure training and testing Datasets \n",
        "\n",
        "saved_dir = './saved/'\n",
        "data_dir  = './data/OASIS/'\n",
        "\n",
        "test_healthy_path = './data/OASIS_all/OASIS_Flair_Test.npy'\n",
        "\n",
        "test_path = './data/BraTS/s0/BraTS_Flair.npy'\n",
        "brainmask_path = './data/BraTS/s0/BraTS_Brainmask.npy'\n",
        "x_prior_path = './data/BraTS/s0/BraTS_prior_52.npy.npy'\n",
        "label_path = './data/BraTS/s0/BraTS_GT.npy'\n",
        "\n",
        "'''\n",
        "#-- If using MSLUB as test-set\n",
        "test_path = './data/MSLUB/MSLUB_Flair.npy'\n",
        "brainmask_path = './data/MSLUB/MSLUB_Brainmask.npy'\n",
        "x_prior_path = './data/MSLUB/MSLUB_prior_57.npy'\n",
        "label_path = './data/MSLUB/MSLUB_GT.npy'\n",
        "'''\n",
        "\n",
        "list_len = [4805, 3078]    #-- BraTS and MSLUB test-set sizes, respectively.\n",
        "len_testset = list_len[0]  #-- 0 for BraTS and 1 for MSLUB\n",
        "\n",
        "train_paths = list_of_paths(data_dir)\n",
        "\n",
        "nb_train_files = 66\n",
        "data_gen = data_generator(train_paths[:nb_train_files], batch_size)\n",
        "training_steps = (256 / batch_size) * nb_train_files\n",
        "\n",
        "nb_val_files = 5\n",
        "val_gen = data_generator(train_paths[-nb_val_files:], batch_size)\n",
        "validation_steps = (256 / batch_size) * nb_val_files\n",
        "\n",
        "\n",
        "#-- Train, Test and Evaluate\n",
        "\n",
        "for param in list_PARAMS:\n",
        "   try:   \n",
        "\n",
        "        #-- Checkpoints dir\n",
        "\n",
        "        date = datetime. now(). strftime(\"%Y_%m_%d-%I:%M:%S_%p\")\n",
        "        ckpts_dir = os.path.join(saved_dir, f'Ckpts_{date}')\n",
        "        os.makedirs(ckpts_dir)\n",
        "\n",
        "        ckpts_path = os.path.join(ckpts_dir, 'Model_Ckpts.h5')\n",
        "        params_path = os.path.join(ckpts_dir, 'Parameters.txt')\n",
        "        results_path = os.path.join(ckpts_dir, 'Results.txt')\n",
        "        fig_path = os.path.join(ckpts_dir, 'History_plot.png')\n",
        "        dice_plot_path = os.path.join(ckpts_dir, 'Dice_plot.png')\n",
        "        predicted_path = os.path.join(ckpts_dir, 'Predicted.npy')\n",
        "        residual_path = os.path.join(ckpts_dir, 'Residuals.npy')\n",
        "        residual_BP_path = os.path.join(ckpts_dir, 'Residuals_BP.npy')\n",
        "      \n",
        "\n",
        "        #-- Configure the parameters\n",
        "\n",
        "        transformer_layers = param['transformer_layers']\n",
        "        patch_size = param['patch_size']\n",
        "        num_heads = param['num_heads']\n",
        "        input_resolution = image_size // patch_size\n",
        "        num_patches = input_resolution ** 2\n",
        "        embed_shape = (num_patches, patch_size*patch_size*num_channels)\n",
        "\n",
        "\n",
        "        #-- Configure the training\n",
        "\n",
        "        opt = tf.keras.optimizers.Adam(learning_rate = learning_rate)\n",
        "            \n",
        "        calbks = tf.keras.callbacks.ModelCheckpoint(filepath=ckpts_path, monitor='val_loss', save_best_only=True, save_weights_only=True, verbose=2)\n",
        "        tqdm_callback = tfa.callbacks.TQDMProgressBar() \n",
        "\n",
        "        model = TAE()\n",
        "        model.summary()        \n",
        "        model.compile(optimizer=opt, loss='mae', metrics=['mse', SSIMLoss, MS_SSIMLoss])\n",
        "\n",
        "\n",
        "        #-- Print & Write model Parameters\n",
        "\n",
        "        parameters = (f'\\nSelected model \"{model_name}\" with :\\n - {num_heads}: Heads,\\n - {patch_size}: Patch size,\\n - {transformer_layers}: Transformer Layer,\\n - ({embed_shape[0]}, {embed_shape[1]}): Embedding Shape,\\n - {batch_size}: Batche(s)\\n - {numEpochs}: Epochs\\n')\n",
        "        print(parameters)\n",
        "        \n",
        "        \n",
        "        #-- Train \n",
        "\n",
        "        print('\\nTrain =>\\n')\n",
        "        history = model.fit(x = data_gen,\n",
        "                            steps_per_epoch = training_steps,\n",
        "                            validation_data = val_gen,\n",
        "                            validation_steps = validation_steps,\n",
        "                            verbose = 0,\n",
        "                            epochs = numEpochs,\n",
        "                            callbacks = [calbks, tqdm_callback]\n",
        "                            )\n",
        "        \n",
        "                        \n",
        "        #-- Get training and test loss histories\n",
        "\n",
        "        plot_history(history, path=fig_path)\n",
        "        plt.close()\n",
        "        time.sleep(2)       \n",
        "               \n",
        "        \n",
        "        #-- Test \n",
        "\n",
        "        print('\\nTest ===>\\n')\n",
        "        my_test = np.load(test_path)\n",
        "        brainmask = np.load(brainmask_path)\n",
        "        x_prior = np.load(x_prior_path)\n",
        "        my_labels = np.load(label_path)                     \n",
        "        \n",
        "        healthy_test = np.load(test_healthy_path)\n",
        "        steps = healthy_test.shape[0]        \n",
        "            \n",
        "        score = model.evaluate(x=healthy_test, y=healthy_test, verbose=0, steps=steps, callbacks = [tqdm_callback])\n",
        "        score_out = (f'\\nTest on healthy unseen data :\\n - Test loss (MAE): {score[0]},\\n - Test MSE : {score[1]},\\n - Test SSIM: {score[2]},\\n - Test MS_SSIMLoss: {score[3]},\\n')\n",
        "        print(score_out)\n",
        "                      \n",
        "                        \n",
        "        #-- Predict\n",
        "\n",
        "        print('\\nPredict =====>\\n')\n",
        "        predicted = model.predict(x=my_test, verbose=1, steps=len_testset)\n",
        "        np.save(predicted_path, predicted)\n",
        "        time.sleep(4)\n",
        "        \n",
        "        \n",
        "        #-- Calculate, Post-process and Save Residuals\n",
        "\n",
        "        print('\\nCalculate, Post-process and Save Residuals =====>\\n')     \n",
        "        residual_BP = calculate_residual_BP(my_test, predicted, brainmask)  #-- You can use either brainmask or x_prior\n",
        "        np.save(residual_BP_path, residual_BP)\n",
        "        \n",
        "        residual = calculate_residual(my_test, predicted, brainmask)  #-- You can use either brainmask or x_prior\n",
        "        np.save(residual_path, residual)\n",
        "        \n",
        "\n",
        "        #-- Evaluation\n",
        "\n",
        "        print('\\nEvaluate =========>\\n')        \n",
        "        [AUROC, AUPRC, AVG_DICE, MAD, STD], DICE = eval_residuals(my_labels, residual)     \n",
        "        results = (f'\\nResults after median_filter :\\n - AUROC = {AUROC}\\n - AUPRC = {AUPRC}\\n - AVG_DICE = {AVG_DICE}\\n - MEDIAN_ABSOLUTE_DEVIATION = {MAD}\\n - STANDARD_DEVIATION = {STD}')\n",
        "        print(results)\n",
        "                       \n",
        "        plt.figure()\n",
        "        hor_axis = [x for x in range(len_testset)]\n",
        "        plt.scatter(hor_axis, DICE, s = 5, marker = '.', c = 'blue')\n",
        "        plt.ylabel('Dice Score')\n",
        "        plt.xlabel('N° Samples')\n",
        "        plt.title('Dice scores')\n",
        "        plt.savefig(dice_plot_path)\n",
        "        time.sleep(2)\n",
        "\n",
        "\n",
        "        #-- Save\n",
        "\n",
        "        print('\\nSave Results and Parameters =============>\\n')\n",
        "        f = open(results_path, \"w\")\n",
        "        f.write(results)       \n",
        "        f.close()   \n",
        "                              \n",
        "        f = open(params_path, \"w\")\n",
        "        f.write(parameters)\n",
        "        f.write(score_out)\n",
        "        f.close()\n",
        "        \n",
        "         \n",
        "        #-- End\n",
        "\n",
        "        print('\\nEnd !\\n')\n",
        "\n",
        "\n",
        "   except:\n",
        "        print('Error encountered !')\n",
        "        tf.keras.backend.clear_session()\n",
        "        continue"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QucRsBiN7b5F"
      },
      "source": [
        "## H_TAE"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FnW34a5k7d02"
      },
      "source": [
        "#-- Dense Layer   \n",
        "\n",
        "class TruncatedDense(layers.Dense):\n",
        "    def __init__(self, units, use_bias=True, initializer = tf.keras.initializers.TruncatedNormal(mean=0., stddev=.02)):\n",
        "        super().__init__(units, use_bias=use_bias, kernel_initializer=initializer)\n",
        "\n",
        "\n",
        "#-- Patch Merging Layer\n",
        "\n",
        "class PatchMerging(layers.Layer):\n",
        "\n",
        "    def __init__(self, input_resolution, dim):\n",
        "        super().__init__()\n",
        "        self.input_resolution = input_resolution\n",
        "        self.dim = dim    #--- refer to the projection_dim (nC)\n",
        "        self.reduction = TruncatedDense(2 * dim, use_bias=False)\n",
        "        self.norm = layers.LayerNormalization(epsilon=1e-5)\n",
        "    \n",
        "    def get_config(self):\n",
        "        config = super().get_config().copy()\n",
        "        config.update({\n",
        "            'input_resolution': self.input_resolution,\n",
        "            'dim': self.dim,\n",
        "            'reduction': self.reduction,\n",
        "            'norm': self.norm\n",
        "          })\n",
        "        return config    \n",
        "\n",
        "    def call(self, x):\n",
        "        \"\"\"\n",
        "        x: B, H*W, C\n",
        "        \"\"\"\n",
        "        H = W = self.input_resolution\n",
        "        B = tf.shape(x)[0]\n",
        "\n",
        "        x = tf.reshape(x, [B, H, W, self.dim])\n",
        "        x0 = x[:, 0::2, 0::2, :]  # B H/2 W/2 C\n",
        "        x1 = x[:, 1::2, 0::2, :]  # B H/2 W/2 C\n",
        "        x2 = x[:, 0::2, 1::2, :]  # B H/2 W/2 C\n",
        "        x3 = x[:, 1::2, 1::2, :]  # B H/2 W/2 C\n",
        "        x = tf.concat([x0, x1, x2, x3], -1)  # B H/2 W/2 4*C\n",
        "        x = tf.reshape(x, [B, (H//2)*(W//2), 4 * self.dim])  # B H/2*W/2 4*C\n",
        "        x = self.norm(x)\n",
        "        x = self.reduction(x)\n",
        "\n",
        "        return x\n",
        "\n",
        "\n",
        "#-- Patch Expanding Layer  \n",
        "\n",
        "class PatchExpanding(layers.Layer):\n",
        "\n",
        "    def __init__(self, input_resolution, dim):\n",
        "        super().__init__()\n",
        "        self.input_resolution = input_resolution\n",
        "        self.dim = dim\n",
        "        self.expand = TruncatedDense(2 * dim, use_bias=False)\n",
        "        self.norm = layers.LayerNormalization(epsilon=1e-5)\n",
        "    \n",
        "    def get_config(self):\n",
        "        config = super().get_config().copy()\n",
        "        config.update({\n",
        "            'input_resolution': self.input_resolution,\n",
        "            'dim': self.dim,\n",
        "            'expand': self.expand,\n",
        "            'norm': self.norm\n",
        "          })\n",
        "        return config\n",
        "\n",
        "    def call(self, x):\n",
        "        \"\"\"\n",
        "        x: B, H*W, C\n",
        "        \"\"\"\n",
        "        H = W = self.input_resolution\n",
        "        x = self.expand(x)\n",
        "        B = tf.shape(x)[0]\n",
        "\n",
        "        x = tf.reshape(x, [B, H, W, self.dim * 2])\n",
        "        x = rearrange(x, 'b h w (p1 p2 c)-> b (h p1) (w p2) c', p1=2, p2=2, c=self.dim // 2)\n",
        "\n",
        "        x = tf.reshape(x, [B, (H*2)*(W*2), self.dim//2])\n",
        "        x = self.norm(x)\n",
        "\n",
        "        return x\n",
        "\n",
        "\n",
        "#-- Final Patch Expanding Layer\n",
        "\n",
        "class FinalPatchExpand_X4(layers.Layer):\n",
        "\n",
        "    def __init__(self, input_resolution, dim):\n",
        "        super().__init__()\n",
        "        self.input_resolution = input_resolution\n",
        "        self.dim = dim\n",
        "        self.expand = TruncatedDense(16 * dim, use_bias=False)\n",
        "        self.norm = layers.LayerNormalization(epsilon=1e-5)\n",
        "    \n",
        "    def get_config(self):\n",
        "        config = super().get_config().copy()\n",
        "        config.update({\n",
        "            'input_resolution': self.input_resolution,\n",
        "            'dim': self.dim,\n",
        "            'expand': self.expand,\n",
        "            'norm': self.norm\n",
        "          })\n",
        "        return config\n",
        "\n",
        "    def call(self, x):\n",
        "        \"\"\"\n",
        "        x: B, H*W, C\n",
        "        \"\"\"\n",
        "        H = W = self.input_resolution\n",
        "        x = self.expand(x)\n",
        "        B = tf.shape(x)[0]\n",
        "\n",
        "        x = tf.reshape(x, [B, H, W, self.dim * 16])\n",
        "        x = rearrange(x, 'b h w (p1 p2 c)-> b (h p1) (w p2) c', p1=4, p2=4, c=self.dim)\n",
        "        x = tf.reshape(x, [B, (H*4)*(W*4), self.dim])\n",
        "\n",
        "        x = self.norm(x)\n",
        "\n",
        "        return x\n",
        "\n",
        "\n",
        "#-- Mlp Head\n",
        "\n",
        "class Mlp(layers.Layer):\n",
        "\n",
        "    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=layers.Activation(tf.nn.gelu), drop=0.):\n",
        "        super().__init__()\n",
        "        out_features = out_features or in_features\n",
        "        hidden_features = hidden_features or in_features\n",
        "        self.fc1 = TruncatedDense(hidden_features)\n",
        "        self.act = act_layer\n",
        "        self.fc2 = TruncatedDense(out_features)\n",
        "        self.drop = layers.Dropout(drop)\n",
        "    \n",
        "    def get_config(self):\n",
        "        config = super().get_config().copy()\n",
        "        config.update({\n",
        "            'fc1': self.fc1,\n",
        "            'act': self.act,\n",
        "            'fc2': self.fc2,\n",
        "            'drop': self.drop\n",
        "          })\n",
        "        return config\n",
        "\n",
        "    def call(self, x):\n",
        "        x = self.fc1(x)\n",
        "        x = self.act(x)\n",
        "        x = self.drop(x)\n",
        "        x = self.fc2(x)\n",
        "        x = self.drop(x)\n",
        "        return x\n",
        "\n",
        "\n",
        "#-- Model implementation : Hierarchical Transformer Autoencoder\n",
        "\n",
        "def transformer_autoencoder(encoded_patches, embed_shape, input_resolution):\n",
        "       \n",
        "    for iter1 in range(transformer_layers//2):\n",
        "\n",
        "      if iter1 != 0:\n",
        "        encoded_patches = PatchMerging(input_resolution, embed_shape[1])(encoded_patches)\n",
        "        input_resolution //= 2       \n",
        "        embed_shape = (input_resolution*input_resolution, embed_shape[1]*2)\n",
        "      \n",
        "      \n",
        "      for _ in range(2):\n",
        "        query = key = encoded_patches\n",
        "        #---Multi-head Attention Layer\n",
        "        attn_encoded_patches = layers.MultiHeadAttention(\n",
        "            num_heads = num_heads,\n",
        "            key_dim = embed_shape[1],\n",
        "            dropout=rate)(query=query, value=encoded_patches, key=key)\n",
        "            \n",
        "        attn_encoded_patches = layers.Dropout(rate)(attn_encoded_patches)\n",
        "        #---Skip Connection 1\n",
        "        encoded_patches += attn_encoded_patches\n",
        "        #---Layer Normalization 1\n",
        "        encoded_patches = layers.LayerNormalization(epsilon=1e-6)(encoded_patches)\n",
        "        #---Feed-Forward Network\n",
        "        ffn_out = Mlp(in_features=embed_shape[1], hidden_features=4*embed_shape[1], drop=rate)(encoded_patches)\n",
        "        #---Skip Connection \n",
        "        encoded_patches += ffn_out\n",
        "        #---Layer Normalization 2\n",
        "        encoded_patches = layers.LayerNormalization(epsilon=1e-6)(encoded_patches)\n",
        "      \n",
        "      \n",
        "    for iter2 in range(transformer_layers//2):\n",
        "\n",
        "      if iter2 != 0:\n",
        "        encoded_patches = PatchExpanding(input_resolution, embed_shape[1])(encoded_patches)\n",
        "        input_resolution *= 2       \n",
        "        embed_shape = (input_resolution*input_resolution, embed_shape[1]//2)\n",
        "      \n",
        "      \n",
        "      target = encoded_patches  \n",
        "      for _ in range(2):\n",
        "        \n",
        "        query_tgt = key_tgt = target\n",
        "            \n",
        "        #---First Multi-Head Attention\n",
        "        attn_target1 = layers.MultiHeadAttention(\n",
        "            num_heads = num_heads,\n",
        "            key_dim = embed_shape[1],\n",
        "            dropout=rate)(query=query_tgt, value=target, key=key_tgt) \n",
        "                                    \n",
        "        attn_target1 = layers.Dropout(rate)(attn_target1)\n",
        "        #---Skip Connection 1\n",
        "        target += attn_target1\n",
        "        #---Layer Normalization 1\n",
        "        target = layers.LayerNormalization(epsilon=1e-6)(target)\n",
        "\n",
        "        query_tgt = target\n",
        "            \n",
        "        #---Second Multi-Head Attention\n",
        "        attn_target2 = layers.MultiHeadAttention(\n",
        "            num_heads = num_heads,\n",
        "            key_dim = embed_shape[1],\n",
        "            dropout=rate)(query=query_tgt, value=encoded_patches, key=encoded_patches)\n",
        "                                              \n",
        "        attn_target2 = layers.Dropout(rate)(attn_target2)\n",
        "        #---Skip Connection 2\n",
        "        target += attn_target2\n",
        "        #---Layer Normalization 2\n",
        "        target = layers.LayerNormalization(epsilon=1e-6)(target)\n",
        "        #---Feed-Forward Network\n",
        "        ffn_out = Mlp(in_features=embed_shape[1], hidden_features=4*embed_shape[1], drop=rate)(encoded_patches)\n",
        "        #---Skip Connection 3\n",
        "        target +=ffn_out\n",
        "        #---Layer Normalization 3\n",
        "        target = layers.LayerNormalization(epsilon=1e-6)(target)\n",
        "      \n",
        "      encoded_patches =  target\n",
        "\n",
        "    return encoded_patches\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\"\"\"# Build the H_TAE Model\"\"\"\n",
        "\n",
        "def TAE():\n",
        "   \n",
        "  input_img = tf.keras.Input(shape=(image_size, image_size, num_channels))\n",
        "  \n",
        "  projection = layers.Conv2D(filters=filters, kernel_size=patch_size, strides=patch_size)(input_img)\n",
        "\n",
        "  projection = layers.LayerNormalization(epsilon=1e-5)(projection) #-- si without\n",
        "\n",
        "  B = tf.shape(projection)[0]\n",
        "  enc_patches = tf.reshape(projection, [B, input_resolution*input_resolution, filters])\n",
        "\n",
        "  decoded_patches = transformer_autoencoder(enc_patches, embed_shape, input_resolution)\n",
        "  \n",
        "  decoded_patches = FinalPatchExpand_X4(input_resolution, filters)(decoded_patches)\n",
        "  \n",
        "  B = tf.shape(decoded_patches)[0]\n",
        "  decoded_patches = tf.reshape(decoded_patches, [B, image_size, image_size, filters])\n",
        "  \n",
        "  reconstructed = layers.Conv2D(filters=num_channels, kernel_size=1, use_bias=False)(decoded_patches)\n",
        "\n",
        "  model = tf.keras.Model(input_img, reconstructed)\n",
        "  \n",
        "  return model\n",
        "  \n",
        "    \n",
        "#-- Configure the hyperparameters\n",
        "\n",
        "model_name = 'Hierarchical Transformer'\n",
        "numEpochs = 50\n",
        "learning_rate = 0.00001\n",
        "rate = 0.  \n",
        "image_size = 256\n",
        "num_channels = 1\n",
        "batch_size = 1\n",
        "\n",
        "\n",
        "#-- Transformer parameters variation\n",
        "\n",
        "param_grid = {\n",
        "              'transformer_layers': [8],\n",
        "              'patch_size' : [4],\n",
        "              'num_heads' : [4],\n",
        "              'filters' : [96]\n",
        "             }\n",
        "             \n",
        "            \n",
        "PARAMS = ParameterGrid(param_grid)\n",
        "list_PARAMS = list(PARAMS)\n",
        "\n",
        "\n",
        "#-- Configure training and testing Datasets \n",
        "\n",
        "saved_dir = './saved/'\n",
        "data_dir  = './data/OASIS/'\n",
        "\n",
        "test_healthy_path = './data/OASIS_all/OASIS_Flair_Test.npy'\n",
        "\n",
        "test_path = './data/BraTS/s0/BraTS_Flair.npy'\n",
        "brainmask_path = './data/BraTS/s0/BraTS_Brainmask.npy'\n",
        "x_prior_path = './data/BraTS/s0/BraTS_prior_52.npy.npy'\n",
        "label_path = './data/BraTS/s0/BraTS_GT.npy'\n",
        "\n",
        "'''\n",
        "#-- If using MSLUB as test-set\n",
        "test_path = './data/MSLUB/MSLUB_Flair.npy'\n",
        "brainmask_path = './data/MSLUB/MSLUB_Brainmask.npy'\n",
        "x_prior_path = './data/MSLUB/MSLUB_prior_57.npy'\n",
        "label_path = './data/MSLUB/MSLUB_GT.npy'\n",
        "'''\n",
        "\n",
        "list_len = [4805, 3078]    #-- BraTS and MSLUB test-set sizes, respectively.\n",
        "len_testset = list_len[0]  #-- 0 for BraTS and 1 for MSLUB\n",
        "\n",
        "train_paths = list_of_paths(data_dir)\n",
        "\n",
        "nb_train_files = 66\n",
        "data_gen = data_generator(train_paths[:nb_train_files], batch_size)\n",
        "training_steps = (256 / batch_size) * nb_train_files\n",
        "\n",
        "nb_val_files = 5\n",
        "val_gen = data_generator(train_paths[-nb_val_files:], batch_size)\n",
        "validation_steps = (256 / batch_size) * nb_val_files\n",
        "\n",
        "\n",
        "#-- Train, Test and Evaluate\n",
        "\n",
        "for param in list_PARAMS:\n",
        "   try:\n",
        "\n",
        "        #-- Checkpoints dir\n",
        "\n",
        "        date = datetime. now(). strftime(\"%Y_%m_%d-%I:%M:%S_%p\")\n",
        "        ckpts_dir = os.path.join(saved_dir, f'Ckpts_{date}')\n",
        "        os.makedirs(ckpts_dir)\n",
        "        \n",
        "        ckpts_path = os.path.join(ckpts_dir, 'Model_Ckpts.h5')\n",
        "        params_path = os.path.join(ckpts_dir, 'Parameters.txt')\n",
        "        results_path = os.path.join(ckpts_dir, 'Results.txt')\n",
        "        fig_path = os.path.join(ckpts_dir, 'History_plot.png')\n",
        "        dice_plot_path = os.path.join(ckpts_dir, 'Dice_plot.png')\n",
        "        predicted_path = os.path.join(ckpts_dir, 'Predicted.npy')\n",
        "        residual_path = os.path.join(ckpts_dir, 'Residuals.npy')\n",
        "        residual_BP_path = os.path.join(ckpts_dir, 'Residuals_BP.npy')\n",
        "\n",
        "\n",
        "        #-- Configure the parameters\n",
        "\n",
        "        transformer_layers = param['transformer_layers']\n",
        "        patch_size = param['patch_size']\n",
        "        num_heads = param['num_heads']\n",
        "        filters = param['filters']\n",
        "        input_resolution = image_size // patch_size\n",
        "        num_patches = input_resolution ** 2\n",
        "        embed_shape = (num_patches, filters)\n",
        "\n",
        "\n",
        "        #-- Configure the training\n",
        "\n",
        "        opt = tf.keras.optimizers.Adam(learning_rate = learning_rate)\n",
        "\n",
        "        calbks = tf.keras.callbacks.ModelCheckpoint(filepath=ckpts_path, monitor='val_loss', save_best_only=True, save_weights_only=True, verbose=2)\n",
        "        tqdm_callback = tfa.callbacks.TQDMProgressBar()\n",
        "\n",
        "        model = TAE()\n",
        "        model.summary()\n",
        "        model.compile(optimizer=opt, loss='mse', metrics=['mae', SSIMLoss, MS_SSIMLoss])\n",
        "\n",
        "\n",
        "        #-- Print & Write model Parameters\n",
        "\n",
        "        parameters = (f'\\nSelected model \"{model_name}\" with :\\n - {num_heads}: Heads,\\n - {patch_size}: Patch size,\\n - {transformer_layers}: Transformer Layer,\\n - ({embed_shape[0]}, {embed_shape[1]}): Embedding Shape,\\n - {batch_size}: Batche(s)\\n - {numEpochs}: Epochs\\n')\n",
        "        print(parameters)\n",
        "\n",
        "        \n",
        "        #-- Train\n",
        "\n",
        "        print('\\nTrain =>\\n')\n",
        "        history = model.fit(x = data_gen,\n",
        "                            steps_per_epoch = training_steps,\n",
        "                            validation_data = val_gen,\n",
        "                            validation_steps = validation_steps,\n",
        "                            verbose = 0,\n",
        "                            epochs = numEpochs,\n",
        "                            callbacks = [calbks, tqdm_callback]\n",
        "                            )\n",
        "        \n",
        "\n",
        "        #-- Get training and test loss historie\n",
        "        plot_history(history, path=fig_path)\n",
        "        plt.close()\n",
        "        time.sleep(2)               \n",
        "        \n",
        "        \n",
        "        #-- Test \n",
        "\n",
        "        print('\\nTest ===>\\n')\n",
        "        my_test = np.load(test_path)\n",
        "        brainmask = np.load(brainmask_path)\n",
        "        x_prior = np.load(x_prior_path)\n",
        "        my_labels = np.load(label_path)              \n",
        "        \n",
        "        \n",
        "        healthy_test = np.load(test_healthy_path)\n",
        "        steps = healthy_test.shape[0]\n",
        "                \n",
        "        score = model.evaluate(x=healthy_test, y=healthy_test, verbose=0, steps=steps, callbacks = [tqdm_callback])\n",
        "        score_out = (f'\\nTest on healthy unseen data :\\n - Test loss (MSE): {score[0]},\\n - Test MAE : {score[1]},\\n - Test SSIM: {score[2]},\\n - Test MS_SSIMLoss: {score[3]},\\n')\n",
        "        print(score_out)\n",
        "        \n",
        "                \n",
        "        #-- Predict\n",
        "\n",
        "        print('\\nPredict =====>\\n')\n",
        "        predicted = model.predict(x=my_test, verbose=1, steps=len_testset)\n",
        "        np.save(predicted_path, predicted)\n",
        "        time.sleep(4)\n",
        "        \n",
        "        \n",
        "        #-- Calculate, Post-process and Save Residuals\n",
        "\n",
        "        print('\\nCalculate, Post-process and Save Residuals =====>\\n')     \n",
        "        residual_BP = calculate_residual_BP(my_test, predicted, brainmask)  #-- You can use either brainmask or x_prior\n",
        "        np.save(residual_BP_path, residual_BP)\n",
        "        \n",
        "        residual = calculate_residual(my_test, predicted, brainmask)  #-- You can use either brainmask or x_prior\n",
        "        np.save(residual_path, residual)\n",
        "        \n",
        "\n",
        "        #-- Evaluation\n",
        "\n",
        "        print('\\nEvaluate =========>\\n')        \n",
        "        [AUROC, AUPRC, AVG_DICE, MAD, STD], DICE = eval_residuals(my_labels, residual)     \n",
        "        results = (f'\\nResults after median_filter :\\n - AUROC = {AUROC}\\n - AUPRC = {AUPRC}\\n - AVG_DICE = {AVG_DICE}\\n - MEDIAN_ABSOLUTE_DEVIATION = {MAD}\\n - STANDARD_DEVIATION = {STD}')\n",
        "        print(results)\n",
        "                     \n",
        "        plt.figure()\n",
        "        hor_axis = [x for x in range(len_testset)]\n",
        "        plt.scatter(hor_axis, DICE, s = 5, marker = '.', c = 'blue')\n",
        "        plt.ylabel('Dice Score')\n",
        "        plt.xlabel('N° Samples')\n",
        "        plt.title('Dice scores')\n",
        "        plt.savefig(dice_plot_path)\n",
        "        time.sleep(2)\n",
        "\n",
        "\n",
        "        #-- Save\n",
        "\n",
        "        print('\\nSave Results and Parameters =============>\\n')\n",
        "        f = open(results_path, \"w\")\n",
        "        f.write(results)       \n",
        "        f.close()   \n",
        "                \n",
        "                \n",
        "        f = open(params_path, \"w\")\n",
        "        f.write(parameters)\n",
        "        f.write(score_out)\n",
        "        f.close()\n",
        "        \n",
        "        \n",
        "        #-- End\n",
        "\n",
        "        print('\\nEnd !\\n')\n",
        "\n",
        "\n",
        "   except:\n",
        "        print('Error encountered !')\n",
        "        tf.keras.backend.clear_session()\n",
        "        continue"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rx-OXyQG7eDv"
      },
      "source": [
        "## H_TAE_S"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vVrCY9DF7gMm"
      },
      "source": [
        "#-- Dense Layer   \n",
        "\n",
        "class TruncatedDense(layers.Dense):\n",
        "    def __init__(self, units, use_bias=True, initializer = tf.keras.initializers.TruncatedNormal(mean=0., stddev=.02)):\n",
        "        super().__init__(units, use_bias=use_bias, kernel_initializer=initializer)\n",
        "\n",
        "\n",
        "#-- Patch Merging Layer  \n",
        "\n",
        "class PatchMerging(layers.Layer):\n",
        "\n",
        "    def __init__(self, input_resolution, dim):\n",
        "        super().__init__()\n",
        "        self.input_resolution = input_resolution\n",
        "        self.dim = dim    #--- refer to the projection_dim (nC)\n",
        "        self.reduction = TruncatedDense(2 * dim, use_bias=False)\n",
        "        self.norm = layers.LayerNormalization(epsilon=1e-5)\n",
        "    \n",
        "    def get_config(self):\n",
        "        config = super().get_config().copy()\n",
        "        config.update({\n",
        "            'input_resolution': self.input_resolution,\n",
        "            'dim': self.dim,\n",
        "            'reduction': self.reduction,\n",
        "            'norm': self.norm\n",
        "          })\n",
        "        return config    \n",
        "\n",
        "    def call(self, x):\n",
        "        \"\"\"\n",
        "        x: B, H*W, C\n",
        "        \"\"\"\n",
        "        H = W = self.input_resolution\n",
        "        B = tf.shape(x)[0]\n",
        "\n",
        "        x = tf.reshape(x, [B, H, W, self.dim])\n",
        "        x0 = x[:, 0::2, 0::2, :]  # B H/2 W/2 C\n",
        "        x1 = x[:, 1::2, 0::2, :]  # B H/2 W/2 C\n",
        "        x2 = x[:, 0::2, 1::2, :]  # B H/2 W/2 C\n",
        "        x3 = x[:, 1::2, 1::2, :]  # B H/2 W/2 C\n",
        "        x = tf.concat([x0, x1, x2, x3], -1)  # B H/2 W/2 4*C\n",
        "        x = tf.reshape(x, [B, (H//2)*(W//2), 4 * self.dim])  # B H/2*W/2 4*C\n",
        "        x = self.norm(x)\n",
        "        x = self.reduction(x)\n",
        "\n",
        "        return x\n",
        "\n",
        "\n",
        "#-- Patch Expanding Layer \n",
        "\n",
        "class PatchExpanding(layers.Layer):\n",
        "\n",
        "    def __init__(self, input_resolution, dim):\n",
        "        super().__init__()\n",
        "        self.input_resolution = input_resolution\n",
        "        self.dim = dim\n",
        "        self.expand = TruncatedDense(2 * dim, use_bias=False)\n",
        "        self.norm = layers.LayerNormalization(epsilon=1e-5)\n",
        "    \n",
        "    def get_config(self):\n",
        "        config = super().get_config().copy()\n",
        "        config.update({\n",
        "            'input_resolution': self.input_resolution,\n",
        "            'dim': self.dim,\n",
        "            'expand': self.expand,\n",
        "            'norm': self.norm\n",
        "          })\n",
        "        return config\n",
        "\n",
        "    def call(self, x):\n",
        "        \"\"\"\n",
        "        x: B, H*W, C\n",
        "        \"\"\"\n",
        "        H = W = self.input_resolution\n",
        "        x = self.expand(x)\n",
        "        B = tf.shape(x)[0]\n",
        "\n",
        "        x = tf.reshape(x, [B, H, W, self.dim * 2])\n",
        "        x = rearrange(x, 'b h w (p1 p2 c)-> b (h p1) (w p2) c', p1=2, p2=2, c=self.dim // 2)\n",
        "\n",
        "        x = tf.reshape(x, [B, (H*2)*(W*2), self.dim//2])\n",
        "        x = self.norm(x)\n",
        "\n",
        "        return x\n",
        "\n",
        "\n",
        "#-- Final Patch Expanding Layer \n",
        "\n",
        "class FinalPatchExpand_X4(layers.Layer):\n",
        "\n",
        "    def __init__(self, input_resolution, dim):\n",
        "        super().__init__()\n",
        "        self.input_resolution = input_resolution\n",
        "        self.dim = dim\n",
        "        self.expand = TruncatedDense(16 * dim, use_bias=False)\n",
        "        self.norm = layers.LayerNormalization(epsilon=1e-5)\n",
        "    \n",
        "    def get_config(self):\n",
        "        config = super().get_config().copy()\n",
        "        config.update({\n",
        "            'input_resolution': self.input_resolution,\n",
        "            'dim': self.dim,\n",
        "            'expand': self.expand,\n",
        "            'norm': self.norm\n",
        "          })\n",
        "        return config\n",
        "\n",
        "    def call(self, x):\n",
        "        \"\"\"\n",
        "        x: B, H*W, C\n",
        "        \"\"\"\n",
        "        H = W = self.input_resolution\n",
        "        x = self.expand(x)\n",
        "        B = tf.shape(x)[0]\n",
        "\n",
        "        x = tf.reshape(x, [B, H, W, self.dim * 16])\n",
        "        x = rearrange(x, 'b h w (p1 p2 c)-> b (h p1) (w p2) c', p1=4, p2=4, c=self.dim)\n",
        "        x = tf.reshape(x, [B, (H*4)*(W*4), self.dim])\n",
        "\n",
        "        x = self.norm(x)\n",
        "\n",
        "        return x\n",
        "\n",
        "\n",
        "#-- Mlp Head\n",
        "\n",
        "class Mlp(layers.Layer):\n",
        "\n",
        "    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=layers.Activation(tf.nn.gelu), drop=0.):\n",
        "        super().__init__()\n",
        "        out_features = out_features or in_features\n",
        "        hidden_features = hidden_features or in_features\n",
        "        self.fc1 = TruncatedDense(hidden_features)\n",
        "        self.act = act_layer\n",
        "        self.fc2 = TruncatedDense(out_features)\n",
        "        self.drop = layers.Dropout(drop)\n",
        "    \n",
        "    def get_config(self):\n",
        "        config = super().get_config().copy()\n",
        "        config.update({\n",
        "            'fc1': self.fc1,\n",
        "            'act': self.act,\n",
        "            'fc2': self.fc2,\n",
        "            'drop': self.drop\n",
        "          })\n",
        "        return config\n",
        "\n",
        "    def call(self, x):\n",
        "        x = self.fc1(x)\n",
        "        x = self.act(x)\n",
        "        x = self.drop(x)\n",
        "        x = self.fc2(x)\n",
        "        x = self.drop(x)\n",
        "        return x\n",
        "\n",
        "\n",
        "#-- Model implementation : Hierarchical Transformer Autoencoder with Skip connections\n",
        "\n",
        "def transformer_autoencoder(encoded_patches, embed_shape, input_resolution):\n",
        "    \n",
        "    encoded_list = []\n",
        "    \n",
        "    for iter1 in range(transformer_layers//2):\n",
        "\n",
        "      if iter1 != 0:\n",
        "        encoded_patches = PatchMerging(input_resolution, embed_shape[1])(encoded_patches)\n",
        "        input_resolution //= 2       \n",
        "        embed_shape = (input_resolution*input_resolution, embed_shape[1]*2)\n",
        "\n",
        "      for _ in range(2):\n",
        "        query = key = encoded_patches\n",
        "        #---Multi-head Attention Layer\n",
        "        attn_encoded_patches = layers.MultiHeadAttention(\n",
        "            num_heads = num_heads,\n",
        "            key_dim = embed_shape[1],\n",
        "            dropout=rate)(query=query, value=encoded_patches, key=key)\n",
        "            \n",
        "        attn_encoded_patches = layers.Dropout(rate)(attn_encoded_patches)\n",
        "        #---Skip Connection 1\n",
        "        encoded_patches += attn_encoded_patches\n",
        "        #---Layer Normalization 1\n",
        "        encoded_patches = layers.LayerNormalization(epsilon=1e-6)(encoded_patches)\n",
        "        #---Feed-Forward Network\n",
        "        ffn_out = Mlp(in_features=embed_shape[1], hidden_features=4*embed_shape[1], drop=rate)(encoded_patches)\n",
        "        #---Skip Connection \n",
        "        encoded_patches += ffn_out\n",
        "        #---Layer Normalization 2\n",
        "        encoded_patches = layers.LayerNormalization(epsilon=1e-6)(encoded_patches)\n",
        "      \n",
        "      encoded_list += [encoded_patches]\n",
        "      \n",
        "      \n",
        "    for iter2 in range(transformer_layers//2):\n",
        "\n",
        "      if iter2 != 0:\n",
        "        \n",
        "        encoded_patches = PatchExpanding(input_resolution, embed_shape[1])(encoded_patches)\n",
        "        input_resolution *= 2       \n",
        "        embed_shape = (input_resolution*input_resolution, embed_shape[1]//2)\n",
        "        \n",
        "        encoded_patches += encoded_list[3 - iter2]\n",
        "        encoded_patches = layers.LayerNormalization(epsilon=1e-6)(encoded_patches)\n",
        "      \n",
        "      target = encoded_patches  \n",
        "      \n",
        "      for _ in range(2):\n",
        "        \n",
        "        query_tgt = key_tgt = target\n",
        "            \n",
        "        #---First Multi-Head Attention\n",
        "        attn_target1 = layers.MultiHeadAttention(\n",
        "            num_heads = num_heads,\n",
        "            key_dim = embed_shape[1],\n",
        "            dropout=rate)(query=query_tgt, value=target, key=key_tgt) \n",
        "                                    \n",
        "        attn_target1 = layers.Dropout(rate)(attn_target1)\n",
        "        #---Skip Connection 1\n",
        "        target += attn_target1\n",
        "        #---Layer Normalization 1\n",
        "        target = layers.LayerNormalization(epsilon=1e-6)(target)\n",
        "\n",
        "        query_tgt = target\n",
        "            \n",
        "        #---Second Multi-Head Attention\n",
        "        attn_target2 = layers.MultiHeadAttention(\n",
        "            num_heads = num_heads,\n",
        "            key_dim = embed_shape[1],\n",
        "            dropout=rate)(query=query_tgt, value=encoded_patches, key=encoded_patches)\n",
        "                                              \n",
        "        attn_target2 = layers.Dropout(rate)(attn_target2)\n",
        "        #---Skip Connection 2\n",
        "        target += attn_target2\n",
        "        #---Layer Normalization 2\n",
        "        target = layers.LayerNormalization(epsilon=1e-6)(target)\n",
        "        #---Feed-Forward Network\n",
        "        ffn_out = Mlp(in_features=embed_shape[1], hidden_features=4*embed_shape[1], drop=rate)(encoded_patches)\n",
        "        #---Skip Connection 3\n",
        "        target +=ffn_out\n",
        "        #---Layer Normalization 3\n",
        "        target = layers.LayerNormalization(epsilon=1e-6)(target)\n",
        "      \n",
        "      encoded_patches =  target \n",
        "\n",
        "    return encoded_patches\n",
        "\n",
        "\n",
        "#-- Build the H_TAE_S Model\n",
        "\n",
        "def TAE():\n",
        "   \n",
        "  input_img = tf.keras.Input(shape=(image_size, image_size, num_channels))\n",
        "  \n",
        "  projection = layers.Conv2D(filters=filters, kernel_size=patch_size, strides=patch_size)(input_img)\n",
        "  projection = layers.LayerNormalization(epsilon=1e-5)(projection) #-- si without\n",
        "  B = tf.shape(projection)[0]\n",
        "  enc_patches = tf.reshape(projection, [B, input_resolution*input_resolution, filters])\n",
        "  decoded_patches = transformer_autoencoder(enc_patches, embed_shape, input_resolution) \n",
        "  decoded_patches = FinalPatchExpand_X4(input_resolution, filters)(decoded_patches) \n",
        "  B = tf.shape(decoded_patches)[0]\n",
        "  decoded_patches = tf.reshape(decoded_patches, [B, image_size, image_size, filters]) \n",
        "  reconstructed = layers.Conv2D(filters=num_channels, kernel_size=1, use_bias=False)(decoded_patches)\n",
        "\n",
        "  model = tf.keras.Model(input_img, reconstructed)\n",
        "  \n",
        "  return model\n",
        "  \n",
        "    \n",
        "#-- Configure the hyperparameters\n",
        "\n",
        "model_name = 'Hierarchical Transformer with Skip connections'\n",
        "numEpochs = 50\n",
        "learning_rate = 0.00001\n",
        "rate = 0.  \n",
        "image_size = 256\n",
        "num_channels = 1\n",
        "batch_size = 1\n",
        "\n",
        "\n",
        "#-- Transformer parameters variation\n",
        "\n",
        "param_grid = {\n",
        "              'transformer_layers': [8],\n",
        "              'patch_size' : [4],\n",
        "              'num_heads' : [8],\n",
        "              'filters' : [96]\n",
        "             }\n",
        "             \n",
        "            \n",
        "PARAMS = ParameterGrid(param_grid)\n",
        "list_PARAMS = list(PARAMS)\n",
        "\n",
        "\n",
        "#-- Configure training and testing Datasets \n",
        "\n",
        "saved_dir = './saved/'\n",
        "data_dir  = './data/OASIS/'\n",
        "\n",
        "test_healthy_path = './data/OASIS_all/OASIS_Flair_Test.npy'\n",
        "\n",
        "test_path = './data/BraTS/s0/BraTS_Flair.npy'\n",
        "brainmask_path = './data/BraTS/s0/BraTS_Brainmask.npy'\n",
        "x_prior_path = './data/BraTS/s0/BraTS_prior_52.npy.npy'\n",
        "label_path = './data/BraTS/s0/BraTS_GT.npy'\n",
        "\n",
        "'''\n",
        "#-- If using MSLUB as test-set\n",
        "test_path = './data/MSLUB/MSLUB_Flair.npy'\n",
        "brainmask_path = './data/MSLUB/MSLUB_Brainmask.npy'\n",
        "x_prior_path = './data/MSLUB/MSLUB_prior_57.npy'\n",
        "label_path = './data/MSLUB/MSLUB_GT.npy'\n",
        "'''\n",
        "\n",
        "list_len = [4805, 3078]    #-- BraTS and MSLUB test-set sizes, respectively.\n",
        "len_testset = list_len[0]  #-- 0 for BraTS and 1 for MSLUB\n",
        "\n",
        "train_paths = list_of_paths(data_dir)\n",
        "\n",
        "nb_train_files = 66\n",
        "data_gen = data_generator(train_paths[:nb_train_files], batch_size)\n",
        "training_steps = (256 / batch_size) * nb_train_files\n",
        "\n",
        "nb_val_files = 5\n",
        "val_gen = data_generator(train_paths[-nb_val_files:], batch_size)\n",
        "validation_steps = (256 / batch_size) * nb_val_files\n",
        "\n",
        "\n",
        "#-- Train, Test and Evaluate\n",
        "\n",
        "for param in list_PARAMS:\n",
        "   try:\n",
        "\n",
        "        #-- Checkpoints dir\n",
        "\n",
        "        date = datetime. now(). strftime(\"%Y_%m_%d-%I:%M:%S_%p\")\n",
        "        ckpts_dir = os.path.join(saved_dir, f'Ckpts_{date}')\n",
        "        os.makedirs(ckpts_dir)\n",
        "\n",
        "        ckpts_path = os.path.join(ckpts_dir, 'Model_Ckpts.h5')\n",
        "        params_path = os.path.join(ckpts_dir, 'Parameters.txt')\n",
        "        results_path = os.path.join(ckpts_dir, 'Results.txt')\n",
        "        fig_path = os.path.join(ckpts_dir, 'History_plot.png')\n",
        "        dice_plot_path = os.path.join(ckpts_dir, 'Dice_plot.png')\n",
        "        predicted_path = os.path.join(ckpts_dir, 'Predicted.npy')\n",
        "        residual_path = os.path.join(ckpts_dir, 'Residuals.npy')\n",
        "        residual_BP_path = os.path.join(ckpts_dir, 'Residuals_BP.npy')\n",
        "\n",
        "\n",
        "        #-- Configure the parameters\n",
        "\n",
        "        transformer_layers = param['transformer_layers']\n",
        "        patch_size = param['patch_size']\n",
        "        num_heads = param['num_heads']\n",
        "        filters = param['filters']\n",
        "        input_resolution = image_size // patch_size\n",
        "        num_patches = input_resolution ** 2\n",
        "        embed_shape = (num_patches, filters)\n",
        "\n",
        "\n",
        "        #-- Configure the training\n",
        "\n",
        "        opt = tf.keras.optimizers.Adam(learning_rate = learning_rate)\n",
        "\n",
        "        calbks = tf.keras.callbacks.ModelCheckpoint(filepath=ckpts_path, monitor='val_loss', save_best_only=True, save_weights_only=True, verbose=2)\n",
        "        tqdm_callback = tfa.callbacks.TQDMProgressBar()\n",
        "\n",
        "        model = TAE()\n",
        "        model.summary()\n",
        "        model.compile(optimizer=opt, loss='mse', metrics=['mae', SSIMLoss, MS_SSIMLoss])\n",
        "\n",
        "\n",
        "        #-- Print & Write model Parameters\n",
        "\n",
        "        parameters = (f'\\nSelected model \"{model_name}\" with :\\n - {num_heads}: Heads,\\n - {patch_size}: Patch size,\\n - {transformer_layers}: Transformer Layer,\\n - ({embed_shape[0]}, {embed_shape[1]}): Embedding Shape,\\n - {batch_size}: Batche(s)\\n - {numEpochs}: Epochs\\n')\n",
        "        print(parameters)\n",
        "\n",
        "        \n",
        "        #-- Train\n",
        "\n",
        "        print('\\nTrain =>\\n')\n",
        "        history = model.fit(x = data_gen,\n",
        "                            steps_per_epoch = training_steps,\n",
        "                            validation_data = val_gen,\n",
        "                            validation_steps = validation_steps,\n",
        "                            verbose = 0,\n",
        "                            epochs = numEpochs,\n",
        "                            callbacks = [calbks, tqdm_callback]\n",
        "                            )\n",
        "        \n",
        "\n",
        "        #-- Get training and test loss history\n",
        "\n",
        "        plot_history(history, path=fig_path)\n",
        "        plt.close()\n",
        "        time.sleep(2)               \n",
        "        \n",
        "        \n",
        "        #-- Test  \n",
        "\n",
        "        print('\\nTest ===>\\n')\n",
        "        my_test = np.load(test_path)\n",
        "        brainmask = np.load(brainmask_path)\n",
        "        x_prior = np.load(x_prior_path)\n",
        "        my_labels = np.load(label_path)              \n",
        "             \n",
        "        healthy_test = np.load(test_healthy_path)\n",
        "        steps = healthy_test.shape[0]\n",
        "                \n",
        "        score = model.evaluate(x=healthy_test, y=healthy_test, verbose=0, steps=steps, callbacks = [tqdm_callback])\n",
        "        score_out = (f'\\nTest on healthy unseen data :\\n - Test loss (MSE): {score[0]},\\n - Test MAE : {score[1]},\\n - Test SSIM: {score[2]},\\n - Test MS_SSIMLoss: {score[3]},\\n')\n",
        "        print(score_out)\n",
        "        \n",
        "               \n",
        "        #-- Predict\n",
        "\n",
        "        print('\\nPredict =====>\\n')\n",
        "        predicted = model.predict(x=my_test, verbose=1, steps=len_testset)\n",
        "        np.save(predicted_path, predicted)\n",
        "        time.sleep(4)\n",
        "\n",
        "\n",
        "        #-- Calculate, Post-process and Save Residuals\n",
        "\n",
        "        print('\\nCalculate, Post-process and Save Residuals =====>\\n')     \n",
        "        residual_BP = calculate_residual_BP(my_test, predicted, brainmask)  #-- You can use either brainmask or x_prior\n",
        "        np.save(residual_BP_path, residual_BP)\n",
        "        \n",
        "        residual = calculate_residual(my_test, predicted, brainmask)  #-- You can use either brainmask or x_prior\n",
        "        np.save(residual_path, residual)\n",
        "        \n",
        "\n",
        "        #-- Evaluation\n",
        "\n",
        "        print('\\nEvaluate =========>\\n')        \n",
        "        [AUROC, AUPRC, AVG_DICE, MAD, STD], DICE = eval_residuals(my_labels, residual)     \n",
        "        results = (f'\\nResults after median_filter :\\n - AUROC = {AUROC}\\n - AUPRC = {AUPRC}\\n - AVG_DICE = {AVG_DICE}\\n - MEDIAN_ABSOLUTE_DEVIATION = {MAD}\\n - STANDARD_DEVIATION = {STD}')\n",
        "        print(results)\n",
        "                     \n",
        "        plt.figure()\n",
        "        hor_axis = [x for x in range(len_testset)]\n",
        "        plt.scatter(hor_axis, DICE, s = 5, marker = '.', c = 'blue')\n",
        "        plt.ylabel('Dice Score')\n",
        "        plt.xlabel('N° Samples')\n",
        "        plt.title('Dice scores')\n",
        "        plt.savefig(dice_plot_path)\n",
        "        time.sleep(2)\n",
        "\n",
        "\n",
        "        #-- Save\n",
        "\n",
        "        print('\\nSave Results and Parameters =============>\\n')\n",
        "        f = open(results_path, \"w\")\n",
        "        f.write(results)       \n",
        "        f.close() \n",
        "                \n",
        "        f = open(params_path, \"w\")\n",
        "        f.write(parameters)\n",
        "        f.write(score_out)\n",
        "        f.close()\n",
        "        \n",
        "        \n",
        "        #-- End\n",
        "\n",
        "        print('\\nEnd !\\n')\n",
        "\n",
        "\n",
        "   except:\n",
        "        print('Error encountered !')\n",
        "        tf.keras.backend.clear_session()\n",
        "        continue"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}