# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + id="DyhbDNfq4ulu"
# from google.colab import drive
# drive.mount("/content/drive", force_remount=True)

# + colab={"base_uri": "https://localhost:8080/"} id="lbnbzSwN40b7" outputId="ae220976-71ff-43b6-a8c0-4c359dca45a3"
# %cd '/content/drive/MyDrive/Cov19/input'

# + id="jF9jsZwY4153"
import pandas as pd
import numpy as np
import scipy
import matplotlib.pyplot as plt
import seaborn
import os as os
import cv2 as cv
import glob as glob
import nibabel as nib
import pickle
import imgaug as ia
import imgaug.augmenters as iaa
import tqdm.notebook as tqdm
import gc
import warnings
import tensorflow as tf
from keras import backend as K
from keras import losses, metrics
from keras import optimizers
from keras import callbacks
from keras.models import Model
from keras.layers import Input, BatchNormalization, Activation, Dense, Dropout
from keras.layers import concatenate, Conv2D, MaxPooling2D, Conv2DTranspose, add, multiply
from keras.layers import Multiply, UpSampling2D, core
from keras.layers.merge import concatenate
from keras.layers.core import Lambda
from sklearn.model_selection import train_test_split
from sklearn.cluster import KMeans
from skimage import morphology as morph
from skimage import measure

# + [markdown] id="fRnstAY55Rjk"
# # 1. Loading data
#

# + id="6ClKVctx43gh"
infile = open('Processed_Data.cp','rb')
data_dict = pickle.load(infile)
all_cts = data_dict['cts']
all_lungs = data_dict['lungs']
all_inf = data_dict['infects']
infile.close()

# + colab={"base_uri": "https://localhost:8080/"} id="UeUflHhi5t8G" outputId="d097b69e-f034-4fb4-cacb-a33627e0649b"
print(all_cts.shape)
print(all_lungs.shape)
print(all_inf.shape)

# + id="JuHjBH6rY4jy"
from sklearn.utils import shuffle
all_cts, all_lungs, all_inf = shuffle(all_cts, all_lungs, all_inf)

# + colab={"base_uri": "https://localhost:8080/"} id="QMOfAhVuY-HC" outputId="d337f9cb-d57c-46ab-ada6-8c5e0377c6ff"
all_cts = (all_cts - all_cts.min())/(all_cts.max()-all_cts.min())
all_lungs = (all_lungs - all_lungs.min())/(all_lungs.max()-all_lungs.min())
all_inf = (all_inf - all_inf.min())/(all_inf.max()-all_inf.min())

print("{} {}".format(all_cts.min(), all_cts.max()))
print("{} {}".format(all_lungs.min(), all_lungs.max()))
print("{} {}".format(all_inf.min(), all_inf.max()))

# + [markdown] id="qmWCu4QuuB9J"
# # 2. Splitting data into training and validation sets

# + colab={"base_uri": "https://localhost:8080/"} id="kF0uPJAYuIcH" outputId="d5aa2317-133e-41a6-8e33-e8c39b84b6b8"
train_size = int(0.65*all_cts.shape[0])
X_train, yl_train, yi_train = (all_cts[:train_size]/255, all_lungs[:train_size], all_inf[:train_size])
X_valid, yl_valid, yi_valid = (all_cts[train_size:int(0.8*all_cts.shape[0])]/255, 
                               all_lungs[train_size:int(0.8*all_cts.shape[0])], 
                               all_inf[train_size:int(0.8*all_cts.shape[0])])
test_size = int(0.8*all_cts.shape[0])
X_test, yl_test, yi_test = (all_cts[test_size:]/255, all_lungs[test_size:], all_inf[test_size:])
print(X_train.shape, yl_train.shape, yi_train.shape)
print(X_valid.shape, yl_valid.shape, yi_valid.shape)
print(X_test.shape, yl_test.shape, yi_test.shape)


# + [markdown] id="_CbeWd38rouS"
# # 3. Define evaluation metrics
# We will use dice coefficient as a metric for infection segmentation. Dice coefficient is (2TP/2TP+FN+FP) where TP, FN and FP correspond to true positive, false negative and false positive. Code taken from
# https://medium.com/@karan_jakhar/100-days-of-code-day-7-84e4918cb72c

