AI in Healthcare: Using DL and TL to Locate Tumors in Brain Scans
Introduction
Artificial Intelligence is revolutionizing industries in many aspects especially in the case of healthcare. Common areas in how AI is branching into healthcare include:
Disease Diagnostics
Maximizing Efficiency in Hospitals
Imaging Automation
This project focuses on a combination of these aspects.
Pipeline
In this project I wanted to take an image of a brain scan, and run it through a classification model that first determines if a tumor exists in the brain scan, and if one exists, then run it through a segmentation model that pinpoints where the tumor is located. The classification model follows the Resnet architecture and the segmentation model follows the Resunet architecture.
Classification Model
The classification model is based on a Convolutional Neural Network. The first layer of this model is used to extract general insights and features at a high level. The last couple of layers is used to perform the actual classification task. What happens from start to finish, is the image is scanned to first search for simple shapes like edges or lines. Then the edges that are found are passed to the subsequent layer to generate/extract more complex features.
However, as CNNs grow deeper, an issue known as the vanishing gradient occurs which causes poor model performance. The answer to this is creating a Residual Neural Network, which includes a “skip connection” feature which allows for training of 152 layers without any vanishing gradient problems. It does this by adding an identity mapping on top of the CNN.
For this project I used transfer learning, in which a network has already been trained to perform a specific task but is now being repurposed to use towards tumor classification. Doing this greatly reduces the computational time required to train the model.
By using transfer learning, and starting off with better tuned weights for a similar task, the model performance greatly improved from the model with randomly initialized weights.
Image Segmentation
The purpose of image segmentation is to extract the information from the image at the pixel level and associate each pixel with a certain class. It’s provides great value in this context to locate the tumor by training a neural network to produce a pixel-wise mask of the tumor.
As I mentioned previously, the segmentation model I used for this project is based on a Unet architecture, which first encodes the image into a vector followed up by decoding it back into an image, all while formulating a loss function for every pixel in the original brain scan image. Starting with the encoding step, the first block consists of a 3x3 convolution layer + Relu activation function + Batch Normalization. The remaining three blocks consists of Res-blocks followed by a Max-pooling 2x2. Feature maps after each block doubles, which helps the model learn complex features effectively. After that it reaches the ‘bottleneck’ in which this serves as a connection between the encoding and decoding paths. This consists of a Res-block followed by a 2x2 up-sampling convolution layer. Then the decoder steps starts with the three blocks following the bottleneck consisting of Res-blocks followed by a 2x2 up-sampling convolutional layer. The final block is a Res-block which is followed by a 1x1 convolutional layer to get the same size as the orignal image. Reason for this is to help ensure the features learned while contracting are used while reconstructing an image. This is architecture is better shown in the image below.
After this process you are left with the predicted mask in the shape of the tumor.
Results
The image shows pic 1) original/input image 2) actual shape and location of tumor 3) predicted shape and location of tumor 4) actual tumor overlayed on input image 5) predicted tumor overlayed on original image
With the models scoped out, all that was left was to create a simple web app that takes a brain scan image and then runs through the entire pipeline and showcases the original brain image and then where the tumor is localized (if no tumor is seen in image, then only prints out ‘No Tumor’).
With the implementation of the pipeline we see very promising results. And with the more brain scans taken and used for training, the more accurate the model will be. This will only help hospitals and doctors as it will greatly reduce resources needed to diagnose tumors in the brain and leave more time available for other issues.
This was a very informative project for me, and I enjoyed every minute of it. If there are any other more efficient ways to implement this pipeline, I would be more than happy to know.
Please share all your thoughts with me!
If you would like to see my code for this project or for any of my other projects, they are all listed on my GitHub
