Medical Image Segmentation

1. Problem Statement

Medical Image Segmentation is a computer vision task that involves dividing a medical image into multiple segments. In this context, the task is the segmentation of healthy organs in medical scans, particularly in the gastrointestinal (GI) tract, to enhance cancer treatment. For patients eligible for radiation therapy, oncologists aim to deliver high doses of radiation using X-ray beams targeted at tumors while avoiding the stomach and intestines. The goal is to effectively segment the stomach and intestines in MRI scans to improve cancer treatment, eliminating the need for the time-consuming and labor-intensive process in which radiation oncologists must manually outline the position of the stomach and intestines, addressing the challenge of daily variations in the tumor and intestines' positions.

As shown in the figure, the tumor (pink thick line) is close to the stomach (red thick line). High doses of radiation are directed at the tumor while avoiding the stomach. The dose levels are represented by a range of outlines, with higher doses shown in red and lower doses in green.

The main challenge is to provide better assistance to patients. The issue lies in the tumor size, which often results in radiation X-rays inadvertently coming into contact with healthy organs. The segmentations must be as precise as possible to prevent any unintended harm to the patient. The problem is to develop a deep learning solution that automates the segmentation of the stomach and intestines in MRI scans of cancer patients who undergo 1-5 MRI scans on separate days during their radiation treatment.

2. Related Works

Date	Title	Code	Link
2015	U-Net: Convolutional Networks for Biomedical Image Segmentation	Code Code	Link
2015	SegNet: A Deep Convolutional Encoder-Decoder Architecture for Image Segmentation		Link
2016	V-Net: Fully Convolutional Neural Networks for Volumetric Medical Image Segmentation	Code	Link
2018	UNet++: A Nested U-Net Architecture for Medical Image Segmentation	Code	Link
2019	CE-Net: Context Encoder Network for 2D Medical Image Segmentation	Code	Link
2022	Medical Image Segmentation using LeViT-UNet++: A Case Study on GI Tract Data	None	Link
2023	3D TransUNet: Advancing Medical Image Segmentation through Vision Transformers	Code	Link
2023	DA-TransUNet: Integrating Spatial and Channel Dual Attention with Transformer U-Net for Medical Image Segmentation	Code	Link
2023	GI Tract Image Segmentation with U-Net and Mask R-CNN	None	Link

3. The Proposed Method

U-Net is a popular and effective architecture for medical image segmentation tasks, including segmenting different parts of the gastrointestinal tract. It is known for its ability to produce accurate segmentations, especially when dealing with limited training data. The architecture of U-Net gives the model the ability of precise localization, meaning the model output a class label for each pixel, and therefore achieve image segmentation.

Figure shows network artitecture in the original paper. It consists of the repeated application of two 3x3 convolutions (unpadded convolutions), each followed by a rectified linear unit (RELU) and a 2x2 max pooling operation with stride 2 for downsampling. At each downsampling step we double the number of feature channels. Every step in the expansive path consists of an upsampling of the feature map followed by a 2x2 convolution (“up-convolution”) that halves the number of feature channels, a concatenation with the correspondingly cropped feature map from the contracting path, and two 3x3 convolutions, each followed by a ReLU. The cropping is necessary due to the loss of border pixels in every convolution. At the final layer a 1x1 convolution is used to map each 64-component feature vector to the desired number of classes. In total the network has 23 convolutional layers.

For this project, As shown in the figure, the model takes MRI scans from cancer patients as input images, then uses the U-Net method to obtain predicted segmented areas of patients' MRI scans for "stomach", "large bowel", and "small bowel". By employing the loss function, it compares the predicted mask to the true mask, which we aim to minimize.

The evaluation metric is used the Dice.

4. Implementation

This section delves into the practical aspects of the project's implementation.

4.1. Dataset

Under this subsection, you'll find information about the dataset used for the medical image segmentation task. It includes details about the dataset source, size, composition, preprocessing, and loading applied to it. Dataset

The dataset is MRIs of patients provided by the UW-Madison Carbone Cancer Center. Specifically, the dataset contains 85 cases with 38496 scan slices of organs represented in 16-bit grayscale PNG format. Each case is represented by multiple sets of scan slices (each set is identified by the day the scan took place). Each scan slices is repeted 3 times, with large_bowel, small_bowel, and stomach classes. The annotations are provided in a csv format with the segmented areas represented as RLE-encoded masks. It would typically need to decode the RLE encoded masks to create pixel-wise binary masks. An empty segmentation entry represents no mask presented for the class in the MRI scan slice. The dataset has missing values.

The most common size across all images in the dataset is 266 × 266, and the rest are of sizes 310×360, 276×276, and 234×234 in frequency descending order. For this project, all images reshape to size 256 × 256. In order to feed the image, and mask into training models, process them as tensors.

Files

• train.csv - IDs and masks for all training objects.
• train.txt - case IDs for training objects.
• validation.txt - case IDs for validation objects.
• test.txt - case IDs for test objects.
• train - a folder of case/day folders, each containing slice images for a particular case on a given day.

Columns of train CSV file

• id - unique identifier for object
• class - the predicted class for the object
• segmentation - RLE-encoded pixels for the identified object

Split the dataset to training, validation, and testing sets based on provided text files. (data folder)

• train_data.csv: IDs, class, segmentation, and image_paths for training objects.
• validation_data.csv: IDs, class, segmentation, and image_paths for validation objects.
• test_data.csv: IDs, class, segmentation, and image_paths for test objects.

Visualization of a batch of images and target masks in the training set

4.2. Model

In this project, the SegmentationModelsPytorch library is used along with Pytorch to ceate a UNet model.

create segmentation model with pretrained encoder
in_channels = 3, # model input channels
classes = 3, # model output channels

model = smp.Unet(encoder_name='efficientnet-b1',
                                in_channels=3,
                                encoder_weights='imagenet',
                                classes=3,
                                activation=None)

4.3. Train

• Finding Hyper-parameters
• Step 1: Calculate the loss for an untrained model using a few batches.
• Step 2: Train the model for a limited number of epochs, experimenting with various learning rates.

• Main Loop
• Define model and optimizer and Set learning rate and weight decay.
• train the model for epoches.

4.4. Evaluate

• Plot learning curves

• model's segmentation result