# + id="kj1Gxu9r53Ee"
def dice(y_true, y_pred, smooth=1):
  intersection = K.sum(y_true * y_pred, axis=[1,2,3])
  union = K.sum(y_true, axis=[1,2,3]) + K.sum(y_pred, axis=[1,2,3])
  dice = K.mean((2. * intersection + smooth)/(union + smooth), axis=0)
  return dice

def dice_loss(y_true, y_pred):
  loss = 1 - dice(y_true, y_pred)
  return loss

def bce_dice_loss(y_true, y_pred):
  loss = 0.5*losses.binary_crossentropy(y_true, y_pred) + 0.5*dice_loss(y_true, y_pred)
  return loss

class CosineAnnealingLearningRateSchedule(callbacks.Callback):
  def __init__(self, n_epochs, n_cycles, lrate_max, verbose=0):
    self.epochs = n_epochs
    self.cycles = n_cycles
    self.lr_max = lrate_max
    self.lrates = list()

  def cosine_annealing(self, epoch, n_epochs, n_cycles, lrate_max):
    epochs_per_cycle = np.floor(n_epochs/n_cycles)
    cos_inner = (np.pi * (epoch % epochs_per_cycle)) / (epochs_per_cycle)
    return lrate_max/2 * (np.cos(cos_inner) + 1)

  def on_epoch_begin(self, epoch, logs=None):
    lr = self.cosine_annealing(epoch, self.epochs, self.cycles, self.lr_max)
    K.set_value(self.model.optimizer.lr, lr)
    self.lrates.append(lr)


# + [markdown] id="W5AXZrg0wqS6"
# # 4. CNN
#

# + id="zS82kxUkATB0"
def block1 (input_shape, filtersize, poolsz=(2,2)) :
    x = Conv2D(filtersize, (3,3), activation='relu', padding='same', kernel_initializer="he_normal") (input_shape)
    x = Conv2D(filtersize, (3,3), activation='relu', padding='same', kernel_initializer="he_normal") (x)
    x_inter = BatchNormalization()(x)
    x = MaxPooling2D(poolsz) (x_inter) 
    x = Dropout(0.2)(x) 
    return x, x_inter

def block2 (input_shape, filtersize) :
    x = BatchNormalization() (input_shape)
    x = Conv2D(filtersize, (3, 3), activation='relu', padding='same', kernel_initializer="he_normal") (x)
    x = Conv2D(filtersize, (3, 3), activation='relu', padding='same', kernel_initializer="he_normal") (x) 
    return x


# + colab={"base_uri": "https://localhost:8080/"} id="QeEUhS-hwQJR" outputId="487709be-0f7b-4c96-e804-b06dc060d36a"
def infection_segmentation(input_shape) :
  x_input = Input(input_shape)

  x, Xa = block1(x_input, 32)
  x, Xb = block1(x, 64)
  x, _ = block1(x, 128, poolsz=(1,1))
  x, _ = block1(x, 256, poolsz=(1,1))
  x = block2(x, 256)

  x = Conv2DTranspose(128, (2, 2), strides=(2,2), padding='same') (x)
  x = block2(x, 128)

  x = Conv2DTranspose(64, (2, 2), padding='same') (x)
  x = concatenate([x, Xb])
  x = block2(x, 64)

  x = Conv2DTranspose(32, (2, 2), strides=(2,2), padding='same') (x)
  x = concatenate([x, Xa], axis=3)
  x = block2(x, 32)

  infection_segmentation = Conv2D(1, (1, 1), activation='sigmoid', name='infect_output') (x)

  model = Model(inputs=x_input, outputs=infection_segmentation, name='infect_model')
    
  return model

strategy = tf.distribute.MirroredStrategy()
print('Devices {}'.format(strategy.num_replicas_in_sync))
with strategy.scope() :
    infection_segmentation = infection_segmentation(all_cts.shape[1:])

infection_segmentation.summary()

