Medical images segmentation with Keras: U-net architecture | by Soriba D. | Towards Data Science

github.com
Medical Image Segmentation is the process of automatic or semi-automatic detection of boundaries within a 2D or 3D image
The segmentation of medical images has long been an active research subject because AI can help fight many diseases like cancer. Performing this task automatically, precisely and quickly would facilitate the word of specialists and better be able to treat patients with something reliable.
In this post, we’ll see how to automatically segment livers using the U-net architecture on an actual patients scans dataset. Hope you will enjoy it.
The data is available in NifTi format here. The NifTi (Neuroimaging Informatics Technology Initiative) is an open file format commonly used to store brain imaging data obtained using Magnetic Resonance Imaging methods. There is a Python packaged called nibabel that we’ll use to deal with this kind of data.
The dataset we’ll use consists of 20 medical examinations in 3D, it contains the source images as well as the masks of segmentation of the liver for each. This data come from IRCAD, a medical research center in France.
To carry out our task, we will obviously not use only 20 data (which is really very little in Deep Learning). In order to augment our data, we’ll use axial 2-d slices.
We’ll train a U-net architecture, a fully convolutional network. The principle of this architecture is to add to a usual contracting network layers with upsampling operators instead of pooling.
This allow the network to learn context (contracting path), then localization (expansive path). Context information is propagated to higher resolution layers thanks to skip-connexions. So we have images of the same size as input:
As metric we use the Dice coefficient (which is quite similar to the Jaccard coefficient). It is a statistical indicator that measures the similarity of two samples:
We can see clearly in this diagram: if the dice coefficient equals to 1, it signifies that the two samples we compare are exactly equal! So the closer it is to 1, the better.
Now that we have set up the framework, we can start with the project.
In your project folder create a “raw” folder:
In the ‘raw’ folder, create a “test” folder, and a “train” folder.
Then separate the data in two sets (train and test, typically we use 13 samples for the train set and 7 for the test set), and put them in the corresponding directories.
We’ll see what to write in “data.py” and “train.ipynb” in the two last parts.
The preprocessing part is crucial in deep learning tasks because neural networks expect data in a certain format.
In this part, I show you the python code that turns the data into numpy arrays that we’ll feed to the neural network. To do so, we’ll transform the 3-dimensional data into 2-d cuts. Then save them in “.npy” format in order to use them for further training.
Now that we can get our data prepared, let’s go to the last step: the training!
In this part I’ll show you the notebook to set up the architecture, train the neural network and store the results in a file (the code is inspired from this repo), don’t hesitate to refer to the comments to understand the code:
from __future__ import print_function

import os
from skimage.transform import resize
from skimage.io import imsave
import numpy as np
from skimage.segmentation import mark_boundaries
from keras.models import Model
from keras.layers import Input, concatenate, Conv2D, MaxPooling2D, Conv2DTranspose
from keras.optimizers import Adam, SGD
from keras.callbacks import ModelCheckpoint
from keras import backend as K
from skimage.exposure import rescale_intensity
from keras.callbacks import History
from skimage import io
from data import load_train_data, load_test_data

K.set_image_data_format('channels_last')  # TF dimension ordering in this code

img_rows = int(512/2)
img_cols = int(512/2)
smooth = 1.
#We divide here the number of rows and columns by two because we undersample our data (We take one pixel over two) 

def dice_coef(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)
    return (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)


def dice_coef_loss(y_true, y_pred):
    return -dice_coef(y_true, y_pred)

#The functions return our metric and loss

def get_unet():
    inputs = Input((img_rows, img_cols, 1))
    conv1 = Conv2D(32, (3, 3), activation='relu', padding='same')(inputs)
    conv1 = Conv2D(32, (3, 3), activation='relu', padding='same')(conv1)
    pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)

    conv2 = Conv2D(64, (3, 3), activation='relu', padding='same')(pool1)
    conv2 = Conv2D(64, (3, 3), activation='relu', padding='same')(conv2)
    pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)

    conv3 = Conv2D(128, (3, 3), activation='relu', padding='same')(pool2)
    conv3 = Conv2D(128, (3, 3), activation='relu', padding='same')(conv3)
    pool3 = MaxPooling2D(pool_size=(2, 2))(conv3)

    conv4 = Conv2D(256, (3, 3), activation='relu', padding='same')(pool3)
    conv4 = Conv2D(256, (3, 3), activation='relu', padding='same')(conv4)
    pool4 = MaxPooling2D(pool_size=(2, 2))(conv4)

    conv5 = Conv2D(512, (3, 3), activation='relu', padding='same')(pool4)
    conv5 = Conv2D(512, (3, 3), activation='relu', padding='same')(conv5)

    up6 = concatenate([Conv2DTranspose(256, (2, 2), strides=(2, 2), padding='same')(conv5), conv4], axis=3)
    conv6 = Conv2D(256, (3, 3), activation='relu', padding='same')(up6)
    conv6 = Conv2D(256, (3, 3), activation='relu', padding='same')(conv6)

    up7 = concatenate([Conv2DTranspose(128, (2, 2), strides=(2, 2), padding='same')(conv6), conv3], axis=3)
    conv7 = Conv2D(128, (3, 3), activation='relu', padding='same')(up7)
    conv7 = Conv2D(128, (3, 3), activation='relu', padding='same')(conv7)

    up8 = concatenate([Conv2DTranspose(64, (2, 2), strides=(2, 2), padding='same')(conv7), conv2], axis=3)
    conv8 = Conv2D(64, (3, 3), activation='relu', padding='same')(up8)
    conv8 = Conv2D(64, (3, 3), activation='relu', padding='same')(conv8)

    up9 = concatenate([Conv2DTranspose(32, (2, 2), strides=(2, 2), padding='same')(conv8), conv1], axis=3)
    conv9 = Conv2D(32, (3, 3), activation='relu', padding='same')(up9)
    conv9 = Conv2D(32, (3, 3), activation='relu', padding='same')(conv9)

    conv10 = Conv2D(1, (1, 1), activation='sigmoid')(conv9)

    model = Model(inputs=[inputs], outputs=[conv10])

    model.compile(optimizer=Adam(lr=1e-3), loss=dice_coef_loss, metrics=[dice_coef])

    return model

