Medical Image Segmentation Using HuggingFace & PyTorch

Medical image segmentation is an innovative process that enables surgeons to have a virtual “x-ray vision.” It is a highly valuable tool in healthcare, providing non-invasive diagnostics and in-depth analysis. With this in mind, in this post, we will explore the UW-Madison GI Tract Image Segmentation Kaggle challenge dataset. As part of this project, we will utilize PyTorch along with PyTorch-Lightning. We will use HuggingFace transformers to load and fine-tune the Segformer transformer-based model on the medical segmentation dataset. Finally, we will create a Gradio app for image inference and deploy it on HuggingFace spaces.

1. What is Medical Image Segmentation?
2. What Are The Problems Faced In Medical Image Segmentation?
3. Building The Medical Image Segmentation Dataset
4. Install & Import Required Libraries
5. Set Hyperparameters For The Project
6. Loading The Medical Image Segmentation Dataset
7. Using SegFormer Transformer Model From HuggingFace
8. Evaluation Metric & Loss Function
9. Creating Custom LightningModule Class
10. Start Training
11. Inference on the Medical Segmentation Dataset
12. Building The Gradio App
13. Summary

Designed for those familiar with the basics of multi-head attention, PyTorch, and image segmentation, this blog post aims to enrich your knowledge of medical image segmentation. We will guide you through every step, from dataset creation and training to deploying an inference model. We’ve crafted our explanations to be accessible, ensuring learners at all levels can benefit.

What is Medical Image Segmentation?

Medical image segmentation is a process that involves dividing medical images, such as CT scans or MRI scans, into distinct regions or structures of interest. This technique is used to identify and isolate specific areas within the image, which is crucial for diagnosis, treatment planning, and monitoring of diseases. It can be done manually by experts or automated using computer algorithms and machine learning. Medical image segmentation plays a vital role in various medical specialties and enables quantitative analysis and precise measurements.

The dataset for this project is taken from the UW-Madison GI Tract Image Segmentation Kaggle competition. The dataset consists of 3 classes: the stomach, small bowel, and large bowel.

About the competition:

In 2019, an estimated 5 million people were diagnosed with a cancer of the gastro-intestinal tract worldwide. Of these patients, about half are eligible for radiation therapy, usually delivered over 10-15 minutes a day for 1-6 weeks. Radiation oncologists try to deliver high doses of radiation using X-ray beams pointed to tumors while avoiding the stomach and intestines. With newer technology such as integrated magnetic resonance imaging and linear accelerator systems, also known as MR-Linacs, oncologists are able to visualize the daily position of the tumor and intestines, which can vary day to day. In these scans, radiation oncologists must manually outline the position of the stomach and intestines in order to adjust the direction of the x-ray beams to increase the dose delivery to the tumor and avoid the stomach and intestines. This is a time-consuming and labor intensive process that can prolong treatments from 15 minutes a day to an hour a day, which can be difficult for patients to tolerate—unless deep learning could help automate the segmentation process. A method to segment the stomach and intestines would make treatments much faster and would allow more patients to get more effective treatment.

– Motivation behind the UW-Madison GI Tract Image Segmentation Kaggle competition

In this competition, the task was to develop a model that can automatically segment the stomach and intestines in MRI scans. These scans come from real cancer patients who underwent 1-5 MRI scans during radiation treatment on different days. By utilizing a dataset of these scans, you will leverage deep learning techniques to create innovative solutions that improve the quality of care for cancer patients.

What Are The Problems Faced In Medical Image Segmentation?

Medical image segmentation faces several challenges, including:

1. Image variability: Diverse image characteristics make creating a universal segmentation algorithm challenging.
2. Ambiguity and complexity: Complex structures, ambiguous boundaries, and overlapping regions complicate accurate segmentation.
3. Limited labeled data: The scarcity of annotated medical images hampers developing and evaluating segmentation algorithms.
4. Computational complexity: Large datasets and high-dimensional images require substantial computational resources and time.
5. Validation and generalization: Ensuring accuracy and applicability across different modalities, populations, and clinical settings is difficult yet crucial.

Building The Medical Image Segmentation Dataset

As the task is to segment organ cells in images, the competition provides the dataset in the form of 16-bit grayscale PNG format images, and the ground-truth segmentation annotations are RLE-encoded masks provided in a CSV file. 

Some observations about the dataset.

