Denoising MRI data with Tensorflow-Python
Tags: python, deeplearning, neuroimaginddata, ai
Published: 2022-02-18 17:17:31


You may find a lot of information on the internet about why deep learning is useful in medical imaging analysis and reconstruction. A simple answer is - ‘Medical data contains numerous amount of data points and where deep learning comes into play.’

> Data always follows a pattern, that can be analysed using different AI based models.

What next for this post??

This post is to show a basic example on how deeplearning can be used in medical domain supporting the healthcare professionals with denoising medical imaging data. It can also be used in data analysis, disease detection, artifacts removal, reconstructing undersampled data and lot more.

**<u>prerequisite</u>**: Basic knowledge on python, Numpy, Basics in dataprocessing(Image)

For our example let us take few medical imaging data samples then create and train a basic deep learning model (convolutional autoencoder) that reduces the noise in our medical images.

**Steps involved:**
- Reading data
- Preprocessing data
- Generating noisy data using scikit-learn random noise generator.
- Build a model and train
- Prediction and model evaluation


**Read data:** Here I am using a total of 20 MRI nifty volumes(T1 weighted MRI scans). We use nibabel a python module to load the data. Each volume contains a sequence of 128 brain slices which i have discussed in my previous blog @https://dev.to/narendraanupoju/mri-data-processing-with-python-1jgg
The following function can load the data from .nii or .nii.gz format files.

```
import nibabel as nib
def FileRead(file_path):
    data = nib.load(file_path).get_fdata()
    return data[:,:,14:114]
```
			
**Dataset structure** - Each sample had a shape (256, 256, 128) (represented as a 3 dimensional numpy array). Here we can observe that we have 128 slices of brain scans sequentially. We also need to perform normalization for each sample because of the wide range of data distribution

**Preprocess the data:** - Preprocessing is one of the most important step before training any machine learning or deep learning models. Preprocessing techniques can include data cleaning, data transformation, outliers detection, handling missing data ..soon a lot other depends on the input data and end result you need. We are using data transformation technique called normalization to reduce the variance in our data.

Now we stack all the volumes and generate a NumPy array before it can be used for training a deep learning model. I am only considering the centric slices of the volumes as the initial and the end volumes don't contain more useful information for the model training, so we are neglecting them. 100 center slices are considered for each sample.

<u>Note: Please note that, we are training the model with 2D slices rather than training with 3D volumes. </u>

```
import glob
import numpy as np
def normalize(x):
    x = (x-np.min(x)) / (np.max(x)-np.min(x))
    return x
listFiles = glob.glob('/*.nii.gz')
totalData = np.dstack([normalize(FileRead(file)) for file in totalFiles])
```
Finally, after the processing is done, you will get a NumPy array something like the below shape:

`Final shape of data : (256, 256, 2000)`

From the above snippet, we can observe 2000 samples of MRI slices aligned across third dimension axis. We are only working w.r.t 2D slices.

**Generating Noisy Data:** - Now, we have noise free images and we need noisy images to train our model w.r.t ground truths. For this, I am using sklearn random_noise generator to generate the noise on each image. The noise generated on each slice is based on the Gaussian noise considering the local variance for each data sample.

```
from skimage.util import random_noise

print('____Generating noisy data______')
noisyData= np.zeros(totalData.shape)
for img in range(samplesCount):
    noisy= random_noise(totalData[:,:,img],mode='localvar' ,seed=None,clip=True)
    noisyData[:,:,img] = noisy
```

![generated_noise_01](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/agrpp5y9jpb64s1oxyhy.PNG)![generated_noise_02](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/pvb2293rrzf9lpl4x6pc.PNG)
After generating noisy data, now we have noisy samples w.r.t ground truths. We can observe from the above noisy images that most of the brains internal structures are distorted. Now we use a simple convolution autoencoder model built on top of TensorFlow to remove the noise from our data reconstructs the image with better internal structures.

`Noisy samples data shape : (256, 256, 2000) #containing noise`
`Ground truth data shape : (256, 256, 200)`

One last step before model training is we have to structure our data with standard representation `(batch_size, lenght, width, height)`. For that, we expand **totalData** and **noisyData** with 4th dimension to train and roll the third axis to starting index as shown below

```
totalData = np.rollaxis(totalData, -1)
totalData = np.expand_dims(totalData, axis=-1)
noisyData = np.rollaxis(noisyData, -1)
noisyData= np.expand_dims(noisyData, axis=-1)

X_train, X_test, y_train, y_test = train_test_split(noisyData, totalData , test_size=0.2, random_state=42)
```
<u>Note</u> You can also consider converting numpy arrays to tensors and use the advantage of tensorflow Dataloader class.
For simple understandinf i am only considering only numpy arrays.

After spliting the data shapes are as follows:
X_train : (1600, 256, 256, 1)
y_train : (1600, 256, 256, 1)
X_test  : (400 , 256, 256, 1)
X_test  : (400 , 256, 256, 1)


