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

# + _cell_guid="b1076dfc-b9ad-4769-8c92-a6c4dae69d19" _uuid="8f2839f25d086af736a60e9eeb907d3b93b6e0e5" colab={"base_uri": "https://localhost:8080/", "height": 73} colab_type="code" id="ZupQDr9J-icw" outputId="615c2281-5a37-4c21-f9ba-5a2a0d8ff70e"
import pandas as pd
import numpy as np
import scipy
import matplotlib.pyplot as plt
import seaborn
import cv2 as cv
import nibabel as nib
import pickle
import glob
import imgaug as ia
import imgaug.augmenters as iaa
import tqdm
import gc, os
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
from keras.layers import Multiply, UpSampling2D
from sklearn.model_selection import train_test_split
from sklearn.cluster import KMeans
from skimage import morphology
from skimage import measure
#import keras_segmentation as ks

warnings.filterwarnings('ignore')
# %matplotlib inline
print("Version: ", tf.version.VERSION)
physical_devices = tf.config.list_physical_devices()
print(physical_devices)

# + [markdown] colab_type="text" id="26-GKjVy-2d6"
# ## 1. Loading data

# +
datadir = '../training-dir/'

glob_search = os.path.join(datadir, 'patient*')
train_files = sorted(glob.glob(glob_search))

print('num of train patients {}'.format(len(train_files)))

# +
img_size = 128

cts_all = []
lungs_all = []
infects_all = []

for fnum in tqdm.tqdm(range(len(train_files))) :
    glob_search = os.path.join(train_files[fnum], '*ctscan*.nii')
    ct_file = glob.glob(glob_search)
    cts = nib.load(ct_file[0])
    data_cts = cts.get_fdata()

    glob_search = os.path.join(train_files[fnum], '*lung*.nii')
    lung_file = glob.glob(glob_search)
    lung_masks = nib.load(lung_file[0])
    data_lungs = lung_masks.get_fdata()

    glob_search = os.path.join(train_files[fnum], '*infection*.nii')
    infect_file = glob.glob(glob_search)
    infect_masks = nib.load(infect_file[0])
    data_infects = infect_masks.get_fdata()
    
    height, width, slices = data_cts.shape
    sel_slices = range(round(slices*0.2), round(slices*0.8))

    data_cts = np.rot90(np.array(data_cts))
    data_lungs = np.rot90(np.array(data_lungs))
    data_infects = np.rot90(np.array(data_infects))
    
    data_cts = np.reshape(np.rollaxis(data_cts, 2),(slices,height,width))
    data_lungs = np.reshape(np.rollaxis(data_lungs, 2),(slices,height,width))
    data_infects = np.reshape(np.rollaxis(data_infects, 2),(slices,height,width))

    data_cts = data_cts[sel_slices, :, :]
    data_lungs = data_lungs[sel_slices, :, :]
    data_infects = data_infects[sel_slices, :, :]
    
    for ii in range(data_cts.shape[0]):
        img = cv.resize(data_cts[ii], dsize=(img_size, img_size), interpolation=cv.INTER_AREA)
        img = np.reshape(img, (img_size, img_size, 1))
        cts_all.append(img)
        
        img = cv.resize(data_lungs[ii], dsize=(img_size, img_size), interpolation=cv.INTER_AREA)
        img = np.reshape(img, (img_size, img_size, 1))
        lungs_all.append(img)

        img = cv.resize(data_infects[ii], dsize=(img_size, img_size), interpolation=cv.INTER_AREA)
        img = np.reshape(img, (img_size, img_size, 1))
        infects_all.append(img)
# -

print(len(cts_all))

# + [markdown] colab_type="text" id="BOj54mfnQ9eT"
# # 2. Data Augmentation

# + colab={} colab_type="code" id="rNld7hkEy7qj"
ia.seed(1)

seq = iaa.Sequential([
    iaa.Fliplr(0.5), # horizontal flips
    iaa.Flipud(0.5), # vertical flips
    
    # Apply affine transformations to each image.
    # Scale/zoom them, translate/move them, rotate them and shear them.
    iaa.Affine(
        scale={"x": (0.8, 1.2), "y": (0.8, 1.2)}, # scale images
        translate_percent={"x": (-0.2, 0.2), "y": (-0.2, 0.2)},
        rotate=(-15, 15),
        shear=(-15, 15)
    )
], random_order=True) # apply augmenters in random order

# + colab={"base_uri": "https://localhost:8080/", "height": 1000} colab_type="code" id="mmUWCRR1y7t9" outputId="6fa91350-1420-48fa-9f60-339e24def85d"
num_augs = round(len(cts_all))
rand_idx = np.random.randint(0, len(cts_all), size=num_augs)

sample_cts = [cts_all[ii] for ii in rand_idx]
sample_lungs = [lungs_all[ii] for ii in rand_idx]
sample_infects = [infects_all[ii] for ii in rand_idx]
# -

seq_det = seq.to_deterministic()
cts_aug = seq_det.augment_images(sample_cts)
lungs_aug = seq_det.augment_images(sample_lungs)
infects_aug = seq_det.augment_images(sample_infects)