1. It contains 115488 samples with multiple rows for each patient or case.
2. Out of 115488 rows, only 33913 (29.364%) have RLE annotations for the class. This number denotes the total number of annotations available.
3. The total number of images with corresponding annotations is 16590.
4. These 16590 images belong to 85 cases/patients, where each patient has multiple scans for each day on multiple days.
5. In these images:
   1. We have 8627 (~52%) annotations images for the Stomach.
   2. We have 11201 (~67.5%) annotations images for Small bowel.
   3. We have 14085 (~84.9%) annotations images for Large bowel.
6. There are 2468 (~6.41%) samples with one annotation present.
   1. Of these, 2286 (~92.6%) are Stomach.
   2. Of these, 123 (~4.98%) are Large bowel.
   3. Of these, 59 (~2.39%) are Small bowel.
7. There are 10921 (28.37%) samples with two annotations present.
   1. Of these, 7781 (~71.3%) are ‘Large bowel, Small bowel.’
   2. Of these, 2980 (~27.3%) are ‘Large bowel, Stomach.’
   3. Of these, 160 (~1.47%) are ‘Small bowel, Stomach.’
8. Finally, 3,201 (8.32%) examples with all three classes are present.

UW-Madison GI Tract Medical Segmentation Dataset Preprocessing

Below is the entire Python script to preprocess the raw dataset to get the final one we used for the project.

The algorithm used is as follows:

1. Traverse each “case” folder and record paths of all the images with annotations in the “train.csv” file using the get_folder_files(...) function.
2. Split the image paths list in an 80:20 ratio for the training and validation set.
3. Then we use the create_and_write_img_msk(...) function to iterate over files of each set to:
   • Load and convert the image from uint16 to uint8 format using the load_img(...) method.
   • Get the dataframe group (rows in “train.csv” file) for the image and decode the RLE-encoded mask for each class using rle_decode(...) method for each image. Then the RGB mask is converted to a grayscale mask (single channel) where each pixel value indicates the class ID of the pixel. This process is done using the rgb_to_onehot_to_gray(...) function.
   • Save image-mask pair in their respective folders.

Note:

1. There are additional intermediate steps, such as extracting the folder name, sub-folder name, and image path of the image.
2. The train.csv file contains image ids in the format "case<num>_day<num>_slice_<slice_id>". The get_folder_files(...) function, along with returning all the relevant image files in a “case” folder, returns a list of img_ids which has the same format as the “train.csv” id column. This is done to make it easier to fetch the relevant data frame rows during image and mask processing steps.

Script: all_in_one.py

Beginning from the next section, we’ll focus on the code implementation of the project. We’ll add all the explanations required to understand the procedures.

Install & Import Required Libraries

Before we begin the coding part, we must ensure all the required libraries are installed. For this project, apart from PyTorch, we are installing additional tools to help ease the implementation process. 

The major ones are:

1. transformers: To load the SegFormer transformer model.
2. lightning: To simplify and structure code implementations.
3. torchmetrics: For evaluating the model’s performance.
4. wandb: For experiment tracking. 
5. albumentations:  For applying augmentations. 

Let’s start

For this project, instead of the default tensorboard used by pytorch-lightning for tracking experiments, we will use a proper MLOps tool: Weights & Biases (WandB).

The following code cell will help us to log into our wandb account.

The code cell will ask you to paste your API key in the dialogue box. You need to click on the Sign In with Auth0 link provided.

Set Hyperparameters For The Project

Next, we will declare all the different hyperparameters used for the project. For this, we are defining three dataclasses. They will be used throughout the notebook.

1. DatasetConfig  – A class that holds all the hyperparameters we will use to process images. It contains the following information:
   1. Image size to use.
   2. Number of classes present in the dataset,
   3. The mean and standard deviation to use for image normalization.
   4. URL of the preprocessed dataset.
   5. Directory path to download the dataset to. 

2. Paths – This class contains the locations of the images and masks of the train and validation sets. It uses the “root dataset path”  set DatasetConfig as the base.

3. TrainingConfig –  A class that holds all the hyperparameters we will use for training and evaluation.  It contains the following information:
   1. Batch size.
   2. Initial learning rate.
   3. The number of epochs to train the model.
   4. The number of workers to use for data loading.
   5. Model, optimizer & learning rate scheduler-related configurations.
4. InferenceConfig – This class contains the (optional) batch size and the number of batches we will use to display our inference results at the end.

Note: We’ve uploaded the preprocessed dataset to our Dropbox and Kaggle accounts. There are two options. You can manually download the dataset and move it to your workstation or utilize the data download code we’ve written below to do it automatically.

Loading The Medical Image Segmentation Dataset

As mentioned before, we use the pytorch-lightning library to remove the boilerplate code necessary when writing pure PyTorch code. To do so, we must ensure we follow the format as Lightning requires. To that end, we will create our custom class MedicalSegmentationDataModule that inherits from LightningDataModule.

