Brain tumor segmentation using MONAI and W&B Models
Supercharging your Deep Learning workflows for Medical Imaging using MonAI and Weights & Biases
Created on May 3|Last edited on March 19
Comment
Introduction
Machine learning is revolutionizing many fields of medicine, with success stories ranging from accurate disease diagnosis to early prediction of adverse events to automatic discovery of antibiotics to LLMs deployed on clinical trial reporting and a whole lot more.
Add to that list the automatic analysis of medical images. This is one of the most important applications of machine learning in the domain of medicine and healthcare. Applications can range from automatic detection of the type of melanoma from images of skin lesions to localized detection of pneumothorax from images of chest X-rays.
In this report, we will be discussing a similar type of application: the segmentation of brain tumors from MRI scans:
-
Semantic segmentation refers to the process of transforming raw medical images into clinically relevant, spatially structured information, such as outlining tumor boundaries.
-
Semantic segmentation is an essential prerequisite for a number of clinical applications, such as radiotherapy planning and treatment response monitoring.
-
Semantic segmentation is, right now, the most widely investigated medical image processing task, with ~70% of all biomedical image analysis challenges dedicated to it, according to the paper The Medical Segmentation Decathlon.
-
We will be performing semantic segmentation of MRI scans, which are 3D volumetric scans and is a little different from an X-Ray.
A 3D MRI Scan of the Human Brain
Semantic Segmentation of 3D MRI scans for Tumor Detection
1
Table of Contents →
Introduction🛠️ Tools of the TradeThe Medical Decathlon Dataset🗻 Training a Baseline Segmentation Model🚅 HyperParameter Optimization using Sweeps🏋️ Training the Best Model⛺️ Using MONAI and Weights & Biases via NVIDIA Base Command🙏 Acknowledgements and Further Reading
🛠️ Tools of the Trade
Before we jump into solving the problem of semantic segmentation of brain tumors, let us take a look at the tools we will be using for this task:
🩺 MONAI
Medical Open Network for Artificial Intelligence or MONAI in short, is a set of open-source, freely available collaborative frameworks built for accelerating research and clinical collaboration in medical imaging. The goal of MONAI is to accelerate the pace of innovation and clinical translation by building a robust software framework that benefits nearly every level of medical imaging, deep learning research, and deployment.
In this report, we will be specifically using MONAI Core, the flagship library of Project MONAI that provides domain-specific capabilities for training AI models for healthcare imaging. These capabilities include:
- Medical-specific image transforms
- State-of-the-art transformer-based 3D Segmentation algorithms like UNETR
- An AutoML framework named DiNTS
🐝 Weights & Biases
Weighs & Biases is the developer-first MLOps platform that makes it easy to track your experiments, manage, and version your data, plus collaborate with your team so you can focus on building the best models. In this report, we will be using the following key features of Weights & Biases:
- Data visualization and exploration.
- Experiment tracking and reproducibility.
- Model checkpoint management using Artifacts.
- Hyperparameter optimization using Sweeps.
- Sharing results using Reports.
The Medical Decathlon Dataset
In this report, we will be focusing on training a deep neural network on the dataset for the Brain Tumor Segmentation task from the Medical Segmentation Decathlon. The goal of this challenge is to produce segmentation labels of the different glioma sub-regions:
- Tumor core
- Whole tumor
- Enhancing tumor
👉 Clinical Relevance →
Glioma, particularly glioblastoma, and diffuse astrocytic glioma with molecular features of glioblastoma (WHO Grade 4 astrocytoma), are the most common and aggressive malignant primary tumors of the central nervous system in adults, with extreme intrinsic heterogeneity in appearance, shape, and histology.
Glioblastoma patients have a very poor prognosis, and the current standard of care comprises surgery, followed by radiotherapy and chemotherapy. MGMT (O[6]-methylguanine-DNA methyltransferase) is a DNA repair enzyme the methylation of its promoter in newly diagnosed glioblastoma has been identified as a favorable prognostic factor and a predictor of chemotherapy response. Thus the determination of MGMT promoter methylation status in newly diagnosed glioblastoma can influence treatment decision-making.
💿 Loading and Exploring the Data
Loading and Exploring the Brain Tumor Dataset
1
🗻 Training a Baseline Segmentation Model
For establishing our baseline experiment, we train a simple SegResNet model based on the paper 3D MRI brain tumor segmentation using autoencoder regularization. Let's explore the steps to write a simple training and validation loop in PyTorch:
Training and Validation Steps
1
🚅 HyperParameter Optimization using Sweeps
Now that we have set up our baseline training and validation experiment, let us try to find how we can get the best performance out of our current model using the most optimal set of hyperparameters. To this end, we use Weights & Biases Sweeps to automate our hyperparameter search and explore the space of possible models with just a few lines of code.
Hyperparameter Optimization using Sweeps
8
🏋️ Training the Best Model
Now that we have found out our optimal set of hyperparameters from the Sweep, let's use it to train the final and best iteration of our model.
Final Experiment with the best set of hyperparameters
2
⛺️ Using MONAI and Weights & Biases via NVIDIA Base Command
NVIDIA Base Command Platform is a comprehensive platform for businesses, data scientists, and IT teams that accelerates ROI for AI initiatives. It manages the end-to-end lifecycle of AI development including workload management and resource sharing with both a graphical user interface and command line APIs with integrated monitoring and reporting dashboards. Offered in a cloud-hosted solution that continuously delivers NVIDIA innovations directly into your AI workflow, Base Command Platform works across on-prem and cloud resources with a single pane of glass view into your AI development process. NVIDIA Base Command Platform is the hub of your AI lifecycle.
👉 Watch Demo →
You can develop your medical imaging training workflows directly on Nvidia Base Command Platform using MONAI and Weights & Biases and deploy such models into clinical production easily.
Get Started with a Free Hands-on Lab →
🙏 Acknowledgements and Further Reading
3D Segmentation with MONAI and PyTorch Supercharged by Weights & Biases
A tutorial on how to use Weights & Biases with MONAI and PyTorch to accelerate your medical research
Brain tumor 3D segmentation with MONAI
Brain tumor 3D segmentation with MONAI
How to Build a Robust Medical Model Using Weights & Biases
A tutorial on how to build a cancer detection model all while tracking of your experiments, storing model weights, optimizing hyperparameters, and more in W&B
Illuminating X-Rays: Deep Learning for Medical Imaging, Part I
Taking a closer look at detecting lung disease from chest Xrays
Add a comment
Tags: Articles, Health Care, Intermediate, Sweeps, Semantic Segmentation, Computer Vision, Experiment
Iterate on AI agents and models faster. Try Weights & Biases today.