with tf.device('/cpu:0') :
    cts_all = tf.convert_to_tensor(cts_all)
    cts_aug = tf.convert_to_tensor(np.asarray(cts_aug))
    lungs_all = tf.convert_to_tensor(np.asarray(lungs_all))
    lungs_aug = tf.convert_to_tensor(np.asarray(lungs_aug))
    infects_all = tf.convert_to_tensor(np.asarray(infects_all))
    infects_aug = tf.convert_to_tensor(np.asarray(infects_aug))

    cts = tf.concat([cts_all, cts_aug], axis=0)
    lungs = tf.concat([lungs_all, lungs_aug], axis=0)
    infects = tf.concat([infects_all, infects_aug], axis=0)

# +
with tf.device('/cpu:0') :
    indices = tf.range(start=0, limit=cts.shape[0], dtype=tf.int32)
    shuffled_indices = tf.random.shuffle(indices)

    cts = tf.gather(cts, shuffled_indices)
    lungs = tf.gather(lungs, shuffled_indices)
    infects = tf.gather(infects, shuffled_indices)

print(cts.shape, lungs.shape, infects.shape)
# -

with open('../outputs/augmented_data.cp', 'wb') as myfile : 
    pickle.dump({'cts': cts, 'lungs': lungs, 'infects': infects}, myfile)


# # 3. Defining evaluation metrics
#
# The most commonly used metrics for image segmentation are the IoU and the Dice Coefficient. The Dice Coefficient is 2 times the Area of overlap divided by the total number of pixels in both images (true and predicted). The Dice coefficient is very similar to the IoU, and they both range from 0 to 1, with 1 signifying the greatest similarity between predicted and truth.
#
# See https://github.com/keras-team/keras/issues/9395
#

# + colab={} colab_type="code" id="EAI13B4H9F_X"
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): #Binary Cross-Entropy
    loss = 0.5*losses.binary_crossentropy(y_true, y_pred) + 0.5*dice_loss(y_true, y_pred)
    return loss

def tversky_loss(y_true, y_pred):
    alpha, beta  = 0.5, 0.5
    ones = K.ones(K.shape(y_true))
    p0 = y_pred
    p1 = ones-y_pred
    g0 = y_true
    g1 = ones-y_true
    num = K.sum(p0*g0, (0,1,2))
    den = num + alpha*K.sum(p0*g1,(0,1,2)) + beta*K.sum(p1*g0,(0,1,2))
    T = K.sum(num/den) 
    Ncl = K.cast(K.shape(y_true)[-1], 'float32')
    return Ncl-T

def weighted_bce_loss(y_true, y_pred, weight):
    epsilon = 1e-7
    y_pred = K.clip(y_pred, epsilon, 1. - epsilon)
    logit_y_pred = K.log(y_pred / (1. - y_pred))
    loss = weight * (logit_y_pred * (1. - y_true) + 
                     K.log(1. + K.exp(-K.abs(logit_y_pred))) + 
                     K.maximum(-logit_y_pred, 0.))
    return K.sum(loss) / K.sum(weight)

def weighted_dice_loss(y_true, y_pred, weight):
    smooth = 1.
    w, m1, m2 = weight, y_true, y_pred
    intersection = (m1*m2)
    score = (2.*K.sum(w*intersection) + smooth) / (K.sum(w*m1) + K.sum(w*m2) + smooth)
    loss = 1. - K.sum(score)
    return loss

def weighted_bce_dice_loss(y_true, y_pred):
    y_true = K.cast(y_true, 'float32')
    y_pred = K.cast(y_pred, 'float32')
    averaged_mask = K.pool2d(y_true, pool_size=(50, 50), strides=(1, 1), 
                             padding='same', pool_mode='avg')
    weight = K.ones_like(averaged_mask)
    w0 = K.sum(weight)
    weight = 5. * K.exp(-5.*K.abs(averaged_mask - 0.5))
    w1 = K.sum(weight)
    weight *= (w0 / w1)
    loss = 0.5*weighted_bce_loss(y_true, y_pred, weight) + 0.5*dice_loss(y_true, y_pred)
    return loss


# -

# # 4. Cosine Annealing Learning Rate
#
# An effective snapshot ensemble requires training a neural network with an aggressive learning rate schedule.
#
# The cosine annealing schedule is an example of an aggressive learning rate schedule where learning rate starts high and is dropped relatively rapidly to a minimum value near zero before being increased again to the maximum.
#
# We can implement the schedule as described in the 2017 paper “Snapshot Ensembles: Train 1, get M for free.” The equation requires the total training epochs, maximum learning rate, and number of cycles as arguments as well as the current epoch number. The function then returns the learning rate for the given epoch.
#
# See https://machinelearningmastery.com/snapshot-ensemble-deep-learning-neural-network/

# + colab={} colab_type="code" id="H14hooGX-ct_"
# define custom learning rate schedule
class CosineAnnealingLearningRateSchedule(callbacks.Callback):
    # constructor
    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()
 
    # calculate learning rate for an epoch
    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)
 
    # calculate and set learning rate at the start of the epoch
    def on_epoch_begin(self, epoch, logs=None):
        # calculate learning rate
        lr = self.cosine_annealing(epoch, self.epochs, self.cycles, self.lr_max)
        # set learning rate
        K.set_value(self.model.optimizer.lr, lr)
        # log value
        self.lrates.append(lr)