Before we begin, we will declare a dictionary to map each class ID to their respective RGB color values.

A Custom PyTorch Dataset Class For Medical Image Segmentation

First, we will define our custom PyTorch Dataset class. This custom is designed to load images and masks for each image. The Dataset class is essential for efficient and organized data handling in machine learning tasks. It provides a standardized interface to load and preprocess data samples from various sources. Encapsulating the dataset into a single object simplifies data management. It enables seamless integration with other PyTorch components like data loaders and models. 

The custom class performs the following functions:

1. Load each image-mask pair.
2. Apply geometric and pixel augmentations if the pair belongs to the training set.
3. Apply preprocessing transformations such as normalization and standardization.

Defining The Custom LightningDataModule Class 

In this section, we will define the custom MedicalSegmentationDataModule class inherited from Lightning’s LightningDataModule class. It helps organize and encapsulate all the data-related operations and logic in a PyTorch project. It acts as a bridge between your data and Lightning’s training pipeline. It is a convenient abstraction that encapsulates data-related operations, promotes code organization, and facilitates seamless integration with other Lightning components for efficient and reproducible deep-learning experiments.

The class will perform the following functions:

1. Download the dataset from Dropbox.
2. Create a MedicalDataset class object for each set.
3. Create and return the DataLoader objects for each set.

The class methods we need to define are as follows:

1. prepare_data(..): This method is used for data preparation, like downloading and one-time preprocessing with the dataset. When training in a distributed setting, this will be called from each GPU machine.
2. setup(...):  When you want to perform data operations on every GPU, this method is apt for it will call from every GPU. For example, perform train/val/test splits.
3. train_dataloader(...): This method returns the train dataloader.
4. val_dataloader(...): This method returns validation dataloader.

Usage: Let’s download the dataset and initialize train and validation data loaders. We’ll use them to visualize the dataset.

Some Helper Function To Visualize the Medical Segmentation Dataset

To help visualize our dataset, we need to define some additional helper functions. They are as follows:

A) num_to_rgb(...): Function will be used to convert single-channel mask representations to an integrated RGB mask for visualization purposes.

B) image_overlay(...): This function overlays an RGB segmentation map on top of an RGB image.

C) display_image_and_mask(...): The convenience function below will display the original image, the ground truth mask, and the ground truth mask overlayed on the original image.

D) denormalize(...): This function is used to denormalize the image tensors and clip values between 0 and 1. It is used to denormalize the images for visualization.

Display Sample Images From The Validation Dataset

We will loop over the first batch in the validation dataset and display the ground truth image, ground truth mask, and the ground truth mask overlayed on the image. The overlay helps us to better understand the segmented classes in the context of the original image.

Here are some samples:

Using SegFormer Transformer Model From HuggingFace

The SegFormer model was proposed in the paper titled SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers. The model consists of a hierarchical Transformer encoder made of efficient multi-head attention modules and a final lightweight all-MLP decoder head.

Paper abstract

“We present SegFormer, a simple, efficient yet powerful semantic segmentation framework which unifies Transformers with lightweight multilayer perception (MLP) decoders. SegFormer has two appealing features: 1) SegFormer comprises a novel hierarchically structured Transformer encoder which outputs multiscale features. It does not need positional encoding, thereby avoiding the interpolation of positional codes which leads to decreased performance when the testing resolution differs from training. 2) SegFormer avoids complex decoders. The proposed MLP decoder aggregates information from different layers, and thus combining both local attention and global attention to render powerful representations. We show that this simple and lightweight design is the key to efficient segmentation on Transformers. We scale our approach up to obtain a series of models from SegFormer-B0 to SegFormer-B5, reaching significantly better performance and efficiency than previous counterparts. For example, SegFormer-B4 achieves 50.3% mIoU on ADE20K with 64M parameters, being 5x smaller and 2.2% better than the previous best method. Our best model, SegFormer-B5, achieves 84.0% mIoU on Cityscapes validation set and shows excellent zero-shot robustness on Cityscapes-C.”

– SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers

You can check all the trained weights available for the SegFormer model on HuggingFace.

Loading a pre-trained model version and getting it ready for inference or finetuning is very easy, thanks to HuggingFace. 

We only have to pass the following:

1. pretrained_model_name_or_path: (string) The id/path of a pre-trained model hosted on the Huggingface model zoo.
2. num_labels: (int) The number of channels (one for each class) we want the model to give as output. Suppose the number differs from the original number. In that case, the layer will be replaced with a new layer with randomly initialized weights.    
3. ignore_mismatched_sizes: (bool) Boolean value to whether or not to ignore the weight key mismatch. Here, it occurs because we change the num_labels value.  

