# ---
# 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
# ---

# + [markdown] colab_type="text" id="SKD7jSvIn1T9"
# # Importing Modules
#
# The necessary modules are : os, opencv, numpy, tqdm, matplotlib, imgaug, tensorflow and sklearn

# + colab={"base_uri": "https://localhost:8080/", "height": 34} colab_type="code" id="nNtz_YYVnzcZ" outputId="21000b0d-93c5-4dcf-96f3-4bedc5d90a3c"
import os
import cv2
import numpy as np
from tqdm import tqdm
import matplotlib.pyplot as plt

import imgaug.augmenters as iaa
from imgaug.augmentables.segmaps import SegmentationMapsOnImage

from tensorflow.keras.layers import Input, Conv2D, MaxPooling2D, Conv2DTranspose, concatenate, BatchNormalization, Activation, add
from tensorflow.keras.models import Model, model_from_json
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.layers import ELU, LeakyReLU
from tensorflow.keras.utils import plot_model
from tensorflow.keras import backend as K 

from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report

# + [markdown] colab_type="text" id="xLkeS7Lxo_Po"
# # Constructing Training and Test Datasets

# + [markdown] colab_type="text" id="gHo6ymwFpJzY"
# ## Loading the Images
#
# We first load all the images and the corresponding segmentation masks. 
#
# They are stored in two lists X, Y and respectively
#
# Moreover, the images are resized to 64x64

# + colab={"base_uri": "https://localhost:8080/", "height": 68} colab_type="code" id="YbeQ1aytHTm3" outputId="d43c1e52-d889-4a4f-96aa-5ef4d20b649f"
img_files = next(os.walk('data/cxr'))[2]
msk_files = next(os.walk('data/masks'))[2]

img_files.sort()
msk_files.sort()

print(len(img_files))
print(len(msk_files))




X = []
Y = []

for img_fl in tqdm(img_files):    
    if(img_fl.split('.')[-1]=='png'):
                
        
        img = cv2.imread('data/cxr/{}'.format(img_fl), cv2.IMREAD_GRAYSCALE)
        resized_img = cv2.resize(img,(64, 64), interpolation = cv2.INTER_CUBIC)


        X.append(resized_img)

        msk = cv2.imread('data/masks/{}'.format(img_fl), cv2.IMREAD_GRAYSCALE)
        resized_msk = cv2.resize(msk,(64, 64), interpolation = cv2.INTER_CUBIC)

        Y.append(resized_msk)
# -

# ## Data augmentation
#
# We generate more data for training using the original dataset.
#
# Operations:
#
# * dropout pixels
# * change scale
# * rotate
# * horizontal flip
# * crop
# * contrast
# * and others

# +
imgs = []
masks = []

for i in range(len(X)):
    # Load an example image (uint8, 64x64).
    image = np.array(X[i], dtype=np.uint8)

    # Define an example segmentation map (int32, 64x64).
    segmap = np.array(Y[i])
    segmap = segmap / 255
    segmap = np.around(segmap, decimals=0)
    segmap = np.array(segmap, dtype=np.int32)

    segmap = SegmentationMapsOnImage(segmap, shape=image.shape)

    # Define our augmentation pipeline.
    seq = iaa.Sequential([
        iaa.Dropout([0.0, 0.004]),      # drop 0% to 0.4% of all pixels
        iaa.Sharpen((0.0, 1.0)),       # sharpen the image
        iaa.Affine(
            scale={"x": (0.9, 1.1), "y": (0.9, 1.1)},
            translate_percent={"x": (-0.05, 0.05), "y": (-0.05, 0.05)},
            rotate=(-10, 10),
            shear=(-8, 8)
        ),
        iaa.ElasticTransformation(alpha=0.1, sigma=0.03),
        iaa.Fliplr(0.5),
        iaa.LinearContrast((0.75, 1.5)),
        iaa.Crop(percent=(0, 0.05)),

        iaa.Sometimes(
            0.5,
            iaa.GaussianBlur(sigma=(0, 0.5))
        ),
    ], random_order=True)

    # Augment images and segmaps.
    images_aug = []
    segmaps_aug = []
    for _ in range(15): #e.g. range(15) = 15 new image are generate per 1 original image
        images_aug_i, segmaps_aug_i = seq(image=image, segmentation_maps=segmap)
        images_aug.append(images_aug_i)
        segmaps_aug.append(segmaps_aug_i)


    for image_aug, segmap_aug in zip(images_aug, segmaps_aug):
        imgs.append(image_aug)  
        
        mask = cv2.cvtColor(segmap_aug.draw(size=image_aug.shape[:2])[0], cv2.COLOR_BGR2GRAY)
        mask = cv2.medianBlur(mask, 5)
        masks.append(mask)