# + [markdown] colab_type="text" id="sadUwQdyVXZ4"
# * All the hyperparameters are put in place after repeating trial and error for a fixed number of epochs.
# -

# # 5. Split data into train and validation sets

# +
train_size = int(0.9*cts.shape[0])

with tf.device('/cpu:0') :
    X_train, yl_train, yi_train = (cts[:train_size]/255, 
                                   lungs[:train_size], 
                                   infects[:train_size])

    X_valid, yl_valid, yi_valid = (cts[train_size:]/255, 
                                   lungs[train_size:],
                                   infects[train_size:])

print(X_train.shape, yl_train.shape, yi_train.shape)
print(X_valid.shape, yl_valid.shape, yi_valid.shape)


# -

# # 6. Convolutional Neural Network

# +
def lung_seg(input_shape, num_filters=[16,32,128]) :
    
    x_input = Input(input_shape)
    
    ###  LUNG SEGMENTATION
    x = Conv2D(num_filters[0], kernel_size=3, activation='relu', padding='same')(x_input)
    x = MaxPooling2D(pool_size=2, padding='same')(x)

    x = Conv2D(num_filters[1], kernel_size=3, activation='relu', padding='same')(x)
    x = MaxPooling2D(pool_size=2, padding='same')(x)

    x = Conv2D(num_filters[2], kernel_size=3, activation='relu', padding='same')(x)
    x = MaxPooling2D(pool_size=2, padding='same')(x)

    x = Dense(num_filters[2], activation='relu')(x)

    x = UpSampling2D(size=2)(x)
    x = Conv2D(num_filters[2], kernel_size=3, activation='sigmoid', padding='same')(x)

    x = UpSampling2D(size=2)(x)
    x = Conv2D(num_filters[1], kernel_size=3, activation='sigmoid', padding='same')(x)

    x = UpSampling2D(size=2)(x)
    lung_seg = Conv2D(1, kernel_size=3, activation='sigmoid', padding='same')(x) # identifying lungs

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

strategy = tf.distribute.MirroredStrategy()
print('Number of devices {}'.format(strategy.num_replicas_in_sync))

with strategy.scope() :
    lung_seg = lung_seg(cts.shape[1:])

lung_seg.summary()

# + colab={} colab_type="code" id="R1WJkDS6_WPx"
# define learning rate callback
epochs = 100
lrmax = 5e-5
n_cycles = epochs / 25
lr_cb = CosineAnnealingLearningRateSchedule(epochs, n_cycles, lrmax)

checkpoint_fpath = "../outputs/weights_lungseg.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() :
    lung_seg.compile(optimizer=optim, loss=bce_dice_loss, metrics=[dice])
# -

# # 7. Training

lung_seg_res = lung_seg.fit(x = X_train, 
                            y = yl_train,
                            batch_size = batch_size, 
                            epochs = epochs,
                            verbose = 1,
                            validation_data = (X_valid, yl_valid),
                            callbacks = [cts_checkpoint_cb, lr_cb])

# + colab={"base_uri": "https://localhost:8080/", "height": 813, "resources": {"http://localhost:17327/content/unet_covid_fold1.hdf5": {"data": "<KEY>", "headers": [["content-length", "1461"], ["content-type", "text/html; charset=utf-8"]], "ok": false, "status": 500, "status_text": ""}}} colab_type="code" id="HddKr4NwNq5K" outputId="c6d719ab-98ca-43f1-f406-91c5a1370e56"
plt.style.use('ggplot')

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

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

axes[1].plot(lung_seg_res.history['loss'], color='b', label='train-lung')
axes[1].plot(lung_seg_res.history['val_loss'], color='r', label='valid-lung')
axes[1].set_ylabel('loss')
axes[1].set_xlabel('epoch')
axes[1].legend();
# -

# # 8. Saving Outputs

# +
gc.collect()

with open('../outputs/lungseg-HistoryDict.cp', 'wb') as myfile:
    pickle.dump(lung_seg_res.history, myfile)

# +
lung_seg.load_weights("../outputs/weights_lungseg.hdf5")
with tf.device('/cpu:0') :
    pred_lungs = lung_seg.predict(cts/255, batch_size=512)

print(cts.shape, pred_lungs.shape)


# -

# # 9. Visualizing Outputs

def plot_lung_seg(ct, lung, pred_lung, axes) :

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

    axes[2].imshow(pred_lung[:,:,0], alpha=0.5, cmap='Reds')
    axes[2].set_title('pred Lung mask'); plt.grid(None)
    axes[2].set_xticks([]); axes[2].set_yticks([])


# +
import random
indices = random.choices(range(len(pred_lungs)), k=5)

fig, axes = plt.subplots(3, 5, figsize=(15,9))

for ii, idx in enumerate(indices) :
    plot_lung_seg(cts[idx], lungs[idx], pred_lungs[idx], list(axes[:,ii]))

StarcoderdataJupyterScriptsDedupFiltered