Train_Script.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/AntonioVispi/Recursive_Segmentation/blob/main/Train_Script.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JzKiaA7GJ55E"
      },
      "outputs": [],
      "source": [
        "!pip install -U -q segmentation-models\n",
        "!pip install tensorflow==2.9.3\n",
        "!pip install h5py==2.10.0\n",
        "!pip install plotly==5.3.1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zGkN2rkBVFKT"
      },
      "outputs": [],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9R4TcDAIVB2R"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import os\n",
        "\n",
        "import nibabel as nib\n",
        "import tensorflow as tf\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "from matplotlib import pyplot as plt\n",
        "from tqdm import tqdm\n",
        "\n",
        "import skimage\n",
        "from skimage.io import imread, imshow, imsave\n",
        "from skimage.transform import resize\n",
        "\n",
        "from tensorflow import keras\n",
        "from keras.callbacks import ModelCheckpoint\n",
        "from keras.callbacks import CSVLogger\n",
        "from keras.callbacks import EarlyStopping\n",
        "from keras.utils.np_utils import to_categorical\n",
        "from keras.preprocessing.image import ImageDataGenerator\n",
        "from keras.models import load_model\n",
        "from keras import metrics\n",
        "from keras.callbacks import ReduceLROnPlateau\n",
        "\n",
        "from segmentation_models import Unet\n",
        "import math\n",
        "from math import floor\n",
        "\n",
        "import random\n",
        "from random import seed\n",
        "from random import random"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Useful functions"
      ],
      "metadata": {
        "id": "gTKQuz8T2Zhn"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qjwjAWgrxXHZ"
      },
      "outputs": [],
      "source": [
        "def dispenser(dataset_path, Total_Slices):     # Function to establish how many slices to sample from each\n",
        "                                                                      # subject: the greater the slices of the subject, the more slices will be selected.\n",
        "  dataset_path = dataset_path + '/'\n",
        "  directory = os.listdir(dataset_path)\n",
        "  directory.sort()\n",
        "\n",
        "  num_subj=len(directory)\n",
        "  All_Slices = np.zeros(num_subj)\n",
        "\n",
        "  i = 0\n",
        "  for n in range (0,num_subj):\n",
        "      case_path = dataset_path +'/'+ directory[n]\n",
        "      volume = nib.load(os.path.join(case_path, \"imaging.nii.gz\"))\n",
        "      n_slice,height,width = volume.shape\n",
        "      All_Slices[i] = n_slice\n",
        "      i = i + 1\n",
        "\n",
        "  brick = round((min(All_Slices)/100) * 5)\n",
        "  num_bricks = round(Total_Slices / brick)\n",
        "  aux_sum = sum(All_Slices)\n",
        "  Slice_Map = np.round((All_Slices / aux_sum) * num_bricks) * brick\n",
        "\n",
        "  return Slice_Map\n",
        "\n",
        "\n",
        "def visualizer(segm, IMG_HEIGHT, IMG_WIDTH):     #Function useful for displaying segmentations\n",
        "\n",
        "   segm = resize(segm,(IMG_HEIGHT,IMG_WIDTH,4), mode='constant', preserve_range=True)\n",
        "\n",
        "   back = segm[:,:,0]\n",
        "   kid = segm[:,:,1]\n",
        "   tum = segm[:,:,2]\n",
        "   cys = segm[:,:,3]\n",
        "\n",
        "   back = (back == 1)\n",
        "   kid = (kid == 1)\n",
        "   tum = (tum == 1)\n",
        "   cys = (cys == 1)\n",
        "\n",
        "   all_segments = np.zeros((IMG_HEIGHT, IMG_WIDTH, 3))\n",
        "\n",
        "   all_segments[back] = (1,0,0)\n",
        "   all_segments[kid] = (0,1,0)\n",
        "   all_segments[tum] = (0,0,1)\n",
        "   all_segments[cys] = (1,1,0)\n",
        "\n",
        "   return all_segments\n",
        "\n",
        "\n",
        "def Make_Dataset(dataset_path, Total_Slices, view_dataset):     # Function that actually defines the dataset.\n",
        "\n",
        "    print('Beginning definition of Dataset useful for training...')\n",
        "    print('\\n')\n",
        "\n",
        "    directory = os.listdir(dataset_path)\n",
        "    directory.sort()\n",
        "\n",
        "    num_subj=len(directory)                              # Get the number of subjects\n",
        "\n",
        "    Slice_Map = dispenser(dataset_path, Total_Slices)\n",
        "    Total_Slices = int(sum(Slice_Map))              # Get the actual number of total slices\n",
        "\n",
        "    IMG_HEIGHT= 512\n",
        "    IMG_WIDTH= 512\n",
        "    IMG_CHANNELS = 3\n",
        "    NUM_CLASSES = 4\n",
        "\n",
        "    X_train=np.zeros([Total_Slices, IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS], dtype=np.uint8) # Default training image volume\n",
        "    Y_train=np.zeros([Total_Slices, IMG_HEIGHT, IMG_WIDTH, NUM_CLASSES], dtype=np.float32) # Default training segmentation volume\n",
        "\n",
        "    k=0\n",
        "    checkpoint = 0\n",
        "\n",
        "    for n in range (0,num_subj):\n",
        "        case_path = dataset_path +'/'+ directory[n]\n",
        "        volume = nib.load(os.path.join(case_path, \"imaging.nii.gz\"))\n",
        "        segmentation = nib.load(os.path.join(case_path, \"segmentation.nii.gz\"))\n",
        "\n",
        "        # Selection of some slices\n",
        "        n_slice,height,width = volume.shape  # Subject dimensions\n",
        "\n",
        "        volume_new = volume.slicer[0 : n_slice]\n",
        "        segmentation_new = segmentation.slicer[0 : n_slice]\n",
        "\n",
        "        # Conversion to ndarray\n",
        "        Vol_train_supporto= volume_new.get_fdata().astype(np.int16)\n",
        "        Mask_train_supporto = segmentation_new.get_fdata().astype(np.uint8)\n",
        "        slices_per_subject = int(Slice_Map[k])\n",
        "\n",
        "        k=k+1\n",
        "\n",
        "        # Take n random and unique slices from the subject\n",
        "        Vol_train=np.zeros((slices_per_subject,IMG_HEIGHT,IMG_WIDTH))\n",
        "        Mask_train=np.zeros((slices_per_subject,IMG_HEIGHT,IMG_WIDTH))\n",
        "        fette_randomiche=np.zeros(slices_per_subject)\n",
        "        support = np.random.permutation(n_slice)\n",
        "        for j in range (0,slices_per_subject):\n",
        "          fette_randomiche[j]=support[j]\n",
        "\n",
        "          if height!=IMG_HEIGHT or width!=IMG_WIDTH:\n",
        "\n",
        "            Vol_train[j,:,:] = resize(Vol_train_supporto[int(abs(fette_randomiche[j]-1)),:,:], (IMG_HEIGHT,IMG_WIDTH), mode='constant', preserve_range=True)\n",
        "            support_1 = to_categorical(Mask_train_supporto[int(abs(fette_randomiche[j]-1)),:,:], num_classes=NUM_CLASSES, dtype='float32')\n",
        "            support_2 = resize(support_1, (IMG_HEIGHT,IMG_WIDTH,4), mode='constant', preserve_range=True)\n",
        "            Mask_train[j,:,:] = np.argmax(support_2,axis=2).astype(np.uint8)\n",
        "\n",
        "          else:\n",
        "            Vol_train[j,:,:] = Vol_train_supporto[int(abs(fette_randomiche[j]-1)),:,:]\n",
        "            Mask_train[j,:,:] = Mask_train_supporto[int(abs(fette_randomiche[j]-1)),:,:]\n",
        "\n",
        "\n",
        "        vol_train_2 = np.zeros([slices_per_subject, IMG_HEIGHT, IMG_WIDTH,IMG_CHANNELS], dtype=np.uint8)\n",
        "        mask_train_2=np.zeros([slices_per_subject, IMG_HEIGHT, IMG_WIDTH,NUM_CLASSES], dtype=np.float32)\n",
        "        for i, id_ in tqdm(enumerate(Vol_train+1), total=len(Vol_train)):\n",
        "\n",
        "          mask_train_2[i] = to_categorical(Mask_train[i], num_classes=NUM_CLASSES, dtype='float32')\n",
        "          vol_train_2[i]=resize(Vol_train[i,:,:], (IMG_HEIGHT,IMG_WIDTH,1), mode='constant', preserve_range=True)\n",
        "\n",
        "        X_train[checkpoint:(checkpoint + slices_per_subject),:,:,: ]=vol_train_2     # Final definition X_train\n",
        "        Y_train[checkpoint:(checkpoint + slices_per_subject),:,:,: ]=mask_train_2    # Final definition Y_train\n",
        "\n",
        "        checkpoint = checkpoint + slices_per_subject\n",
        "\n",
        "        print(volume.shape)\n",
        "\n",
        "    if view_dataset == 'on':\n",
        "      for n in range(0,X_train.shape[0]):\n",
        "        fig = plt.figure(figsize=(20, 20))\n",
        "        fig = plt.figure()\n",
        "        ax1 = fig.add_subplot(121)\n",
        "        ax1.imshow(X_train[n,:,:,:]), ax1.set_title('TAC Image')\n",
        "        ax2= fig.add_subplot(122)\n",
        "        ax2.imshow(visualizer(Y_train[n,:,:,:], IMG_HEIGHT, IMG_WIDTH)), ax2.set_title('Manual Segmentation')\n",
        "\n",
        "    print('The definition phase of the Dataset for training has ended.')\n",
        "    print('\\n')\n",
        "\n",
        "    return X_train, Y_train"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3D6VDWglqkj6"
      },
      "outputs": [],
      "source": [
        "dataset_path = '/content/drive/your_dataset_path'\n",
        "\n",
        "print('Definition of the Trainin set')\n",
        "print('\\n')\n",
        "\n",
        "X_train, Y_train = Make_Dataset(dataset_path, 8000, 'off') # the final number of training images is set to 8000.\n",
        "\n",
        "print('Definition of the Validation set')\n",
        "print('\\n')\n",
        "\n",
        "X_val, Y_val = Make_Dataset(dataset_path, 1200, 'off') # the final number of validation images is set to 8000."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "CrYuKulPW9iQ"
      },
      "outputs": [],
      "source": [
        "# Defining the characteristics of the data that are generated from the available images\n",
        "image_datagen = ImageDataGenerator(rotation_range = 23,\n",
        "                                   width_shift_range = 0.21,\n",
        "                                   height_shift_range = 0.21,\n",
        "                                   horizontal_flip = True,\n",
        "                                   vertical_flip = False,\n",
        "                                   zoom_range = 0.28,\n",
        "                                   fill_mode = 'nearest')\n",
        "\n",
        "#Data augmentation\n",
        "val_datagen = ImageDataGenerator()\n",
        "\n",
        "#Generator\n",
        "seed = 1\n",
        "def XYaugmentGenerator(X1, y, seed, batch_size):\n",
        "    genX1 = image_datagen.flow(X1, y, batch_size=batch_size, seed=seed)\n",
        "    genX2 = image_datagen.flow(y, X1, batch_size=batch_size, seed=seed)\n",
        "    while True:\n",
        "        X1i = genX1.next()\n",
        "        X2i = genX2.next()\n",
        "        yield X1i[0], X2i[0]"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Run the following block if the training is from scratch: first round."
      ],
      "metadata": {
        "id": "7fac8PSn6NZU"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kU0TW3TZcds5"
      },
      "outputs": [],
      "source": [
        "NUM_CLASSES = 4\n",
        "\n",
        "BACKBONE = 'efficientnetb5' # <-------- chosen architecture\n",
        "\n",
        "model = Unet(backbone_name=BACKBONE,\n",
        "            input_shape=(512,512,3),\n",
        "            encoder_weights='imagenet',\n",
        "            encoder_freeze=False,\n",
        "            decoder_block_type='transpose',\n",
        "            classes= NUM_CLASSES,\n",
        "            decoder_filters=(512, 256, 128, 64, 32),\n",
        "            decoder_use_batchnorm=True,\n",
        "            activation='sigmoid')\n",
        "\n",
        "# Optimization algorithm definition and loss function\n",
        "\n",
        "m = [ metrics.TruePositives(name='tp') ]\n",
        "model.compile('Nadam', loss='mse', metrics=m)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CqneOw3z1en-"
      },
      "source": [
        "# Run the following block if the training is resumed: rounds following the first."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "l478R9my1k6F"
      },
      "outputs": [],
      "source": [
        "model=load_model('/content/drive/output_training_path/Epoch_15-Val_Loss0.00022.h5',custom_objects=None, compile=False)\n",
        "\n",
        "# Optimization algorithm definition and loss function\n",
        "\n",
        "m = [ metrics.TruePositives(name='tp') ]\n",
        "\n",
        "opt =tf.keras.optimizers.Nadam(learning_rate=1e-03)                    #suggested values for cascade trainings: --> 1e-03 --> 3.16227766e-04 --> 1e-04.\n",
        "\n",
        "model.compile(optimizer=opt, loss='mse', metrics=m)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's get the training going!"
      ],
      "metadata": {
        "id": "X10YRxQi7za0"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cJNZySOsnEEc"
      },
      "outputs": [],
      "source": [
        "path = '/content/drive/output_training_path'\n",
        "batch_size = 16\n",
        "n_epochs = 100     # Reasonably large choice\n",
        "\n",
        "# CVSLogger definition\n",
        "csv_logger = CSVLogger('./log.out', append=True, separator=';')\n",
        "\n",
        "#Earlystopping definition\n",
        "earlystopping = EarlyStopping(monitor = 'val_tp',verbose = 1, min_delta = 0.01, patience = 65, mode = 'max')\n",
        "\n",
        "#Learning rate definition\n",
        "reduce_LR=ReduceLROnPlateau(monitor='val_tp', factor=0.316227766, patience=20, verbose=1, mode='max', min_delta=0.0001, cooldown=0, min_lr=0)\n",
        "\n",
        "#Checkpoint rate definition\n",
        "checkpoint = ModelCheckpoint(path + '/Epoch_{epoch:02d}-Val_Loss{val_loss:.5f}.h5', monitor='val_tp', mode='max', save_best_only=True, verbose=1)\n",
        "\n",
        "callbacks_list = [csv_logger, reduce_LR, earlystopping, checkpoint]\n",
        "\n",
        "# Train model\n",
        "results = model.fit_generator(XYaugmentGenerator(X_train,Y_train,seed, batch_size),\n",
        "                              steps_per_epoch = np.ceil(float(len(X_train))/float(batch_size)),\n",
        "                              validation_data = val_datagen.flow(X_val,Y_val,batch_size),\n",
        "                              validation_steps = np.ceil(float(len(X_val))/float(batch_size)),\n",
        "                              shuffle = True,\n",
        "                              epochs = n_epochs,\n",
        "                              callbacks = callbacks_list)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "This round is done.\n",
        "You are ready for the next one!"
      ],
      "metadata": {
        "id": "LbtvwTZk8e1x"
      }
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "machine_shape": "hm",
      "provenance": [],
      "gpuType": "A100",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}