**Model building and training**
Now we have simple convolutional autoencode model as shown below to train and predict our results:

```
import tensorflow as tf
import keras
import keras.layers as layers
import keras.models as models
from keras.initializers import orthogonal
from tensorflow.keras.optimizers import Adam

def ConvolutionLayer(x, filters, kernel, strides, padding, block_id, kernel_init=orthogonal()):
    prefix = f'block_{block_id}_'
    
    x = layers.Conv2D(filters, kernel_size=kernel, strides=strides, padding=padding,
                      kernel_initializer=kernel_init, name=prefix+'conv')(x)
    
    x = layers.LeakyReLU(name=prefix+'lrelu')(x)
    x = layers.Dropout(0.2, name=prefix+'drop')((x))
    x = layers.BatchNormalization(name=prefix+'conv_bn')(x)
    return x

def DeconvolutionLayer(x, filters, kernel, strides, padding, block_id, kernel_init=orthogonal()):
    prefix = f'block_{block_id}_'
    
    x = layers.Conv2DTranspose(filters, kernel_size=kernel, strides=strides, padding=padding,
                               kernel_initializer=kernel_init, name=prefix+'de-conv')(x)
    
    x = layers.LeakyReLU(name=prefix+'lrelu')(x)
    x = layers.Dropout(0.2, name=prefix+'drop')((x))
    x = layers.BatchNormalization(name=prefix+'conv_bn')(x)
    return x


# Convolutional Autoencoder with skip connections
def noiseModel(input_shape):
    inputs = layers.Input(shape=input_shape)
    
    # 256 x 256
    conv1 = ConvolutionLayer(inputs, 64, 3, strides=2, padding='same', block_id=1)

    conv2 = ConvolutionLayer(conv1, 128, 5, strides=2, padding='same', block_id=2)
    
    # 64 x 64
    conv3 = ConvolutionLayer(conv2, 256, 3, strides=2, padding='same', block_id=3)
    
    # 16 x 16
    deconv1 = DeconvolutionLayer(conv3, 128, 3, strides=2, padding='same', block_id=4)
    
    # 64 x 64
    deconv2 = DeconvolutionLayer(deconv1, 64, 3, strides=2, padding='same', block_id=5)
    
    # 256 x 256
    deconv3 = DeconvolutionLayer(deconv2, 1, 3, strides=2, padding='same', block_id=6)
    
    return models.Model(inputs=inputs, outputs=deconv3)

```
From the above code each **ConvolutionLayer** and **DeconvolutionLayer** block contains sequential operation of convolution/convolution transpose, followed by leakyRelu(activation function), dropout and batchnormalization layers.
Our convolutional autoencoder model consists of six blocks, first three blocks is the **_encoder network _**part where they perform the convolution operations to generate a latent feature map of the noisy data, then the latent feature map is passed to **_decoder _**for generating a noise free image.

Further more detailed information on the model development and parametric optimization can be discussed in the next post.

Now we have our model blocks **ConvolutionLayer** and **DeconvolutionLayer** which will support the model generation **<u>noiseModel</u>**

Next, we have to define an optmizer, metric and the loss functions that is required for our model training for updating the weights and bias.

```
input_shape = (256, 256)    # shape of the sample
model = noiseModel((*input_shape, 1))
model_optimizer = Adam(learning_rate=0.002)

def SSIMMeasure(y_true, y_pred):
    return tf.reduce_mean(tf.image.ssim(y_true, y_pred, 1.0))

def SSIMLoss(y_true, y_pred):
    return 1 - tf.reduce_mean(tf.image.ssim(y_true, y_pred, 1.0))

model.compile(optimizer=model_optimizer, loss=SSIMLoss, metrics=[SSIMMeasure])
```
Now we are good to call the training method to train our model.
Let us view the summary of our simple convolutional autoencoder  model.

