Medical images segmentation with Keras: U-net architecture
Learn how to build a liver segmentation algorithm. A tutorial that includes data preprocessing, modelling and results visualisation.
Introduction
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.
Data
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.
Model
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:
Metric
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.
Project Structure
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.
Data Preprocessing
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!
Training the neural network with Keras
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:
Thanks to the “mark_boundaries” function, we’ll be able to see the results of the segmentation with a yellow marker.
Results
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.
Conclusion
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.