# + [markdown] colab_type="text" id="jYgEz9HQpgsR"
# ## Train-Test Split
#
# The X, Y lists are converted to numpy arrays for convenience. 
# Furthermore, the images are divided by 255 to bring down the pixel values to [0...1] range. On the other hand the segmentations masks are converted to binary (0 or 1) values.
#
# Using Sklearn *train_test_split* we split the data randomly into 80% training and 20% testing data

# + colab={"base_uri": "https://localhost:8080/", "height": 139} colab_type="code" id="CxFIZv3715Dt" outputId="5a93db60-ed50-4122-8aa0-b26e42a601ca"
print(len(imgs))
print(len(masks))

X = np.array(imgs)
Y = np.array(masks)


X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.1, random_state=3)

Y_train = Y_train.reshape((Y_train.shape[0],Y_train.shape[1],Y_train.shape[2],1))
Y_test = Y_test.reshape((Y_test.shape[0],Y_test.shape[1],Y_test.shape[2],1))

X_train = X_train / 255
X_test = X_test / 255

Y_train = np.where(Y_train == 0, 0, 1)
Y_test = np.where(Y_test == 0, 0, 1)

print(X_train.shape)
print(Y_train.shape)
print(X_test.shape)
print(Y_test.shape)


# + [markdown] colab_type="text" id="C-Ajp2QVrMti"
# # MultiResUNet Model

# + [markdown] colab_type="text" id="21cbmiojrYrU"
# ## Model Definition
#
# The MultiResUNet model as described in the [paper](https://arxiv.org/abs/1902.04049) can be found  [here](https://github.com/nibtehaz/MultiResUNet/blob/master/MultiResUNet.py)

# + colab={} colab_type="code" id="2nX7I1Wf_zEy"
def conv2d_bn(x, filters, num_row, num_col, padding='same', strides=(1, 1), activation='relu', name=None):
    '''
    2D Convolutional layers
    
    Arguments:
        x {keras layer} -- input layer 
        filters {int} -- number of filters
        num_row {int} -- number of rows in filters
        num_col {int} -- number of columns in filters
    
    Keyword Arguments:
        padding {str} -- mode of padding (default: {'same'})
        strides {tuple} -- stride of convolution operation (default: {(1, 1)})
        activation {str} -- activation function (default: {'relu'})
        name {str} -- name of the layer (default: {None})
    
    Returns:
        [keras layer] -- [output layer]
    '''

    x = Conv2D(filters, (num_row, num_col), strides=strides, padding=padding, use_bias=False)(x)
    x = BatchNormalization(axis=3, scale=False)(x)

    if(activation == None):
        return x

    x = Activation(activation, name=name)(x)

    return x


def trans_conv2d_bn(x, filters, num_row, num_col, padding='same', strides=(2, 2), name=None):
    '''
    2D Transposed Convolutional layers
    
    Arguments:
        x {keras layer} -- input layer 
        filters {int} -- number of filters
        num_row {int} -- number of rows in filters
        num_col {int} -- number of columns in filters
    
    Keyword Arguments:
        padding {str} -- mode of padding (default: {'same'})
        strides {tuple} -- stride of convolution operation (default: {(2, 2)})
        name {str} -- name of the layer (default: {None})
    
    Returns:
        [keras layer] -- [output layer]
    '''

    x = Conv2DTranspose(filters, (num_row, num_col), strides=strides, padding=padding)(x)
    x = BatchNormalization(axis=3, scale=False)(x)
    
    return x


def MultiResBlock(U, inp, alpha = 1.67):
    '''
    MultiRes Block
    
    Arguments:
        U {int} -- Number of filters in a corrsponding UNet stage
        inp {keras layer} -- input layer 
    
    Returns:
        [keras layer] -- [output layer]
    '''

    W = alpha * U

    shortcut = inp

    shortcut = conv2d_bn(shortcut, int(W*0.167) + int(W*0.333) +
                         int(W*0.5), 1, 1, activation=None, padding='same')

    conv3x3 = conv2d_bn(inp, int(W*0.167), 3, 3,
                        activation='relu', padding='same')

    conv5x5 = conv2d_bn(conv3x3, int(W*0.333), 3, 3,
                        activation='relu', padding='same')

    conv7x7 = conv2d_bn(conv5x5, int(W*0.5), 3, 3,
                        activation='relu', padding='same')

    out = concatenate([conv3x3, conv5x5, conv7x7], axis=3)
    out = BatchNormalization(axis=3)(out)

    out = add([shortcut, out])
    out = Activation('relu')(out)
    out = BatchNormalization(axis=3)(out)

    return out