Function definition:

Usage

• The model’s forward pass takes multiple arguments [Segformer documentation]. The two important ones are pixel_values and labels.
• The pixel_values argument refers to the input images. The labels argument is for passing the ground-truth mask.
• If labels are passed, the model’s forward pass calculates the cross-entropy (CE) loss.
• The output logits are smaller than the input image size. To get the outputs to match the input image size, we need to upsample it simply.

Create dummy inputs.

Generate dummy outputs.

To access the model’s output, we have to use the "logits" key. Similarly, we can access the loss via the "loss" key.

In this project, we won’t be using the CE loss returned by the model for training. Instead, we will define our custom combo loss function that combines the Smooth Dice coefficient & CE to compute the loss.

Print model summary:

Evaluation Metric & Loss Function

The Dice Coefficient (otherwise known as the F1-Score) is a function commonly used in the context of segmentation and is often specifically used as the basis for a loss function for segmentation problems. We will write the custom loss function next based on the Dice Coefficient, but let’s first provide the motivation for why this might be a good idea. 

For a binary classification problem, the metric is defined as follows using set notation, where A and B are segmentation masks representing the ground truth mask and the predicted segmentation map, the Dice coefficient is defined as:

   

Simply put, the metric is twice the overlap area divided by the total number of pixels in both images. In terms of confusion matrix terms, it’s defined as:

   

Note that the Dice Coefficient can also be used as an evaluation metric and is used in the Kaggle competition as an evaluation metric along with 3D Hausdorff distance. But since, for this project, we are focusing on 2D images, we will stick with the Dice coefficient as our primary evaluation metric.

Custom Loss Functions – Smooth Dice + Cross-Entropy

Below, we define a custom loss function often used in segmentation problems when there is an imbalance in the classes within the dataset. The loss is based on the Dice metric and combined with Cross-entropy (CE) loss.

Dice + CE is a good loss function for semantic segmentation as it combines pixel-wise accuracy with boundary alignment, encouraging precise object localization. It addresses the class imbalance issue by incorporating the Dice coefficient, promoting balanced predictions and improving overall segmentation performance.

Evaluation Metric – Dice Coefficient 

To calculate the Dice score for the medical image segmentation task, we will use the MulticlassF1Score class from the torchmetrics library with the “macro” average reduction method.

Macro average refers to a method of calculating average performance in multiclass or multilabel classification problems, which treats all classes equally.

---

In practice, we found that using a combined loss (Dice loss + CCE loss) works better than Dice loss alone. This is also supported by our experiments:

---

Creating Custom LightningModule Class

The final custom class we need to create is the MedicalSegmentationModel which inherits its functionalities from Lightning’s LightningModule class.

The LightningModule class in pytorch-lightning is a higher-level abstraction that simplifies the training and organizing of PyTorch models. It provides a structured, standardized interface for defining and training deep learning models. It separates the concerns of model definition, optimization, and training loop, making the code more modular and readable.

The class methods we need to define are as follows:

1. Model initialization: __init__(...) method where the model and its parameters are defined. This method also includes the initialization of the loss and metric calculation methods.
2. Forward pass: forward(...) method where the forward pass of the model is defined.
3. Training step: training_step(...) method where the training step for each batch is defined. It includes calculating loss and metrics, which are logged for tracking.
4. Validation step: validation_step(...) method where the validation step for each batch is defined. It also includes the calculation of loss and metrics.
5. Optimizer configuration: configure_optimizers(...) method where the optimizer and, optionally, the learning rate scheduler are defined.

Moreover, two methods, on_train_epoch_end(...) and on_validation_epoch_end(...), are defined to log the average loss and f1 score after each epoch for training and validation, respectively.

Start Training

Once we have organized the LightningModule and LightningDataModule classes, we can utilize Lightning’s Trainer class to automate the remaining tasks effortlessly.

The Trainer offers a range of valuable deep-learning training functionalities, such as mixed-precision training, distributed training, deterministic training, profiling, gradient accumulation, batch overfitting, and more. Implementing these functionalities correctly can be time-consuming, but it becomes a swift process with the Trainer class.

By initializing our MedicalSegmentationModel and MedicalSegmentationDataModule classes and passing them to the .fit(...) method of the Trainer class instance, we can promptly commence training. This streamlined approach eliminates the need to implement various training aspects manually, providing convenience and efficiency.

Next, we will define a ModelCheckpoint and a LearningRateMonitor callback for saving the best model during training and the current learning rate of an epoch.

