White_Matter_Hyperintensities_Segmentation.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "White Matter Hyperintensities Segmentation.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "10MH0TkmrTXG",
        "m-e4SiGnAgLv",
        "HtkuJtLwI6w8",
        "fKnJNBsumqq1",
        "xi74WxiMcC7O"
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "rmK4pfA5ePsM"
      },
      "source": [
        "import re\n",
        "import os\n",
        "import sys\n",
        "import shutil\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import nibabel as nib\n",
        "import tensorflow as tf\n",
        "import albumentations as A\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "from Models.models import Prob_Unet, AttXnet, Xnet, Unet\n",
        "\n",
        "from glob import glob\n",
        "from keras import backend as K\n",
        "from sklearn.model_selection import KFold\n",
        "from tensorflow.keras.optimizers import Adam\n",
        "from tensorflow.keras.callbacks import ModelCheckpoint, TensorBoard, ReduceLROnPlateau"
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "10MH0TkmrTXG"
      },
      "source": [
        "# 1. Helpers"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kNkCREt8ra6-"
      },
      "source": [
        "def read_data(file_path, start=0, end=0, step=1):\n",
        "    img = nib.load(file_path)\n",
        "    data = img.get_fdata()\n",
        "    end = data.shape[-1] if end == 0 else end\n",
        "    reverse_channel = [data[:,:,i] for i in range(start,end,step)]\n",
        "    return np.array(reverse_channel, dtype='float32')\n",
        "\n",
        "def adjust_data(img, cf=1, mask=False):\n",
        "    if mask is False:\n",
        "        # Normalized data \n",
        "        img = (img - np.mean(img))/(np.std(img) + K.epsilon())\n",
        "\n",
        "    img = np.expand_dims(img, -1)\n",
        "    # Central crop if needed\n",
        "    if cf<1: img = tf.image.central_crop(img, cf)\n",
        "    return img\n",
        "\n",
        "def check_data(img, seg):\n",
        "    # Select data that have a value (image that contain only background)\n",
        "    idx = [k for k,i in enumerate(img) if len(np.unique(i)) > 1] \n",
        "    return img[idx], seg[idx]\n",
        "\n",
        "def data_augmentation(images, masks):\n",
        "    transform = A.Compose([A.HorizontalFlip(), A.Rotate(p=0.8)])\n",
        "    for k,(i,j) in enumerate(zip(images,masks)):\n",
        "        transformed = transform(image=i, mask=j)\n",
        "        images[k] = transformed['image']\n",
        "        masks[k] = transformed['mask']\n",
        "\n",
        "def find_best(path_model, type_='min'):\n",
        "    # Find best weights in filename\n",
        "    path = glob(path_model+'/*.hdf5') \n",
        "    split = [i.split('-')[-1] for i in path]\n",
        "    split.sort(key=natural_keys)\n",
        "\n",
        "    if type_ == 'min':\n",
        "        file_ = split[0]\n",
        "    else:\n",
        "        file_ = split[-1]\n",
        "    return [i for i in path if file_ in i] [0]\n",
        "\n",
        "def show_img(img, y_true, y_pred, idx=0):\n",
        "    # Visualize image, ground truth, and prediction by slice\n",
        "    plt.figure(figsize=(15,15))\n",
        "    plt.subplot(1,3,1); plt.title('FLAIR image')\n",
        "    plt.imshow(img[idx,:,:,0], cmap='gray'); plt.axis('off')\n",
        "    plt.subplot(1,3,2); plt.title('Ground Truth')\n",
        "    plt.imshow(y_true[idx,:,:,0], cmap='gray'); plt.axis('off')\n",
        "    plt.subplot(1,3,3); plt.title('Prediction')\n",
        "    plt.imshow(y_pred[idx,:,:,0], cmap='gray'); plt.axis('off')\n",
        "    plt.show()\n",
        "\n",
        "def save_wmh(y_pred, file_in, file_out, name='ADNI'):\n",
        "    # Save WMHs predictions \n",
        "    y_true = nib.load(file_in)\n",
        "    wmh = np.zeros(y_true.shape)\n",
        "\n",
        "    if name == 'Singapore':\n",
        "        if wmh.shape[1] != y_pred.shape[2]:\n",
        "            wmh_temp = tf.image.rot90(y_pred, k=3).numpy()\n",
        "            for i in range(wmh.shape[-1]):\n",
        "                wmh[:,12:220,i] = wmh_temp[i,:,:,0]\n",
        "        else:\n",
        "            for i in range(wmh.shape[-1]):\n",
        "                wmh[12:220,:,i] = y_pred[i,:,:,0]\n",
        "\n",
        "    elif name == 'GE3T':\n",
        "        for i in range(wmh.shape[-1]):\n",
        "            wmh[:128,:,i] = y_pred[i,:,:,0]\n",
        "    else:\n",
        "        for i in range(wmh.shape[-1]):\n",
        "            wmh[:,:,i] = y_pred[i,:,:,0]\n",
        "\n",
        "    wmh = nib.Nifti1Image(wmh, y_true.affine, y_true.header)\n",
        "    nib.save(wmh, file_out)\n",
        "\n",
        "def save_wmh_challenge(y_pred, file_in, file_out, name):\n",
        "    # Save WMHs prediction only for Challenge dataset as full\n",
        "    y_true = nib.load(file_in)\n",
        "    wmh = np.zeros(y_true.shape)\n",
        "\n",
        "    if name == 'Singapore': start=4; end=236\n",
        "    elif name == 'GE3T': start=54; end=186\n",
        "    elif name == 'Utrecht': start=8; end=248\n",
        "    else: start=0; end=240\n",
        "\n",
        "    if name == 'Singapore':\n",
        "        if wmh.shape[1] != y_pred.shape[2]:\n",
        "            wmh_temp = tf.image.rot90(y_pred, k=3).numpy()\n",
        "            for i in range(wmh.shape[-1]):\n",
        "                wmh[:,:,i] = wmh_temp[i,:,start:end,0]\n",
        "        else:\n",
        "            for i in range(wmh.shape[-1]):\n",
        "                wmh[:,:,i] = y_pred[i,start:end,:,0]\n",
        "\n",
        "    elif name == 'Utrecht':\n",
        "        for i in range(wmh.shape[-1]):\n",
        "            wmh[:,:,i] = y_pred[i,:,start:end,0]\n",
        "    else:\n",
        "        for i in range(wmh.shape[-1]):\n",
        "            wmh[:,:,i] = y_pred[i,start:end,:,0]\n",
        "\n",
        "    wmh = nib.Nifti1Image(wmh, y_true.affine, y_true.header)\n",
        "    nib.save(wmh, file_out)\n",
        "\n",
        "# https://stackoverflow.com/questions/5967500/how-to-correctly-sort-a-string-with-a-number-inside\n",
        "def atoi(text):\n",
        "    return int(text) if text.isdigit() else text\n",
        "\n",
        "# https://stackoverflow.com/questions/5967500/how-to-correctly-sort-a-string-with-a-number-inside\n",
        "def natural_keys(text):\n",
        "    return [atoi(c) for c in re.split(r'(\\d+)', text)]\n",
        "\n",
        "# https://github.com/MrGiovanni/UNetPlusPlus/blob/master/keras/helper_functions.py#L37-L42\n",
        "def dice_coef(y_true, y_pred):\n",
        "    smooth = 1.\n",
        "    y_true_f = K.flatten(y_true)\n",
        "    y_pred_f = K.flatten(y_pred)\n",
        "    intersection = K.sum(y_true_f * y_pred_f)\n",
        "    return (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)\n",
        "\n",
        "# https://github.com/umbertogriffo/focal-loss-keras/blob/master/src/loss_function/losses.py#L11-L53\n",
        "def binary_focal_loss(gamma=2., alpha=.25):\n",
        "    def binary_focal_loss_fixed(y_true, y_pred):\n",
        "        \"\"\"\n",
        "        :param y_true: A tensor of the same shape as `y_pred`\n",
        "        :param y_pred:  A tensor resulting from a sigmoid\n",
        "        :return: Output tensor.\n",
        "        \"\"\"\n",
        "        y_true = tf.cast(y_true, tf.float32)\n",
        "        # Define epsilon so that the back-propagation will not result in NaN for 0 divisor case\n",
        "        epsilon = K.epsilon()\n",
        "        # Add the epsilon to prediction value\n",
        "        # y_pred = y_pred + epsilon\n",
        "        # Clip the prediciton value\n",
        "        y_pred = K.clip(y_pred, epsilon, 1.0 - epsilon)\n",
        "        # Calculate p_t\n",
        "        p_t = tf.where(K.equal(y_true, 1), y_pred, 1 - y_pred)\n",
        "        # Calculate alpha_t\n",
        "        alpha_factor = K.ones_like(y_true) * alpha\n",
        "        alpha_t = tf.where(K.equal(y_true, 1), alpha_factor, 1 - alpha_factor)\n",
        "        # Calculate cross entropy\n",
        "        cross_entropy = -K.log(p_t)\n",
        "        weight = alpha_t * K.pow((1 - p_t), gamma)\n",
        "        # Calculate focal loss\n",
        "        loss = weight * cross_entropy\n",
        "        # Sum the losses in mini_batch\n",
        "        loss = K.mean(K.sum(loss, axis=1))\n",
        "        return loss\n",
        "    return binary_focal_loss_fixed\n",
        "\n",
        "# https://github.com/baumgach/PHiSeg-code/blob/c43f3b32e1f434aecba936ff994b6f743ba7a5f8/utils.py#L326-L370\n",
        "def ambiguity_map(y_gen, seg=None):\n",
        "    def pixel_wise_xent(m_samp, m_gt, eps=1e-8):\n",
        "        log_samples = np.log(m_samp + eps)\n",
        "        return -1.0*np.sum(m_gt*log_samples, axis=-1)\n",
        "\n",
        "    y_pred = np.average(y_gen, axis=0)\n",
        "    E_arr = np.zeros(y_gen.shape)\n",
        "    for i in range(y_gen.shape[0]):\n",
        "        for j in range(y_gen.shape[1]):\n",
        "            if seg is None:\n",
        "                E_arr[i,j,...] = np.expand_dims(pixel_wise_xent(y_gen[i,j,...], y_pred[j,...]), axis=-1)\n",
        "            else:\n",
        "                E_arr[i,j,...] = np.expand_dims(pixel_wise_xent(y_gen[i,j,...], seg[j,...]), axis=-1)\n",
        "\n",
        "    return np.average(E_arr, axis=0)"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m-e4SiGnAgLv"
      },
      "source": [
        "# 2. Datasets"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DM4oLm2ZOo61"
      },
      "source": [
        "def get_dataset(name, end=0):\n",
        "    if name == 'ADNI' and end >= 1:\n",
        "        start = 5; end = 30; step=1\n",
        "    elif name == 'Singapore' and end >= 1:\n",
        "        start = 0; end = 45; step=1\n",
        "    elif name == 'GE3T' and end >= 1:\n",
        "        start = 25; end = 70; step=1\n",
        "    elif name == 'Utrecht' and end >= 1:\n",
        "        start = 0; end = 45; step=1\n",
        "    else:\n",
        "        start = 0; end = 0; step=1\n",
        "\n",
        "    if name == 'Challenge':\n",
        "        fold_img, fold_seg = read_challenge(start, end, step=1, cf=1)\n",
        "    else:\n",
        "        fold_img, fold_seg = read_dataset(name, start, end, step=step, cf=1)\n",
        "    \n",
        "    return fold_img, fold_seg\n",
        "\n",
        "def read_dataset(name, start, end, step=1, cf=1):\n",
        "    if name == 'ADNI':\n",
        "        flair_path = glob(path_dataset+'ADNI/*/*/*brain.nii.gz')\n",
        "        lesion_path = glob(path_dataset+'ADNI/*/*/*wmh.nii.gz')\n",
        "    else:\n",
        "        flair_path = glob(path_dataset+'Challenge/{}/*/*brain.nii.gz'.format(name))\n",
        "        lesion_path = glob(path_dataset+'Challenge/{}/*/*wmh.nii.gz'.format(name))\n",
        "\n",
        "    flair_path.sort(key=natural_keys) \n",
        "    lesion_path.sort(key=natural_keys)\n",
        "    \n",
        "    if name == 'ADNI':\n",
        "        flair_path = [[i for i in flair_path if i.split('/')[-3] == 'fold'+str(j)] for j in range(1,5)]\n",
        "        lesion_path = [[i for i in lesion_path if i.split('/')[-3] == 'fold'+str(j)] for j in range(1,5)]\n",
        "    else:\n",
        "        flair_path = [flair_path]\n",
        "        lesion_path = [lesion_path]\n",
        "\n",
        "    fold_img, fold_seg = [], []\n",
        "    for fp,lp in zip(flair_path,lesion_path):\n",
        "        if name == 'Singapore':\n",
        "            # Adjust first 2 data at Singapore\n",
        "            img_temp = [np.expand_dims(read_data(i, start=start, end=end, step=step), -1) for i in fp[:2]]\n",
        "            seg_temp = [np.expand_dims(read_data(i, start=start, end=end, step=step), -1) for i in lp[:2]]\n",
        "            img_temp = np.array([[np.squeeze(tf.image.rot90(j).numpy(), -1) for j in i] for i in img_temp])\n",
        "            seg_temp = np.array([[np.squeeze(tf.image.rot90(j).numpy(), -1) for j in i] for i in seg_temp])\n",
        "            del fp[:2]; del lp[:2]\n",
        "\n",
        "        # Read Data\n",
        "        img = np.array([read_data(i, start=start, end=end, step=step) for i in fp])\n",
        "        seg = np.array([read_data(i, start=start, end=end, step=step) for i in lp])\n",
        "\n",
        "        if name == 'Singapore':\n",
        "            # Concat with first 2 data and cropping\n",
        "            img = np.concatenate([img_temp, img])\n",
        "            seg = np.concatenate([seg_temp, seg])\n",
        "            img = np.array([i[:,12:220,:] for i in img])\n",
        "            seg = np.array([i[:,12:220,:] for i in seg])\n",
        "        elif name == 'GE3T':\n",
        "            # Cropping\n",
        "            img = np.array([i[:,:128,:] for i in img])\n",
        "            seg = np.array([i[:,:128,:] for i in seg])\n",
        "\n",
        "        # Normalized data\n",
        "        img = np.array([[adjust_data(j, cf=cf) for j in i] for i in img])\n",
        "        seg = np.array([[adjust_data(j, cf=cf, mask=True) for j in i] for i in seg])\n",
        "\n",
        "        if name not in ['ADNI']:\n",
        "            seg[(seg == 2).all(axis=-1)] = 0\n",
        "\n",
        "        fold_img.append(img); fold_seg.append(seg)\n",
        "    return np.array(fold_img), np.array(fold_seg)\n",
        "\n",
        "def read_challenge(start, end, step=1, cf=1):\n",
        "    # Read Challenge dataset as full dataset\n",
        "    fold_img, fold_seg = [], []\n",
        "    for name in ['Singapore', 'GE3T', 'Utrecht']:\n",
        "        flair_path = glob(path_dataset+'Challenge/{}/*/*brain.nii.gz'.format(name))\n",
        "        lesion_path = glob(path_dataset+'Challenge/{}/*/*wmh.nii.gz'.format(name))\n",
        "        flair_path.sort(key=natural_keys) \n",
        "        lesion_path.sort(key=natural_keys)\n",
        "        flair_path = [flair_path]\n",
        "        lesion_path = [lesion_path]\n",
        "\n",
        "        for fp,lp in zip(flair_path, lesion_path):\n",
        "            if name == 'Singapore':\n",
        "                # Adjust first 2 data at Singapore\n",
        "                img_temp = [np.expand_dims(read_data(i, start=start, end=end, step=step), -1) for i in fp[:2]]\n",
        "                seg_temp = [np.expand_dims(read_data(i, start=start, end=end, step=step), -1) for i in lp[:2]]\n",
        "                img_temp = np.array([[np.squeeze(tf.image.rot90(j).numpy(), -1) for j in i] for i in img_temp])\n",
        "                seg_temp = np.array([[np.squeeze(tf.image.rot90(j).numpy(), -1) for j in i] for i in seg_temp])\n",
        "                del fp[:2]; del lp[:2]\n",
        "\n",
        "            # Read Data\n",
        "            img = np.array([read_data(i, start=start, end=end, step=step) for i in fp])\n",
        "            seg = np.array([read_data(i, start=start, end=end, step=step) for i in lp])\n",
        "            \n",
        "            if name == 'Singapore':\n",
        "                img = np.concatenate([img_temp, img])\n",
        "                seg = np.concatenate([seg_temp, seg])\n",
        "            fold_img.append(img); fold_seg.append(seg)\n",
        "\n",
        "    # Get the max sizes from all insitutions\n",
        "    sizes = [i.shape[-2:] for i in fold_img]\n",
        "    max_sizes = np.max(list(zip(*sizes)), -1).tolist()\n",
        "    diff_sizes = [(int((max_sizes[0]-i)/2), (int((max_sizes[1]-j)/2))) for i,j in sizes]\n",
        "    fold_img_padd, fold_seg_padd = [], []\n",
        "\n",
        "    for k,(i,j) in enumerate(zip(fold_img, fold_seg)):\n",
        "        # Zero padding\n",
        "        img = np.zeros(i.shape[:2] + tuple(max_sizes))\n",
        "        seg = np.zeros(i.shape[:2] + tuple(max_sizes))\n",
        "        x = sizes[k][0]; x_pad = diff_sizes[k][0]\n",
        "        y = sizes[k][1]; y_pad = diff_sizes[k][1]\n",
        "\n",
        "        img[:,:,x_pad:x+x_pad,y_pad:y+y_pad] = i\n",
        "        seg[:,:,x_pad:x+x_pad,y_pad:y+y_pad] = j\n",
        "\n",
        "        # Normalized data\n",
        "        img = np.array([[adjust_data(j, cf=cf) for j in i] for i in img])\n",
        "        seg = np.array([[adjust_data(j, cf=cf, mask=True) for j in i] for i in seg])\n",
        "        seg[(seg == 2).all(axis=-1)] = 0\n",
        "        fold_img_padd.append(img); fold_seg_padd.append(seg)\n",
        "\n",
        "    return fold_img_padd, fold_seg_padd"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yTc9t70jQlWK"
      },
      "source": [
        "# 3. Parameters"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jCrhKuwp-WYk"
      },
      "source": [
        "# Paths config\n",
        "## For saving WMHs prediction\n",
        "result_path = 'Results/{}/{}/{}'\n",
        "## Pre-trained path\n",
        "path = 'pre-trained/'\n",
        "## Dataset path\n",
        "path_dataset = 'Datasets/'\n",
        "## Evaluation path\n",
        "eval_path = 'Evaluation_Excel/'\n",
        "\n",
        "# Hyperparameters\n",
        "reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=2)\n",
        "generate_pred = 30\n",
        "threshold = 0.25\n",
        "batch_size = 16\n",
        "epoch = 50\n",
        "lr = 0.001\n",
        "gamma = 0.25\n",
        "alpha = 0.5\n",
        "\n",
        "#If not use Probabilistic U-Net please fill this variable as None\n",
        "lr_latent = 0.001\n",
        "\n",
        "if lr_latent is None:\n",
        "    uniq = str(lr)+'_'+str(gamma)+'_'+str(alpha)\n",
        "else:\n",
        "    uniq = str(lr)+'_'+str(lr_latent)+'_'+str(gamma)+'_'+str(alpha)\n",
        "\n",
        "def pick_model(model_name):\n",
        "    n_blocks = 4\n",
        "    n_classes = 1 \n",
        "    input_shape = (None, None, 1)\n",
        "    decoder_filters = (256, 128, 64, 32)\n",
        "\n",
        "    # Deterministic\n",
        "    if model_name == 'U_Net':\n",
        "        md = Unet(use_backbone=False, input_shape=input_shape, attention=False,\n",
        "                n_upsample_blocks=n_blocks, decoder_filters=decoder_filters)\n",
        "\n",
        "    elif model_name == 'AttU_Net':\n",
        "        md = Unet(use_backbone=False, input_shape=input_shape, attention=True,\n",
        "                n_upsample_blocks=n_blocks, decoder_filters=decoder_filters)\n",
        "        \n",
        "    elif model_name == 'XNet':\n",
        "        md = Xnet(use_backbone=False, input_shape=input_shape, attention=False,\n",
        "                n_upsample_blocks=n_blocks, decoder_filters=decoder_filters, deep_supervision=False)\n",
        "\n",
        "    elif model_name == 'AttXNet':\n",
        "        md = Xnet(use_backbone=False, input_shape=input_shape, attention=True,\n",
        "                n_upsample_blocks=n_blocks, decoder_filters=decoder_filters, deep_supervision=False)\n",
        "\n",
        "    elif model_name == 'XNet_ds':\n",
        "        md = Xnet(use_backbone=False, input_shape=input_shape, attention=False,\n",
        "                n_upsample_blocks=n_blocks, decoder_filters=decoder_filters, deep_supervision=True)\n",
        "\n",
        "    elif model_name == 'AttXNet_ds':\n",
        "        md = Xnet(use_backbone=False, input_shape=input_shape, attention=True,\n",
        "                n_upsample_blocks=n_blocks, decoder_filters=decoder_filters, deep_supervision=True)\n",
        "\n",
        "    # Probabilistic\n",
        "    elif model_name == 'Prob_U_Net':\n",
        "        md = Prob_Unet(num_classes=n_classes, activation='sigmoid', latent_dim=6)\n",
        "\n",
        "    else:\n",
        "        md = None\n",
        "    return md"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KP5v3szUIvwv"
      },
      "source": [
        "# 4. Experiments"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HtkuJtLwI6w8"
      },
      "source": [
        "## 4.1. K-Fold Cross Validation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "se7uwr3wI6NZ"
      },
      "source": [
        "# U_Net, AttU_Net, XNet, AttXNet, XNet_ds, AttXNet_ds, Prob_U_Net\n",
        "model_name = 'Prob_U_Net'\n",
        "\n",
        "# ADNI, Challenge\n",
        "train = 'ADNI'\n",
        "amount = 2\n",
        "dsc_list = []\n",
        "kf = KFold(n_splits=amount)\n",
        "fold_img, fold_seg = get_dataset(train)\n",
        "\n",
        "for k, (train_idx, test_idx) in enumerate(kf.split(fold_img)):\n",
        "    num = 0 if k == 0 else 10 \n",
        "    # Training data \n",
        "    if train == 'ADNI':\n",
        "        img_temp = np.concatenate([fold_img[train_idx[0]], fold_img[train_idx[1]]])\n",
        "        seg_temp = np.concatenate([fold_seg[train_idx[0]], fold_seg[train_idx[1]]])\n",
        "        img_train = np.concatenate(img_temp[:24])\n",
        "        seg_train = np.concatenate(seg_temp[:24])\n",
        "        img_val = np.concatenate(img_temp[-6:])\n",
        "        seg_val = np.concatenate(seg_temp[-6:])\n",
        "    else:\n",
        "        img_temp = [i[num:num+10] for i in fold_img]\n",
        "        seg_temp = [i[num:num+10] for i in fold_seg]\n",
        "        img_train = np.concatenate([np.concatenate(i[:-2]) for i in img_temp])\n",
        "        seg_train = np.concatenate([np.concatenate(i[:-2]) for i in seg_temp])\n",
        "        img_val = np.concatenate([np.concatenate(i[-2:]) for i in img_temp])\n",
        "        seg_val = np.concatenate([np.concatenate(i[-2:]) for i in seg_temp])\n",
        "\n",
        "    del img_temp, seg_temp\n",
        "    img_train, seg_train = check_data(img_train, seg_train)\n",
        "    img_val, seg_val = check_data(img_val, seg_val)\n",
        "    data_augmentation(img_train, seg_train)\n",
        "    \n",
        "    # Model\n",
        "    logdir = path+'logdir/{}_{}_fold{}_2_{}'.format(model_name,uniq,k+1,train)\n",
        "    temp = logdir.split('/')[-1]\n",
        "    # Becareful with these 4 lines!\n",
        "    ## Will delete existing folders and files\n",
        "    if len(glob(logdir)) >= 1: shutil.rmtree(logdir)\n",
        "    if len(glob(path+temp)) >= 1: shutil.rmtree(path+temp)\n",
        "    ## Will create folders and files\n",
        "    if len(glob(logdir)) == 0: os.mkdir(logdir)\n",
        "    if len(glob(path+temp)) == 0: os.mkdir(path+temp)\n",
        "\n",
        "    tensorboard = TensorBoard(log_dir=logdir)\n",
        "\n",
        "    save_weights_only = True if lr_latent is not None else False\n",
        "    model_checkpoint = ModelCheckpoint(\n",
        "        filepath=path+temp+\"/checkpoint-best-{epoch:02d}-{val_loss:.4f}.hdf5\",\n",
        "        monitor='val_loss', mode='min', save_best_only=True, save_weights_only=save_weights_only)\n",
        "    \n",
        "    # Compile model\n",
        "    md = pick_model(model_name)\n",
        "    if model_name == 'Prob_U_Net':\n",
        "        md.compile(unet_opt=Adam(learning_rate=lr), \n",
        "                   prior_opt=Adam(learning_rate=lr_latent), \n",
        "                   posterior_opt=Adam(learning_rate=lr_latent),\n",
        "                   loss=binary_focal_loss(gamma=gamma, alpha=alpha), metric=dice_coef)\n",
        "    \n",
        "    else:\n",
        "        md.compile(optimizer=Adam(learning_rate=lr), metrics=[dice_coef],\n",
        "                   loss=binary_focal_loss(gamma=gamma, alpha=alpha))\n",
        "\n",
        "    # Training\n",
        "    md.fit(img_train, seg_train, epochs=epoch, batch_size=batch_size,\n",
        "            validation_data=(img_val,seg_val), verbose=2,\n",
        "            callbacks=[model_checkpoint, tensorboard])\n",
        "    print('Finish training on fold-', k+1,'\\n')\n",
        "\n",
        "    # Testing\n",
        "    if lr_latent is not None: md.built = True\n",
        "    best_weight = find_best(path+temp)\n",
        "    md.load_weights(best_weight)\n",
        "    dsc_patient = []\n",
        "\n",
        "    ## ADNI Testing\n",
        "    if train == 'ADNI':\n",
        "        paths = glob(path_dataset+'ADNI/{}/*'.format('fold'+str(test_idx[0]+1)))\n",
        "        paths = paths + glob(path_dataset+'ADNI/{}/*'.format('fold'+str(test_idx[1]+1)))\n",
        "        paths.sort(key=natural_keys)\n",
        "\n",
        "        img_test = np.concatenate([fold_img[test_idx[0]], fold_img[test_idx[1]]])\n",
        "        seg_test = np.concatenate([fold_seg[test_idx[0]], fold_seg[test_idx[1]]])\n",
        "\n",
        "        for idx_p, (img, seg) in enumerate(zip(img_test, seg_test)):\n",
        "            if lr_latent is not None:\n",
        "                y_gen = []\n",
        "                for x in range(generate_pred):\n",
        "                    if model_name == 'Prob_U_Net':\n",
        "                        z_sample, _,_ = md.prior.predict(img)\n",
        "                        y_pred = md.det_unet.predict([img, z_sample])\n",
        "                    else:\n",
        "                        y_pred = md.predict(img)\n",
        "\n",
        "                    y_gen.append(y_pred)\n",
        "\n",
        "                y_gen = np.array(y_gen)\n",
        "                y_ss = ambiguity_map(y_gen)\n",
        "                y_sy = ambiguity_map(y_gen, seg)\n",
        "                y_pred = np.average(y_gen, axis=0)\n",
        "            else:\n",
        "                y_pred = md.predict(img)\n",
        "                y_ss = None; y_sy = None\n",
        "\n",
        "            y_pred[y_pred > threshold] = 1\n",
        "            y_pred[y_pred <= threshold] = 0\n",
        "            dsc = dice_coef(seg, y_pred).numpy()\n",
        "            dsc_patient.append(dsc)\n",
        "            \n",
        "            patient_id = paths[idx_p].split('/')[-1]\n",
        "            fold_id = paths[idx_p].split('/')[-2]\n",
        "            file_in = '{}/{}_wmh.nii.gz'.format(paths[idx_p], patient_id)\n",
        "            file_out = result_path.format('ADNI', fold_id, patient_id)\n",
        "            if len(glob(file_out)) == 0: os.makedirs(file_out)\n",
        "            file_out_pred = file_out+'/{}_wmh_{}.nii.gz'.format(patient_id, model_name)\n",
        "            show_img(img, seg, y_pred, 18)\n",
        "            save_wmh(y_pred, file_in, file_out_pred)\n",
        "\n",
        "            # Ambiguity maps\n",
        "            if y_ss is not None and y_sy is not None:\n",
        "                file_out_1 = file_out+'/{}_am_ss_{}.nii.gz'.format(patient_id, model_name)\n",
        "                file_out_2 = file_out+'/{}_am_sy_{}.nii.gz'.format(patient_id, model_name)\n",
        "                save_wmh(y_ss, file_in, file_out_1)\n",
        "                save_wmh(y_sy, file_in, file_out_2)\n",
        "\n",
        "    ## Challenge Testing\n",
        "    else:\n",
        "        for iter_,inst in enumerate(['Singapore', 'GE3T', 'Utrecht']):\n",
        "            dsc_temp = []\n",
        "            paths = glob(path_dataset+'Challenge/{}/*'.format(inst))\n",
        "            paths.sort(key=natural_keys)\n",
        "            paths = paths[10-num:20-num]\n",
        "            \n",
        "            img_test = np.array(fold_img[iter_][10-num:20-num], dtype='float32')\n",
        "            seg_test = np.array(fold_seg[iter_][10-num:20-num], dtype='float32')\n",
        "    \n",
        "            for idx_p, (img, seg) in enumerate(zip(img_test, seg_test)):\n",
        "                if lr_latent is not None:\n",
        "                    y_gen = []\n",
        "                    for x in range(generate_pred):\n",
        "                        if model_name == 'Prob_U_Net':\n",
        "                            z_sample, _,_ = md.prior.predict(img)\n",
        "                            y_pred = md.det_unet.predict([img, z_sample])\n",
        "                        else:\n",
        "                            y_pred = md.predict(img)\n",
        "\n",
        "                        y_gen.append(y_pred)\n",
        "\n",
        "                    y_gen = np.array(y_gen)\n",
        "                    y_ss = ambiguity_map(y_gen)\n",
        "                    y_sy = ambiguity_map(y_gen, seg)\n",
        "                    y_pred = np.average(y_gen, axis=0)\n",
        "                else:\n",
        "                    y_pred = md.predict(img)\n",
        "                    y_ss = None; y_sy = None\n",
        "\n",
        "                y_pred[y_pred > threshold] = 1\n",
        "                y_pred[y_pred <= threshold] = 0\n",
        "                dsc = dice_coef(seg, y_pred).numpy()\n",
        "                dsc_temp.append(dsc)\n",
        "                \n",
        "                patient_id = paths[idx_p].split('/')[-1]\n",
        "                file_in = '{}/{}_wmh.nii.gz'.format(paths[idx_p], patient_id)\n",
        "                file_out = result_path.format('Challenge', inst, patient_id)\n",
        "                if len(glob(file_out)) == 0: os.makedirs(file_out)\n",
        "                file_out_pred = file_out+'/{}_wmh_{}.nii.gz'.format(patient_id, model_name)\n",
        "                show_img(img, seg, y_pred, 20)\n",
        "                save_wmh_challenge(y_pred, file_in, file_out_pred, name=inst)\n",
        "\n",
        "                # Ambiguity maps\n",
        "                if y_ss is not None and y_sy is not None:\n",
        "                    file_out_1 = file_out+'/{}_am_ss_{}.nii.gz'.format(patient_id, model_name)\n",
        "                    file_out_2 = file_out+'/{}_am_sy_{}.nii.gz'.format(patient_id, model_name)\n",
        "                    save_wmh_challenge(y_ss, file_in, file_out_1, name=inst)\n",
        "                    save_wmh_challenge(y_sy, file_in, file_out_2, name=inst)\n",
        "\n",
        "            dsc_patient.append(dsc_temp)\n",
        "    dsc_list.append(dsc_patient)\n",
        "    print('Average Dice Coefficient on fold-{} : {:.4f}'.format(k+1, np.average(dsc_patient)))\n",
        "print('Average Dice Coefficient : {:.4f}'.format(np.average(dsc_list)))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fKnJNBsumqq1"
      },
      "source": [
        "## 4.2. Cross Dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6lUXqsAGswN6"
      },
      "source": [
        "# U_Net, AttU_Net, XNet, AttXNet, XNet_ds, AttXNet_ds, Prob_U_Net\n",
        "model_name = 'Prob_U_Net'\n",
        "\n",
        "# ADNI, Challenge\n",
        "train = 'ADNI'\n",
        "\n",
        "# ADNI, Singapore, GE3T, Utrecht\n",
        "test = ['Singapore', 'GE3T', 'Utrecht']\n",
        "\n",
        "fold_img, fold_seg = get_dataset(train)\n",
        "dsc_dict = {i:{j:0 for j in test} for i in [train]}\n",
        "\n",
        "for type_data in [train]: # train data\n",
        "    # Training data\n",
        "    if type_data == 'Challenge':\n",
        "        img_train = np.concatenate([np.concatenate(i[:-4]) for i in fold_img])\n",
        "        seg_train = np.concatenate([np.concatenate(i[:-4]) for i in fold_seg])\n",
        "        img_val = np.concatenate([np.concatenate(i[-4:]) for i in fold_img])\n",
        "        seg_val = np.concatenate([np.concatenate(i[-4:]) for i in fold_seg])\n",
        "    else:\n",
        "        if type_data == 'ADNI':\n",
        "            img_temp = np.concatenate(fold_img)\n",
        "            seg_temp = np.concatenate(fold_seg)\n",
        "        else:\n",
        "            img_temp = fold_img[0]\n",
        "            seg_temp = fold_seg[0]\n",
        "\n",
        "        img_train = np.concatenate(img_temp[:-6])\n",
        "        seg_train = np.concatenate(seg_temp[:-6])\n",
        "        img_val = np.concatenate(img_temp[-6:])\n",
        "        seg_val = np.concatenate(seg_temp[-6:])\n",
        "\n",
        "    img_train, seg_train = check_data(img_train, seg_train)\n",
        "    img_val, seg_val = check_data(img_val, seg_val)\n",
        "    data_augmentation(img_train, seg_train)\n",
        "\n",
        "    # Model\n",
        "    logdir = path+'logdir/{}_{}_{}'.format(model_nam, uniq, type_data)\n",
        "    temp = logdir.split('/')[-1]\n",
        "    # # Becareful with these 4 lines!\n",
        "    ## Will delete existing folders and files\n",
        "    if len(glob(logdir)) >= 1: shutil.rmtree(logdir)\n",
        "    if len(glob(path+temp)) >= 1: shutil.rmtree(path+temp)\n",
        "    ## Will create folders and files\n",
        "    if len(glob(logdir)) == 0: os.mkdir(logdir)\n",
        "    if len(glob(path+temp)) == 0: os.mkdir(path+temp)\n",
        "\n",
        "    tensorboard = TensorBoard(log_dir=logdir)\n",
        "\n",
        "    save_weights_only = True if lr_latent is not None else False\n",
        "    model_checkpoint = ModelCheckpoint(\n",
        "        filepath=path+temp+\"/checkpoint-best-{epoch:02d}-{val_loss:.4f}.hdf5\",\n",
        "        monitor='val_loss', mode='min', save_best_only=True, save_weights_only=save_weights_only)\n",
        "\n",
        "    # Compile model\n",
        "    md = pick_model(model_name)\n",
        "    if model_name == 'Prob_U_Net':\n",
        "        md.compile(unet_opt=Adam(learning_rate=lr), \n",
        "                   prior_opt=Adam(learning_rate=lr_latent), \n",
        "                   posterior_opt=Adam(learning_rate=lr_latent),\n",
        "                   loss=binary_focal_loss(gamma=gamma, alpha=alpha), metric=dice_coef)\n",
        "    else:\n",
        "        md.compile(optimizer=Adam(learning_rate=lr), metrics=[dice_coef], \n",
        "                   loss=binary_focal_loss(gamma=gamma, alpha=alpha))\n",
        "\n",
        "    # Training\n",
        "    md.fit(img_train, seg_train, epochs=epoch, batch_size=batch_size,\n",
        "            validation_data=(img_val,seg_val), verbose=2,\n",
        "            callbacks=[reduce_lr, model_checkpoint, tensorboard])\n",
        "    print('Finished Training with Dataset ', type_data, '\\n')\n",
        "\n",
        "    # Testing\n",
        "    if lr_latent is not None: md.built = True\n",
        "    best_weight = find_best(path+temp)\n",
        "    md.load_weights(best_weight)\n",
        "\n",
        "    for i in test: # test data\n",
        "        score = []\n",
        "        if i == 'ADNI': paths = glob(path_dataset+'ADNI/*/*')\n",
        "        else: paths = glob(path_dataset+'Challenge/{}/*'.format(i))\n",
        "        paths.sort(key=natural_keys)\n",
        "        \n",
        "        img_test, seg_test = get_dataset(i)\n",
        "        img_test = np.concatenate(img_test)\n",
        "        seg_test = np.concatenate(seg_test)\n",
        "        \n",
        "        for idx_p, (img, seg) in enumerate(zip(img_test, seg_test)):\n",
        "            if lr_latent is not None:\n",
        "                y_gen = []\n",
        "                for x in range(generate_pred):\n",
        "                    if model_name == 'Prob_U_Net':\n",
        "                        z_sample, _,_ = md.prior.predict(img)\n",
        "                        y_pred = md.det_unet.predict([img, z_sample])\n",
        "                    else:\n",
        "                        y_pred = md.predict(img)\n",
        "\n",
        "                    y_gen.append(y_pred)\n",
        "\n",
        "                y_gen = np.array(y_gen)\n",
        "                y_ss = ambiguity_map(y_gen)\n",
        "                y_sy = ambiguity_map(y_gen, seg)\n",
        "                y_pred = np.average(y_gen, axis=0)\n",
        "            else:\n",
        "                y_pred = md.predict(img)\n",
        "                y_ss = None; y_sy = None\n",
        "\n",
        "            y_pred[y_pred > threshold] = 1\n",
        "            y_pred[y_pred <= threshold] = 0\n",
        "            dsc = dice_coef(seg, y_pred).numpy()\n",
        "            score.append(dsc)\n",
        "\n",
        "            patient_id = paths[idx_p].split('/')[-1]\n",
        "            if i == 'ADNI':\n",
        "                fold_id = paths[idx_p].split('/')[-2]\n",
        "                file_out = result_path.format(i, fold_id, patient_id)\n",
        "            else:\n",
        "                fold_id = 'Challenge'\n",
        "                file_out = result_path.format(fold_id, i, patient_id)\n",
        "\n",
        "            file_in = '{}/{}_wmh.nii.gz'.format(paths[idx_p], patient_id)\n",
        "            if len(glob(file_out)) == 0: os.makedirs(file_out)\n",
        "            file_out_pred = file_out+'/{}_wmh_{}_{}.nii.gz'.format(patient_id, model_name,type_data)\n",
        "            show_img(img, seg, y_pred, 18)\n",
        "            save_wmh(y_pred, file_in, file_out_pred, name=i)\n",
        "\n",
        "            # Ambiguity maps\n",
        "            if y_ss is not None and y_sy is not None:\n",
        "                file_out_1 = file_out+'/{}_am_ss_{}_{}.nii.gz'.format(patient_id, model_name, type_data)\n",
        "                file_out_2 = file_out+'/{}_am_sy_{}_{}.nii.gz'.format(patient_id, model_name, type_data)\n",
        "                save_wmh(y_ss, file_in, file_out_1, name=i)\n",
        "                save_wmh(y_sy, file_in, file_out_2, name=i)\n",
        "\n",
        "        dsc_dict[type_data][i] = np.average(score)\n",
        "        print('Dice Coefficient with dataset {} trained by {}: {:.4f}'.format(i, type_data, dsc_dict[type_data][i]))\n",
        "\n",
        "    avg_ = np.average([dsc_dict[train][i] for i in test])\n",
        "    print('Average: {:.4f}'.format(avg_))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xi74WxiMcC7O"
      },
      "source": [
        "# 5. Evaluation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8e6-f9k-ygT3"
      },
      "source": [
        "# 'ADNI', 'Challenge'\n",
        "train = 'ADNI'\n",
        "# 'ADNI', 'Singapore', 'GE3T', 'Utrecht', 'Challenge'\n",
        "test = ['Singapore', 'GE3T', 'Utrecht']\n",
        "# U_Net, AttU_Net, XNet, AttXNet, KiU_Net, Prob_U_Net\n",
        "model_name = 'Prob_U_Net'\n",
        "df_by_model, dsc, patient = [], [], []\n",
        "\n",
        "for i in test: # testing \n",
        "    if i == 'ADNI': \n",
        "        paths = glob(path_dataset+'ADNI/*/*/*wmh.nii.gz')\n",
        "    elif i == 'Challenge':\n",
        "        paths = []\n",
        "        for j in ['Singapore', 'GE3T', 'Utrecht']:\n",
        "            path = glob(path_dataset+'Challenge/{}/*/*wmh.nii.gz'.format(j))\n",
        "            path.sort(key=natural_keys)\n",
        "            paths += path\n",
        "    else: \n",
        "        paths = glob(path_dataset+'Challenge/{}/*/*wmh.nii.gz'.format(i))\n",
        "\n",
        "    if i != 'Challenge': paths.sort(key=natural_keys)\n",
        "    temp, temp_dsc, temp_patient = [], [], []\n",
        "\n",
        "    for lp in paths:\n",
        "        patient_id = lp.split('/')[-2]\n",
        "        if i == 'ADNI':\n",
        "            fold_id = lp.split('/')[-3]\n",
        "            pp = result_path.format(i, fold_id, patient_id)\n",
        "        elif i == 'Challenge':\n",
        "            fold_id = lp.split('/')[-3]\n",
        "            pp = result_path.format('Challenge', fold_id, patient_id)\n",
        "        else:\n",
        "            pp = result_path.format('Challenge', i, patient_id)\n",
        "        \n",
        "        if i == train:\n",
        "            pp = pp+'/{}_wmh_{}.nii.gz'.format(patient_id, model_name) # training\n",
        "        else:\n",
        "            pp = pp+'/{}_wmh_{}_{}.nii.gz'.format(patient_id, model_name, train) # training\n",
        "\n",
        "        voxel_true = nib.load(lp); wmh_true = read_data(lp)\n",
        "        voxel_pred = nib.load(pp); wmh_pred = read_data(pp)\n",
        "        volume_true = np.prod([abs(voxel_true.affine[i][i]) for i in range(3)])/1000\n",
        "        volume_true = np.count_nonzero(wmh_true)*volume_true\n",
        "        volume_pred = np.prod([abs(voxel_pred.affine[i][i]) for i in range(3)])/1000\n",
        "        volume_pred = np.count_nonzero(wmh_pred)*volume_pred\n",
        "        score = dice_coef(wmh_true, wmh_pred).numpy()\n",
        "        temp_dsc.append(score)\n",
        "        temp_patient.append(patient_id)\n",
        "        temp.append([volume_true, volume_pred])\n",
        "    df_by_model.append(temp)\n",
        "    dsc.append(temp_dsc)\n",
        "    patient.append(temp_patient)\n",
        "\n",
        "df_temp = [pd.DataFrame(i, index=j).transpose() for i,j in zip(df_by_model, patient)]\n",
        "temp_str = '_Cross' if train not in test else ''\n",
        "with pd.ExcelWriter(eval_path+'output{}_{}_{}.xlsx'.format(temp_str, train, model_name)) as writter:\n",
        "    num = 0\n",
        "    for k,i in enumerate(test):\n",
        "        df_temp[k].to_excel(writter, startrow=num, startcol=0, sheet_name='vol')\n",
        "        pd.DataFrame([dsc[k]], columns=df_temp[k].columns).to_excel(\n",
        "            writter, startrow=num, startcol=0, sheet_name='dsc')\n",
        "        num = num + df_temp[k].shape[0] + 3"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}