```
print(model.summary())
Model: "model"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 input_1 (InputLayer)        [(None, 256, 256, 1)]     0         
                                                                 
 block_1_conv (Conv2D)       (None, 128, 128, 64)      640       
                                                                 
 block_1_lrelu (LeakyReLU)   (None, 128, 128, 64)      0         
                                                                 
 block_1_drop (Dropout)      (None, 128, 128, 64)      0         
                                                                 
 block_1_conv_bn (BatchNorma  (None, 128, 128, 64)     256       
 lization)                                                       
                                                                 
 block_2_conv (Conv2D)       (None, 64, 64, 128)       204928    
                                                                 
 block_2_lrelu (LeakyReLU)   (None, 64, 64, 128)       0         
                                                                 
 block_2_drop (Dropout)      (None, 64, 64, 128)       0         
                                                                 
 block_2_conv_bn (BatchNorma  (None, 64, 64, 128)      512       
 lization)                                                       
                                                                 
 block_3_conv (Conv2D)       (None, 32, 32, 256)       295168    
                                                                 
 block_3_lrelu (LeakyReLU)   (None, 32, 32, 256)       0         
                                                                 
 block_3_drop (Dropout)      (None, 32, 32, 256)       0         
                                                                 
 block_3_conv_bn (BatchNorma  (None, 32, 32, 256)      1024      
 lization)                                                       
                                                                 
 block_4_de-conv (Conv2DTran  (None, 64, 64, 128)      295040    
 spose)                                                          
                                                                 
 block_4_lrelu (LeakyReLU)   (None, 64, 64, 128)       0         
                                                                 
 block_4_drop (Dropout)      (None, 64, 64, 128)       0         
                                                                 
 block_4_conv_bn (BatchNorma  (None, 64, 64, 128)      512       
 lization)                                                       
                                                                 
 block_5_de-conv (Conv2DTran  (None, 128, 128, 64)     73792     
 spose)                                                          
                                                                 
 block_5_lrelu (LeakyReLU)   (None, 128, 128, 64)      0         
                                                                 
 block_5_drop (Dropout)      (None, 128, 128, 64)      0         
                                                                 
 block_5_conv_bn (BatchNorma  (None, 128, 128, 64)     256       
 lization)                                                       
                                                                 
 block_6_de-conv (Conv2DTran  (None, 256, 256, 1)      577       
 spose)                                                          
                                                                 
 block_6_lrelu (LeakyReLU)   (None, 256, 256, 1)       0         
                                                                 
 block_6_drop (Dropout)      (None, 256, 256, 1)       0         
                                                                 
 block_6_conv_bn (BatchNorma  (None, 256, 256, 1)      4         
 lization)                                                       
                                                                 
=================================================================
Total params: 872,709
Trainable params: 871,427
Non-trainable params: 1,282
_________________________________________________________________
```
From the above model summary you can observe that we have around 0.9 million trainable parameters and very few non trainable parameters. Adding more layers will result in increase of trainable parameters. 

For training we have to define the batch size and epochs aswell, and also we log the information into tensorboard and csv logger so that the data can be analysed lateron. For this i have the following training loop as below:

```
Path_MODELSAVE = 'saved_models'
Path_LOGS = 'logs'
for i in [Path_MODELSAVE, Path_LOGS]:
    if i not in os.listdir(os.getcwd()):
        os.mkdir(i)

epochs = 100
batch_size = 10
saved_weight = os.path.join(Path_MODELSAVE, 'dataweights.{epoch:02d}-{SSIMMeasure:.2f}.hdf5')

model_checkpoint = keras.callbacks.ModelCheckpoint(saved_weight,
                                        monitor = 'val_SSIMMeasure',
                                        verbose=1,
                                        save_best_only=False,
                                        save_weights_only=True,
                                        mode='auto', save_freq = 'epoch', period = 10)

tensorboard = keras.callbacks.TensorBoard(log_dir=Path_LOGS,
                                            histogram_freq=0,
                                            write_graph=True,
                                            write_images=True)

csv_logger = keras.callbacks.CSVLogger(f'{Path_LOGS}/keras_log.csv' ,append=True)
                                    
model.fit(X_train, y_train, epochs=epochs, batch_size=batch_size, 
            validation_data=(X_test, y_test), 
            callbacks=[model_checkpoint, tensorboard, csv_logger])

```
From the above sniplet **model_checkpoint** saves the model in our local system over a period of 10 epochs.
**tensorboard** will log the tensorboard events to Path_LOGS directory and **csv_logger** will log **loss**, **SSIMMeasure** for both training and validation into a .csv file for each epoc.

model.fit() function will perform all the training process considering the batch_size and callbacks as defined above.

**Prediction and model evaluation**
Once the data is trained, you can perform prediction on any other samples for model evaluation by maintaining the batch size of the trained model.
Also please note that you have to perform the same preprocessing steps on the samples that are passed for model prediction. In our case we have to perform normalization before passing a sample for prediction.

Here are few predictions from the model trained with 130 epochs.
![model_prediction_1](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/akf30hi1rongzvoakrdr.PNG)

![model_prediction_2](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/1d27s4oqqgmimwbyjhpq.PNG)

From above results we can observe that our model could able to remove the noise and also generates few structures that are not exactly visible in our noisy data. It all happened by training upto few 130 epochs with batch_size equal to 10 and model containing only 6 blocks.

If you want to see better results, i suggest to train your model with more number of epochs with a greater batch_size by inspecting whether your  model is overfitting or underfitting.

Below you can observe the model accuracy and loss over training and validation.

Training results:

![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/gudcro0v02tfnnkavker.PNG)
Validation results:

![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/5s2stlk1iyyx34okqobj.PNG)

If you have any questions regarding the post, please comment. I will try to answer asap.