We will also initialize the WandbLogger to upload the training metrics to your wandb project. When the logger is initialized, we pass two parameters:

1. log_model=True – Upload the model as an artifact when the training is completed.
2. project – The project name to use on WandB. A project typically contains logs from multiple experiments along with their trained checkpoints.

When the logger is initialized, it will also print the link for the current experiment, which you open on any device to monitor the training process and also share with your team.

Train

Inference on the Medical Segmentation Dataset

For inference, we will use the same validation data as we did during training. We will plot the ground truth images, the ground truth masks, and the predicted segmentation maps overlayed on the ground truth images.

Load The Best Trained Model

Initialize the model with trained weights.

Evaluate Model On Validation Dataset

Get the best evaluation metrics using the saved model.

Validation metric	DataLoader 0
valid/f1	0.9437294006347656
valid/loss	0.31268423795700073

Log them as experiment summary metrics to WandB

Image Inference Using DataLoader Objects

In the code below, we define a helper function that performs inference given a trained model and a dataloader object. The model prediction will also be uploaded to wandb.

Usage:

Here are a few output samples.

Terminate the wandb experiment run.

Building The Gradio App

Medical Image Segmentation: Next, we’ll build a small UI app using Gradio for demonstration. The app can also be easily deployed on HuggingFace spaces.
If you observe, the size of the trained checkpoint file saved by lightning is around 997 MB. It contains the model’s and optimizer’s state_dict() and additional information such as callbacks, schedulers, etc.

To make things easier and reduce the file size to be deployed, we’ve extracted the model’s state_dict() from the pytorch-lightning checkpoint and processed it such that we can directly pass it to SegformerForSemanticSegmentation.from_pretrained(…) method.

Here are the steps that we followed:

A) Checking the contents of pytorch-lightning .ckpt file.

B) Check the “key” string format for the parameter loaded using the`SegformerForSemanticSegmentation.from_pretrained(...)` method.

Notice that each parameter key string has the prefix “segformer.”.

C) Check the same but for the pytorch-lightning checkpoint.

And here we see the difference between the parameter key strings. The extra “model.” prefix is the name of the variable used inside the MedicalSegmentationModel class to hold the model.

D) To remove the dependency on pytorch-lightning during inference, we will create a new state_dict() dictionary object where the keys will be the same as the ones in step B and the values will be from the pytorch-lightning checkpoint.

E) Next, we’ll use the .save_pretrained(...) method to save the transformer model.

So the next time we want to load the trained Segformer model, we can simply do:

The size of the new file is just 224 MB which is significantly smaller. This will also help to load the model quickly, as we are only loading the model’s parameters.

Script: app.py

The majority of the code remains the same. The only new function is the predict(…). It takes in a Pillow image and returns a list of tuples of length three (one for each class). The tuple object at each index contains 2 elements. One is the binary segmentation mask for the class, and the second is the class name. This output format is required by the gr.AnnotatedImage(...) component to display the image with predicted mask overlays.

---

You can use the Gradio inference demo here –> Medical Image Segmentation Gradio App

All the files required can be accessed from here –> Gradio App Files

---

You can also try out the Gradio app embedded below.

Summary

Medical image segmentation using deep learning provides significant advantages. Deep learning models excel at capturing complex patterns and features, leading to highly accurate and precise segmentation results compared to traditional methods. Additionally, deep learning algorithms automate segmentation, improving efficiency and enabling analysis of large volumes of medical image data. Moreover, deep learning models demonstrate adaptability and generalization, making them suitable for diverse image characteristics, imaging modalities, patient populations, and clinical settings, expanding their utility in medical imaging applications.

To summarise this article, we covered a comprehensive list of related topics:

1. Medical Image Segmentation: Explored the definition and challenges of medical image segmentation.
2. Dataset Preparation: Used the UW-Madison GI Tract segmentation dataset, made observations, and created preprocessed training and validation sets.
3. We defined a few essential functions and classes for PyTorch and PyTorch-Lightning frameworks to facilitate ease of training.
4. We learned how to use the Segformer model from Hugging Face transformers for segmentation and fine-tuned it on our dataset.
5. We defined a custom loss function combining the Dice coefficient with cross-entropy for improved segmentation performance.
6. Training and Metrics Tracking: Trained the model, monitored metrics using WandB, and uploaded the model as an artifact for future use.
7. We designed a user-friendly interface using the Gradio app, making our medical multi-label image classification model accessible to everyone.

References

1. UWM – GI Tract Image Segmentation – EDA
2. AW-Madison: EDA & In Depth Mask Exploration
3. UWMGI: Unet [Train] [PyTorch]