def ResPath(filters, length, inp):
    '''
    ResPath
    
    Arguments:
        filters {int} -- [description]
        length {int} -- length of ResPath
        inp {keras layer} -- input layer 
    
    Returns:
        [keras layer] -- [output layer]
    '''


    shortcut = inp
    shortcut = conv2d_bn(shortcut, filters, 1, 1,
                         activation=None, padding='same')

    out = conv2d_bn(inp, filters, 3, 3, activation='relu', padding='same')

    out = add([shortcut, out])
    out = Activation('relu')(out)
    out = BatchNormalization(axis=3)(out)

    for i in range(length-1):

        shortcut = out
        shortcut = conv2d_bn(shortcut, filters, 1, 1,
                             activation=None, padding='same')

        out = conv2d_bn(out, filters, 3, 3, activation='relu', padding='same')

        out = add([shortcut, out])
        out = Activation('relu')(out)
        out = BatchNormalization(axis=3)(out)

    return out


def MultiResUnet(height, width, n_channels):
    '''
    MultiResUNet
    
    Arguments:
        height {int} -- height of image 
        width {int} -- width of image 
        n_channels {int} -- number of channels in image
    
    Returns:
        [keras model] -- MultiResUNet model
    '''


    inputs = Input((height, width, n_channels))

    mresblock1 = MultiResBlock(32, inputs)
    pool1 = MaxPooling2D(pool_size=(2, 2))(mresblock1)
    mresblock1 = ResPath(32, 4, mresblock1)

    mresblock2 = MultiResBlock(32*2, pool1)
    pool2 = MaxPooling2D(pool_size=(2, 2))(mresblock2)
    mresblock2 = ResPath(32*2, 3, mresblock2)

    mresblock3 = MultiResBlock(32*4, pool2)
    pool3 = MaxPooling2D(pool_size=(2, 2))(mresblock3)
    mresblock3 = ResPath(32*4, 2, mresblock3)

    mresblock4 = MultiResBlock(32*8, pool3)
    pool4 = MaxPooling2D(pool_size=(2, 2))(mresblock4)
    mresblock4 = ResPath(32*8, 1, mresblock4)

    mresblock5 = MultiResBlock(32*16, pool4)

    up6 = concatenate([Conv2DTranspose(
        32*8, (2, 2), strides=(2, 2), padding='same')(mresblock5), mresblock4], axis=3)
    mresblock6 = MultiResBlock(32*8, up6)

    up7 = concatenate([Conv2DTranspose(
        32*4, (2, 2), strides=(2, 2), padding='same')(mresblock6), mresblock3], axis=3)
    mresblock7 = MultiResBlock(32*4, up7)

    up8 = concatenate([Conv2DTranspose(
        32*2, (2, 2), strides=(2, 2), padding='same')(mresblock7), mresblock2], axis=3)
    mresblock8 = MultiResBlock(32*2, up8)

    up9 = concatenate([Conv2DTranspose(32, (2, 2), strides=(
        2, 2), padding='same')(mresblock8), mresblock1], axis=3)
    mresblock9 = MultiResBlock(32, up9)

    conv10 = conv2d_bn(mresblock9, 1, 1, 1, activation='sigmoid')
    
    model = Model(inputs=[inputs], outputs=[conv10])

    return model


# + [markdown] colab_type="text" id="2frZlpmFsv1f"
# ## Auxiliary Functions

# + [markdown] colab_type="text" id="degBGBWYsyNG"
# ### Custom Metrics
#
# Since Keras does not have build-in support for computing Dice Coefficient or Jaccard Index (at the time of writing), the following functions are declared

# + colab={} colab_type="code" id="xq8qfLqDA6q2"
def dice_coef(y_true, y_pred):
    smooth = 0.0
    y_true_f = K.flatten(y_true)
    y_pred_f = K.flatten(y_pred)
    intersection = K.sum(y_true_f * y_pred_f)
    return (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)

def jacard(y_true, y_pred):

    y_true_f = K.flatten(y_true)
    y_pred_f = K.flatten(y_pred)
    intersection = K.sum ( y_true_f * y_pred_f)
    union = K.sum ( y_true_f + y_pred_f - y_true_f * y_pred_f)

    return intersection/union