#The different layers in our neural network model (including convolutions, maxpooling and upsampling)

def preprocess(imgs):
    imgs_p = np.ndarray((imgs.shape[0], img_rows, img_cols), dtype=np.uint8)
    for i in range(imgs.shape[0]):
        imgs_p[i] = resize(imgs[i], (img_cols, img_rows), preserve_range=True)

    imgs_p = imgs_p[..., np.newaxis]
    return imgs_p

#We adapt here our dataset samples dimension so that we can feed it to our network

import matplotlib.pyplot as plt

def train_and_predict():
    print('-'*30)
    print('Loading and preprocessing train data...')
    print('-'*30)
    imgs_train, imgs_mask_train = load_train_data()

    imgs_train = preprocess(imgs_train)
    imgs_mask_train = preprocess(imgs_mask_train)

    imgs_train = imgs_train.astype('float32')
    mean = np.mean(imgs_train)  # mean for data centering
    std = np.std(imgs_train)  # std for data normalization

    imgs_train -= mean
    imgs_train /= std
    #Normalization of the train set

    imgs_mask_train = imgs_mask_train.astype('float32')

    print('-'*30)
    print('Creating and compiling model...')
    print('-'*30)
    model = get_unet()
    model_checkpoint = ModelCheckpoint('weights.h5', monitor='val_loss', save_best_only=True)
    #Saving the weights and the loss of the best predictions we obtained

    print('-'*30)
    print('Fitting model...')
    print('-'*30)
    history=model.fit(imgs_train, imgs_mask_train, batch_size=10, epochs=20, verbose=1, shuffle=True,
              validation_split=0.2,
              callbacks=[model_checkpoint])

    print('-'*30)
    print('Loading and preprocessing test data...')
    print('-'*30)
    imgs_test, imgs_maskt = load_test_data()
    imgs_test = preprocess(imgs_test)

    imgs_test = imgs_test.astype('float32')
    imgs_test -= mean
    imgs_test /= std
    #Normalization of the test set

    print('-'*30)
    print('Loading saved weights...')
    print('-'*30)
    model.load_weights('weights.h5')

    print('-'*30)
    print('Predicting masks on test data...')
    print('-'*30)
    imgs_mask_test = model.predict(imgs_test, verbose=1)
    np.save('imgs_mask_test.npy', imgs_mask_test)
    print('-' * 30)
    print('Saving predicted masks to files...')
    print('-' * 30)
    pred_dir = 'preds'
    if not os.path.exists(pred_dir):
        os.mkdir(pred_dir)

    for k in range(len(imgs_mask_test)):
        a=rescale_intensity(imgs_test[k][:,:,0],out_range=(-1,1))
        b=(imgs_mask_test[k][:,:,0]).astype('uint8')
        io.imsave(os.path.join(pred_dir, str(k) + '_pred.png'),mark_boundaries(a,b))
    #Saving our predictions in the directory 'preds'
    plt.plot(history.history['dice_coef'])
    plt.plot(history.history['val_dice_coef'])
    plt.title('Model dice coeff')
    plt.ylabel('Dice coeff')
    plt.xlabel('Epoch')
    plt.legend(['Train', 'Test'], loc='upper left')
    plt.show()
    #plotting our dice coeff results in function of the number of epochs

if __name__ == '__main__':
    train_and_predict()

