AI-Driven MRI Brain Segmentation Using U-Net Model

Abstract

In this white paper, we present an AI-based approach to brain MRI segmentation utilizing the U-Net architecture. Our solution enables precise identification of regions of interest in MRI scans, enhancing diagnostic capabilities in medical imaging. This paper outlines the design, architecture, and implementation of the model and showcases how it can be integrated into real-world applications using fast and efficient inference pipelines.

Demo: arithescientist.com/unetidiagnosis

Introduction

Medical imaging plays a crucial role in diagnosing diseases, particularly brain abnormalities. Manual segmentation of MRI scans is time-consuming and prone to errors. AI-driven segmentation models, such as U-Net, have proven effective in automating this process with high accuracy. This paper introduces a machine learning–based solution using U-Net for brain MRI segmentation.

• Problem statement: manual segmentation of MRI scans is labor-intensive and error-prone, necessitating an automated approach.
• Objective: develop a robust and efficient AI model to segment brain MRI scans accurately.

Methodology

Dataset preprocessing

The input MRI images are preprocessed to ensure uniformity and enhance segmentation accuracy. This involves resizing the images to a standard resolution of 256×256 pixels — the input size expected by the U-Net model.

# Resizing images to 256x256
def resize_image(input_image_path):
    img = Image.open(input_image_path)
    resized_img = img.resize(target_size)
    resized_img.save(input_image_path)
    return input_image_path

U-Net for brain segmentation

The U-Net architecture was chosen for its demonstrated effectiveness in biomedical image segmentation. Its symmetric encoder-decoder structure with skip connectionsretains fine spatial information from the encoder path while letting the decoder reconstruct a full-resolution mask — segmenting with high precision even when training data is limited.

• Encoder-decoder structure. Captures features from MRI images at different scales and restores image detail for precise segmentation.
• Skip connections. Preserve spatial information from encoder layers to decoder layers, improving segmentation accuracy.

The model is pre-trained on brain MRI datasets, and its weights are loaded through PyTorch Hub:

model = torch.hub.load(
    'arithescientisttt/unet_brainsegmentation',
    'unet',
    in_channels=3,
    out_channels=1,
    init_features=32,
    pretrained=True,
)

Image processing and inference

After preprocessing, we segment MRI images by sampling them into patches, passing each patch through the model, and aggregating the predictions.

• Grid sampler and aggregator. Images are split into smaller patches, processed individually, and combined into a final segmentation map.
• Inference process. Each image patch is processed by the U-Net model to predict the segmented regions.

# Sampling image patches
grid_sampler = tio.inference.GridSampler(subject_preprocessed, patch_size, patch_overlap)
aggregator = tio.inference.GridAggregator(grid_sampler)

Results and evaluation

The output of the U-Net model is a segmentation mask, which is visualized by overlaying it on the original MRI scan. This makes the segmented areas easy to interpret at a glance:

# Plotting MRI and segmentation mask
fig, ax = plt.subplots(figsize=(images.shape[1] / 100, images.shape[0] / 100), dpi=100)
ax.imshow(images)
ax.imshow(mask, cmap='gray', alpha=0.5)

The model was tested on several MRI images, with output confirming its ability to segment key areas in the brain effectively. Accuracy can be quantified using Dice similarity coefficients or Jaccard indices — both of which will be covered in future work.

Integration and application

Gradio interface

To make the segmentation model easy to interact with, a web-based interface was developed using Gradio. Users upload MRI images and the system returns the segmented output. Gradio was selected for its simplicity and its ability to give users immediate visual feedback.

interface = gr.Interface(
    inference,
    gr.Image(label="input image", type="filepath"),
    gr.Plot(),
    description=description,
    title=title,
    examples=examples,
)

FastAPI deployment

The entire solution can be deployed as a web application using FastAPI, allowing integration with other systems and APIs in real time. This makes it suitable for cloud deployment or as part of a hospital’s imaging workflow:

app = FastAPI()
app = gr.mount_gradio_app(app, interface, path="/")

Conclusion and future work

This project demonstrates the effectiveness of AI-based approaches — specifically U-Net — in automating brain MRI segmentation. The ability to accurately identify regions of interest in MRI scans has significant implications for medical diagnostics. Future work will focus on improving accuracy, adding training data, and extending the same U-Net + FastAPI + Gradio chassis to other medical-imaging tasks: tumor segmentation, organ delineation, and lesion tracking.

References

• Ronneberger, O., Fischer, P., & Brox, T. (2015). U-Net: Convolutional Networks for Biomedical Image Segmentation. arXiv:1505.04597
• PyTorch Hub Documentation. pytorch.org/hub
• TorchIO Documentation. torchio.readthedocs.io