# + [markdown] colab_type="text" id="VtvyeBXy8Mk3"
# ### Saving Model 
#
# Function to save the model

# + colab={} colab_type="code" id="GCfX9MUYALar"
def saveModel(model):

    model_json = model.to_json()

    try:
        os.makedirs('models')
    except:
        pass
    
    fp = open('models/modelP.json','w')
    fp.write(model_json)
    model.save_weights('models/modelW.h5')



# + [markdown] colab_type="text" id="AH3znjx-8vDq"
# ### Evaluate the Model
#
# We evaluate the model on test data (X_test, Y_test). 
#
# We compute the values of Jaccard Index and Dice Coeficient, and save the predicted segmentation of first 10 images. The best model is also saved
#
# (This could have been done using keras call-backs as well)

# + colab={} colab_type="code" id="Tkit_YYvBQ7V"
def evaluateModel(model, X_test, Y_test, batchSize):
    
    try:
        os.makedirs('results')
    except:
        pass 
    

    yp = model.predict(x=X_test, batch_size=batchSize, verbose=1)

    yp = np.round(yp,0)

    for i in range(10):

        plt.figure(figsize=(20,10))
        plt.subplot(1,3,1)
        plt.imshow(X_test[i])
        plt.title('Input')
        plt.subplot(1,3,2)
        plt.imshow(Y_test[i].reshape(Y_test[i].shape[0],Y_test[i].shape[1]))
        plt.title('Ground Truth')
        plt.subplot(1,3,3)
        plt.imshow(yp[i].reshape(yp[i].shape[0],yp[i].shape[1]))
        plt.title('Prediction')

        intersection = yp[i].ravel() * Y_test[i].ravel()
        union = yp[i].ravel() + Y_test[i].ravel() - intersection

        jacard = (np.sum(intersection)/np.sum(union))  
        plt.suptitle('Jacard Index'+ str(np.sum(intersection)) +'/'+ str(np.sum(union)) +'='+str(jacard))

        plt.savefig('results/'+str(i)+'.png',format='png')
        plt.close()


    jacard = 0
    dice = 0
    
    
    for i in range(len(Y_test)):
        yp_2 = yp[i].ravel()
        y2 = Y_test[i].ravel()
        
        intersection = yp_2 * y2
        union = yp_2 + y2 - intersection

        jacard += (np.sum(intersection)/np.sum(union))  

        dice += (2. * np.sum(intersection) ) / (np.sum(yp_2) + np.sum(y2))

    
    jacard /= len(Y_test)
    dice /= len(Y_test)
    


    print('Jacard Index : '+str(jacard))
    print('Dice Coefficient : '+str(dice))
    

    fp = open('models/log.txt','a')
    fp.write(str(jacard)+'\n')
    fp.close()

    fp = open('models/best.txt','r')
    best = fp.read()
    fp.close()

    if(jacard>float(best)):
        print('***********************************************')
        print('Jacard Index improved from '+str(best)+' to '+str(jacard))
        print('***********************************************')
        fp = open('models/best.txt','w')
        fp.write(str(jacard))
        fp.close()

        saveModel(model)


# + [markdown] colab_type="text" id="80IwnKtM9NHY"
# ### Training the Model
#
# The model is trained and evaluated after each epochs

# + colab={} colab_type="code" id="wlOOh0-nA05L"
def trainStep(model, X_train, Y_train, X_test, Y_test, epochs, batchSize):

    
    for epoch in range(epochs):
        print('Epoch : {}'.format(epoch+1))
        model.fit(x=X_train, y=Y_train, batch_size=batchSize, epochs=1, verbose=1)     

        evaluateModel(model,X_test, Y_test,batchSize)

    return model 

# + [markdown] colab_type="text" id="4wJ5V9AJ9Ygu"
# ## Define Model, Train and Evaluate

# + colab={"base_uri": "https://localhost:8080/", "height": 1105} colab_type="code" id="GJqeuZPhDZSK" outputId="b1344374-f69e-4805-a22d-b8152f07b754"
model = MultiResUnet(height=64, width=64, n_channels=1)

model.compile(optimizer='adam', loss='binary_crossentropy', metrics=[dice_coef, jacard, 'accuracy'])

saveModel(model)

fp = open('models/log.txt','w')
fp.close()
fp = open('models/best.txt','w')
fp.write('-1.0')
fp.close()
    
trainStep(model, X_train, Y_train, X_test, Y_test, epochs=35, batchSize=32)

StarcoderdataJupyterScriptsDedupFiltered