------------------------------
Loading and preprocessing train data...
------------------------------
------------------------------
Creating and compiling model...
------------------------------
------------------------------
Fitting model...
------------------------------
Train on 1083 samples, validate on 271 samples
Epoch 1/20
1083/1083 [==============================] - 18s 17ms/step - loss: -0.6274 - dice_coef: 0.6274 - val_loss: -0.3438 - val_dice_coef: 0.3438
Epoch 2/20
1083/1083 [==============================] - 17s 16ms/step - loss: -0.6386 - dice_coef: 0.6386 - val_loss: -0.5677 - val_dice_coef: 0.5677
Epoch 3/20
1083/1083 [==============================] - 17s 16ms/step - loss: -0.8072 - dice_coef: 0.8072 - val_loss: -0.6668 - val_dice_coef: 0.6668
Epoch 4/20
1083/1083 [==============================] - 17s 16ms/step - loss: -0.8585 - dice_coef: 0.8585 - val_loss: -0.6800 - val_dice_coef: 0.6800
Epoch 5/20
1083/1083 [==============================] - 17s 16ms/step - loss: -0.8862 - dice_coef: 0.8862 - val_loss: -0.6324 - val_dice_coef: 0.6324
Epoch 6/20
1083/1083 [==============================] - 17s 16ms/step - loss: -0.8991 - dice_coef: 0.8991 - val_loss: -0.7832 - val_dice_coef: 0.7832
Epoch 7/20
1083/1083 [==============================] - 17s 16ms/step - loss: -0.9232 - dice_coef: 0.9232 - val_loss: -0.7244 - val_dice_coef: 0.7244
Epoch 8/20
1083/1083 [==============================] - 17s 16ms/step - loss: -0.9318 - dice_coef: 0.9318 - val_loss: -0.7118 - val_dice_coef: 0.7118
Epoch 9/20
1083/1083 [==============================] - 17s 16ms/step - loss: -0.9367 - dice_coef: 0.9367 - val_loss: -0.7447 - val_dice_coef: 0.7447
Epoch 10/20
1083/1083 [==============================] - 18s 16ms/step - loss: -0.9430 - dice_coef: 0.9430 - val_loss: -0.8505 - val_dice_coef: 0.8505
Epoch 11/20
1083/1083 [==============================] - 17s 16ms/step - loss: -0.9551 - dice_coef: 0.9551 - val_loss: -0.8539 - val_dice_coef: 0.8539
Epoch 12/20
1083/1083 [==============================] - 17s 16ms/step - loss: -0.9598 - dice_coef: 0.9598 - val_loss: -0.8495 - val_dice_coef: 0.8495
Epoch 13/20
1083/1083 [==============================] - 18s 16ms/step - loss: -0.9658 - dice_coef: 0.9658 - val_loss: -0.8600 - val_dice_coef: 0.8600
Epoch 14/20
1083/1083 [==============================] - 17s 16ms/step - loss: -0.9447 - dice_coef: 0.9447 - val_loss: -0.8143 - val_dice_coef: 0.8143
Epoch 15/20
1083/1083 [==============================] - 18s 16ms/step - loss: -0.9482 - dice_coef: 0.9482 - val_loss: -0.8345 - val_dice_coef: 0.8345
Epoch 16/20
1083/1083 [==============================] - 18s 16ms/step - loss: -0.9670 - dice_coef: 0.9670 - val_loss: -0.8631 - val_dice_coef: 0.8631
Epoch 17/20
1083/1083 [==============================] - 17s 16ms/step - loss: -0.9716 - dice_coef: 0.9716 - val_loss: -0.8764 - val_dice_coef: 0.8764
Epoch 18/20
1083/1083 [==============================] - 17s 16ms/step - loss: -0.9712 - dice_coef: 0.9712 - val_loss: -0.8603 - val_dice_coef: 0.8603
Epoch 19/20
1083/1083 [==============================] - 18s 16ms/step - loss: -0.9756 - dice_coef: 0.9756 - val_loss: -0.8670 - val_dice_coef: 0.8670
Epoch 20/20
1083/1083 [==============================] - 18s 16ms/step - loss: -0.9769 - dice_coef: 0.9769 - val_loss: -0.8538 - val_dice_coef: 0.8538
------------------------------
Loading and preprocessing test data...
------------------------------
------------------------------
Loading saved weights...
------------------------------
------------------------------
Predicting masks on test data...
------------------------------
1057/1057 [==============================] - 5s 5ms/step
------------------------------
Saving predicted masks to files...
------------------------------

Thanks to the “mark_boundaries” function, we’ll be able to see the results of the segmentation with a yellow marker.
Finally the results ! Our work has paid off these curves represent train and test dice coefficients over 20 epochs. The test dice coefficient almost reached 0.87 which is quite satisfying.
You’ll see in the “preds” directory this kind of results that represent 2D cuts. What is outlined in yellow is the liver which has been predicted by the algorithm:
We can now segment thousands of scans in a fraction of seconds! A task that would take specialists much longer.
Finally we successfully trained our neural network using a U-net architecture with a Dice coefficient that reaches almost 0.87!
Improvement track:
reconstruct 3-d scans by stacking the 2-d segmented images.
multiply the data even more by making 2-d slices along several oblique axes
Hope that you found this article useful. Thanks for reading.