PyTorch DataLoaders

Let us try to use the transformation functions offered in MedicalTorch. PyTorch offers a native “transforms” module that helps us to stack up and apply many transformations to our data. In the code below, we first re-sample the dataset so that all samples are of the same size and then apply a crop filter, followed by a type transformation (to tensor format).

As you can see, we have loaded the data and printed out the size of a mini-batch. The mini-batch contains 4 images with dimension: 200 X 200 px. You can also visualize by batches (refer to the notebook below for source code).

Constructing The Segmentation Model

We explored the images above, now we want to build the gray matter segmentation model

And here are all the transformations that we apply to both the training and validation dataset:

We load the dataset into our environment and split it into training and validation portions. Again, note that you need to change the directory to point to your own data source.

Here are the final data loaders that we will use during training:

Models and Parameters

Below we declare our model and hyperparameters. Note that we are using GPU in this model. The model used below refers to the U-net convolutional-based architecture proposed by Ronneberger et al., 2015, which essentially aggregates semantic information to perform the image segmentation in the upper layers. See a figure of the U-net architecture below.

We build helper functions, such as accuracy, to produce the desired performance metrics for the model:

Training

Now we finally train the model for spinal cord gray matter segmentation. We report the training and testing accuracy below and train for 10 epochs only.

A lot happened in the code above! I would recommend that you spend some time digesting it and understanding every bit by going through the complete notebook provided at the end of this post. The final results produced by the U-net model is as follows:

Train loss: -0.9267, Training Accuracy: 97.0087

Val Loss: -0.9262, Validation Accuracy: 99.5775

The validation accuracy is 99% after 10 epochs! Not bad for what is considered a very complex and important task.

Final Words

In summary, you learned how to process MRI image scans using a neat and powerful tool known as medicaltorch. In addition, you learned how to pre-process, prepare and load the data using MedicalTorch’s and PyTorch’s built-in data loader functions. Finally, you trained a model based on convolutional neural networks to conduct spinal cord gray matter segmentation. Feel free to explore more of the utility functions provided in the medicaltorch API and explore different types of datasets.

Contributions

I also submitted a pull request where I tried to improve the documentation a bit and suggested minor changes to MedicalTorch’s project page. I hope it helps jump-start others into the cool and fun world of MRI data analysis. If you have any trouble running the code above or spot any bugs, please comment them below.

Things to Try

You can try the following tasks:

• Add suggestions to improve the attached notebook. And then we can propose it as a PR to the official documentation website of MedicalTorch.
• Apply the models to the testing dataset (I only did validation)
• Perform segmentation using the others models offered by the medicaltorch library
• Apply the models to a different type of dataset and contribute to the tutorial section of the medicaltorch documentation

References

Final Notebooks

Please find the Google Colab tutorial here.