# + id="7lIiAF5Aw_QE"
epochs = 100
lrmax = 5e-5
n_cycles = epochs / 25
lr_cb = CosineAnnealingLearningRateSchedule(epochs, n_cycles, lrmax)
checkpoint_fpath = "infection_segmentation_weights.hdf5"
cts_checkpoint_cb = callbacks.ModelCheckpoint(checkpoint_fpath, 
                                              monitor='val_dice', 
                                              save_best_only=True, 
                                              mode='max', 
                                              verbose=1,
                                              save_weights_only=True)

batch_size = 8
optim = optimizers.Adam(lr=5e-5, beta_1=0.9, beta_2=0.99)
with strategy.scope() :
    infection_segmentation.compile(optimizer=optim, loss=bce_dice_loss, metrics=[dice])

# + [markdown] id="UWCIZz9vxiS7"
# # 5. Train Model

# + colab={"base_uri": "https://localhost:8080/"} id="EY9B84-Exk7N" outputId="e691b825-f05f-4ccd-923d-96ba12536c29"
infection_segmentation_res = infection_segmentation.fit(x = X_train, 
                            y = yi_train,
                            batch_size = batch_size, 
                            epochs = epochs,
                            verbose = 1,
                            validation_data = (X_valid, yi_valid),
                            callbacks = [cts_checkpoint_cb, lr_cb])

# + colab={"base_uri": "https://localhost:8080/", "height": 337} id="hV_kk78VzJY3" outputId="0129dc54-3b8a-4ac0-cf9c-fae07cdba78f"
plt.style.use('ggplot')

fig, axes = plt.subplots(1, 2, figsize=(11,5))

axes[0].plot(infection_segmentation_res.history['dice'], color='b', label='train-infection')
axes[0].plot(infection_segmentation_res.history['val_dice'], color='r', label='valid-infection')
axes[0].set_ylabel('Dice coefficient')
axes[0].set_xlabel('Epoch')
axes[0].legend()

axes[1].plot(infection_segmentation_res.history['loss'], color='b', label='train-infection')
axes[1].plot(infection_segmentation_res.history['val_loss'], color='r', label='valid-infection')
axes[1].set_ylabel('Loss')
axes[1].set_xlabel('Epoch')
axes[1].legend();

# + [markdown] id="rKaY03cI0mVw"
# # 6. Saving loads

# + id="n716JP-wbcp1"
infection_segmentation.load_weights('infection_segmentation_weights.hdf5')
prediction = infection_segmentation.predict(X_test)


# + [markdown] id="s5m49RsN08nS"
# # 7. Testing

# + id="zEFH6wvo06VY"
def plot_lung_seg(all_cts, all_lungs, all_inf, pred_infs, axes) :

    axes[0].imshow(all_cts[:,:,0], cmap='bone')
    axes[0].set_title('CT image'); plt.grid(None)
    axes[0].set_xticks([]); axes[0].set_yticks([])
    
    axes[1].imshow(all_lungs[:,:,0], cmap='bone')
    axes[1].imshow(all_inf[:,:,0], alpha=0.5, cmap='Reds')
    axes[1].set_title('Infection mask'); plt.grid(None)
    axes[1].set_xticks([]); axes[1].set_yticks([])

    axes[2].imshow(all_lungs[:,:,0], cmap='bone')
    axes[2].imshow(pred_infs[:,:,0], alpha=0.5, cmap='Reds')
    axes[2].set_title('Pred. Infection mask'); plt.grid(None)
    axes[2].set_xticks([]); axes[2].set_yticks([])


# + colab={"base_uri": "https://localhost:8080/", "height": 541} id="gaUb86Ol1LEh" outputId="f4200be2-72d2-4a01-8f84-05c63b955dcf"
import random
indices = random.choices(range(len(X_test)), k=5)
fig, axes = plt.subplots(3, 5, figsize=(15,9))

for ii, idx in enumerate(indices) :
    plot_lung_seg(X_test[idx], yl_test[idx], yi_test[idx], prediction[idx], list(axes[:,ii]))

StarcoderdataJupyterScriptsDedupFiltered