Jun 29, 2026
Lightweight Brain Tumor Segmentation on Low-Resource Systems: A Step-by-Step Guide with 3D U-Net
- Ayomide Oladele1,
- Raymond Confidence1,2,3,
- Dong Zhang1,4,
- Charity Umoren1,
- Aondona M Iorumbur5,
- Anu Gbadamosi1,
- Farouk Dako1,6,
- Maruf Adewole1,6,
- Udunna Anazodo1,2,3
- 1Medical Artificial Intelligence Laboratory, Lagos, Nigeria;
- 2Montreal Neurological Institute, McGill University, Montreal, Canada;
- 3Department of Biomedical Engineering, McGill University, Montreal, Canada;
- 4Department of Electrical and Computer Engineering, University of British Columbia, Vancouver, Canada;
- 5Department of Medical Physics, University of Lagos, Nigeria;
- 6Department of Radiology, University of Pennsylvania, USA
- Maruf Adewole: contributed equally;
- Udunna Anazodo: contributed equally;
Protocol Citation: Ayomide Oladele, Raymond Confidence, Dong Zhang, Charity Umoren, Aondona M Iorumbur, Anu Gbadamosi, Farouk Dako, Maruf Adewole, Udunna Anazodo 2026. Lightweight Brain Tumor Segmentation on Low-Resource Systems: A Step-by-Step Guide with 3D U-Net. protocols.io https://dx.doi.org/10.17504/protocols.io.14egn4bd6v5d/v1
License: This is an open access protocol distributed under the terms of the Creative Commons Attribution License, which permits unrestricted use, distribution, and reproduction in any medium, provided the original author and source are credited
Protocol status: Working
We use this protocol and it's working
Created: March 24, 2025
Last Modified: June 29, 2026
Protocol Integer ID: 124927
Keywords: conventional deep learning models for brain tumor segmentation, lightweight brain tumor segmentation model, lightweight brain tumor segmentation, lightweight brain tumor segmentation on low, automated brain tumor segmentation, deep learning in low, brain tumor segmentation, brain tumor segmentation in sub, conventional deep learning model, deep learning model development, deep learning method, enabling deep learning, skills training in deep learning method, promising segmentation performance, application of deep learning method, precise identification of tumor region, high computing resource, need for high computing resource, segmentation performance, expensive computational resources such as gpus, computational efficiency with segmentation accuracy, dataset, tumor region, segmentation accuracy, substantial computational resource, expensive computational resource, computing resource, mri, gpus, segmentation output
Disclaimer
About this Tutorial
This tutorial was collaboratively developed by research staff at the Medical Artificial Intelligence (MAI) Lab, Lagos, Nigeria, and participants of the Sprint AI Training for African Medical Imaging Knowledge Translation (SPARK). It provides a hands-on, end-to-end guide for building and training a lightweight deep learning model for automated brain tumor segmentation using the BraTS-Africa 2024 dataset.
Designed with accessibility in mind, this tutorial is optimized for devices with limited computing power and demonstrates how meaningful results can be obtained using only a CPU. By reducing dependence on expensive GPU-based infrastructure, the tutorial makes deep learning more approachable for a broader community of learners, researchers, and healthcare technology practitioners, especially in resource-constrained settings.
Details of the lightweight architecture and its pilot implementation during the 2025 SPARK Academy are summarized in the MIRASOL paper: Lightweight 3D U-Net for Brain Tumor Segmentation on CPUs: Enabling Deep Learning in Low-Resource Environments.
Disclaimer: This tutorial is provided for informational and educational purposes only. Use at your own risk.
DISCLAIMER – FOR INFORMATIONAL PURPOSES ONLY; USE AT YOUR OWN RISK
The content provided in this tutorial is for informational purposes only and does not constitute legal, clinical, medical, or safety advice. It has not undergone peer review or formal approval processes. Always exercise your own professional judgment before acting on any information presented. Neither the authors nor any affiliated organization can be held liable for the use or misuse of this material.
Abstract
The application of deep learning methods in healthcare has gained significant popularity and relevance among researchers, clinicians, consumers, academic institutions, and industry stakeholders, leading to impactful discoveries that are improving human health. One important application is automated brain tumor segmentation, which supports the precise identification of tumor regions on magnetic resonance imaging (MRI) scans for diagnosis, treatment planning, and prognosis.
However, conventional deep learning models for brain tumor segmentation often require substantial computational resources, limiting their use by researchers, students, and practitioners in resource-constrained settings. This computational burden also restricts skills training in deep learning methods, particularly in environments where access to high-performance GPUs is limited or unavailable.
This tutorial presents a practical approach to addressing the need for high computing resources in deep learning model development. It provides a step-by-step guide to developing a lightweight brain tumor segmentation model using a 3D U-Net architecture optimized for low-resource systems. Using the Brain Tumor Segmentation in Sub-Saharan African Populations (BraTS-Africa) 2024 dataset, the model architecture was trained efficiently and evaluated on standard CPUs without relying on GPUs. The approach presented in this tutorial seeks to balance computational efficiency with segmentation accuracy. The lightweight model achieved a Dice score of 0.70 on the validation data, and the segmentation output was visually compared with the ground truth. Despite being trained using limited computing resources, the model demonstrated promising segmentation performance.
Details of the lightweight architecture and its pilot implementation during the 2025 SPARK Academy are summarized in the MIRASOL paper: Lightweight 3D U-Net for Brain Tumor Segmentation on CPUs: Enabling Deep Learning in Low-Resource Environments.
The main objective of this tutorial is to empower researchers, students, and practitioners in resource-constrained settings to learn how to develop, validate, and deploy deep learning methods using existing frameworks without relying on expensive computational resources such as GPUs. More importantly, the tutorial aims to enable a wider audience to gain practical AI skills, facilitate the development of locally relevant tools for early detection, and ultimately contribute to improved patient outcomes.
Introduction
This comprehensive tutorial guides you through the end-to-end process of developing and training a lightweight deep learning model for brain tumor segmentation using the BRATS 2024 dataset. Specifically designed for computers with limited computing resources, this tutorial demonstrates how to achieve promising results on a CPU, making deep learning accessible to a wider range of learners and users.
Brain tumor segmentation is a critical task in medical imaging, where accurate identification and delineation of tumor regions in brain magnetic resonance imaging (MRI) scans are crucial for diagnosis, treatment planning, and monitoring disease progression. However, this process is challenging due to the complex variability of tumor shapes, sizes, and locations, as well as the intricate structure of the brain. In resource-constrained settings, where access to skilled radiologists and medical facilities is limited, delayed diagnosis can lead to poor outcomes including mortalities.
To address this issue, AI-powered tools can facilitate early detection and initial assessment. Nevertheless, most operational AI models require high-performance computing resources (e.g., GPUs or TPUs), which can hinder their deployment on low-resource systems.
To overcome this limitation, lightweight models that balance accuracy and efficiency are essential. These models must be capable of operating effectively on standard CPUs, without relying on high-end GPUs or extensive memory. By developing and deploying such models, we can expand access to timely andaccurate diagnosis, ultimately improving health outcomes in resource-constrained communities.
The tutorial is divided into four phases:
1: Data Collection, Preparation and Preprocessing Phase
2: Data Loading and Model Building Phase
3: Model Training and Evaluation Phase
4: Model Deployment and Practical Stages
| Specification | Minimum System Requirements | System Specification used for this tutorial | |
| Processor | Core i5 and more | Core i5 | |
| Python Version | 3.8+ | 3.11 | |
| RAM | 4GB and more | 8GB | |
| Storage | 5GB of free disk space (can either be internal or external hard drive) | 12GB of free disk space | |
| The model runs on CPU-only systems | |||
‱ Installation Instructions: Required Dependencies
To run this project, please ensure you have Python 3.8 or above installed. Then follow the steps below to set up the required environment.
# Create virtual environment
python -m venv venv
# Activate it
# On Windows:
venv\Scripts\activate
# On macOS/Linux:
source venv/bin/activate
Install dependencies
Use pip to install all required libraries:
pip install numpy pandas tensorflow keras segmentation_models_3D scikit-learn matplotlib split-folders
PHASE 1: DATA COLLECTION, PREPARATION AND PREPROCESSING
The Integrated Development Environment (IDE )used for this tutorial is Visual Studio Code (VS Code)).
An IDE is a software application that helps programmers develop software codes, like a container that allows programmers to type in their python codes. Google colab, Kaggle notebooks, VS Code are some of the popular python IDEs.
A. Dataset Overview
The BraTS-Africa dataset used in this tutorial was collected from The CancerImaging Archive (TCIA), a publicly available repository of medical images (https://www.cancerimagingarchive.net/collection/brats-africa/). The dataset contains 146 MR images (scans) of both glioma (95) and other neoplasms (51) tumor in a NIFTI format (`.nii.gz`), which is a common format for storing 3D medical imaging data. Each scan includes the multiple modalities listed below:
T1n -> (T1-weighted, native) – t1n.nii.gz
T1c ->(T1-weighted, contrast-enhanced - t1c.nii.gz
T2f -> (T2-weighted, fluid sensitive) – t2f.nii.gz
T2w -> (T2 weighted) – t2w.nii.gz
and corresponding segmentation masks of 3 tumor sub-regions generated by expert
readers, as reference standard. - seg.nii.gz
Illustration of the BraTS-Africa Dataset showing brain slices of the four MRI image contrasts and egmentation masks of the three tumor sub-regions in one representative patient (Adewole et al., 2025).
B. Download the Dataset
Before you click on the “DOWNLOAD(1.6GB)” button, you need to download IBM Aspera plugin by clicking on the IBM-Aspera-Connect plugin. This redirects you to the page below:
On this path, select the Download now. link with respect to your PC operating system. For, Windows, Download now for Windows. Then install the IBM Aspera connect application. Navigate back to the TCIA page and click on the DOWNLOAD (1.6GB) button,which will redirect you the page below.
Then click on the Download icon, immediately the download will start and be saved on your local computer as a 1.6 GB zipped folder. Right-click on the folder to extract (or decompress) the folder. The folder downloaded will later be opened in your VS code application.
C. Data Collection and Pre-Processing
The first step is to download the dataset from TCIA. The dataset is organized into patient-specific folders, each containing multiple `. nii.gz` files for different MRI modalities and segmentation masks (See above). The data was later split into three folders of train, val and test data in a “glioma split data”; Data split: Train = 66 cases Validation = 13 cases , test = 9 cases.
Loading and Converting NIfTI Files
NIfTI files are loaded using the `nibabel` library, which provides tools for reading and writing neuroimaging data. The data is then converted into NumPy arrays for easier processing and manipulation.
import numpy as np
import nibabel as nib
#
Define the path to your dataset
TRAIN_DATASET_PATH = '/95_Glioma'
#
Load sample images and visualize
image_t1n = nib.load(TRAIN_DATASET_PATH + '/BraTS-SSA-00008-000/BraTS-SSA-00008-000-t1n.nii.gz').get_fdata()
image_t1c = nib.load(TRAIN_DATASET_PATH + '/BraTS-SSA-00008-000/BraTS-SSA-00008-000-t1c.nii.gz').get_fdata()
image_t2f = nib.load(TRAIN_DATASET_PATH + '/BraTS-SSA-00008-000/BraTS-SSA-00008-000-t2f.nii.gz').get_fdata()
image_t2w = nib.load(TRAIN_DATASET_PATH + '/BraTS-SSA-00008-000/BraTS-SSA-00008-000-t2w.nii.gz').get_fdata()
mask = nib.load(TRAIN_DATASET_PATH + '/BraTS-SSA-00008-000/BraTS-SSA-00008-000-seg.nii.gz').get_fdata()
Scaling the Images
Medical images often have varying intensity ranges. To standardize the data, we use `MinMaxScaler` from `sklearn` to scale the pixel values to a range of [0,1]. Note that the reshape was used to convert the 3D images to a 2D shape so the scaler function can operate on the images. The .reshape (-1, image_t1n.shape[-1]) flattens the first two dimensions while keeping the last one unchanged, converting (H, W, D) into a 2D shape (H*W, D).
from sklearn.preprocessing import MinMaxScaler
scaler = MinMaxScaler()
image_t1n = scaler.fit_transform(image_t1n.reshape(-1,image_t1n.shape[-1])).reshape(image_t1n.shape)
image_t1c = scaler.fit_transform(image_t1c.reshape(-1,image_t1c.shape[-1])).reshape(image_t1c.shape)
image_t2f=scaler.fit_transform(image_t2f.reshape(-1,image_t2f.shape[-1])).reshape(image_t2f.shape)
image_t2w = scaler.fit_transform(image_t2w.reshape(-1,image_t2w.shape[-1])).reshape(image_t2w.shape)
Handling Segmentation Masks
The segmentation masks are labelled with integer values representing different tumour regions. For consistency, we convert the mask to `uint8` and reassign label `4`to `3`.
mask = mask.astype(np.uint8)
mask[mask== 4] = 3 # Reassign mask values 4 to 3
Note that:
0 → Background (Non-brain regions or healthy tissue)
1 → Tumor Core (Necrotic/Cancerous region)
2 → Edema (Swelling or fluid accumulation around the tumor)
3 → Enhancing Tumor (Active tumor region that enhances with contrast agents in MRI)
Visualizing the Data
To ensure the data is correctly loaded and processed, we visualize a random slice
from the MRI scan and its corresponding mask.
import numpy as np
import matplotlib.pyplot as plt
import random
n_slice = random.randint(0, mask.shape[2])
plt.figure(figsize=(12,8))
plt.subplot(231)
plt.imshow(np.rot90(image_t1n[:,:, n_slice]), cmap='gray')
plt.title('Imaget1n')
plt.subplot(232)
plt.imshow(np.rot90(image_t1c[:,:, n_slice]), cmap='gray')
plt.title('Image t1c')
plt.subplot(233)
plt.imshow(np.rot90(image_t2f[:,:, n_slice]), cmap='gray')
plt.title('Image t2f')
plt.subplot(234)
plt.imshow(np.rot90(image_t2w[:,:, n_slice]), cmap='gray')
plt.title('Imaget2w')
plt.subplot(235)
plt.imshow(np.rot90(mask[:,:, n_slice]))
plt.title('Mask')
plt.show()
Visual Output:
Combining Multi-Modality Images
Brain MRI scans often include multiple modalities (e.g., T1, T2, FLAIR), each providing complementary information. We combine these modalities into a single multi-channel NumPy array for input to the model.
combined_x = np.stack([test_image_t1n, test_image_t1c, test_image_t2f, test_image_t2w], axis=3)
Cropping the Images
To reduce computational load, we crop the images to a smaller size (from 256x256x256 to 128x128x128). This step ensures the data fits into memory and is compatible with the model architecture.
combined_x = combined_x[56:184, 56:184, 13:141] #Crop to 128x128x128
test_mask = test_mask[56:184, 56:184, 13:141]
Processing the Entire Dataset
The above steps are repeated for all images in the dataset. We use `glob` to locate all `.nii.gz` files and process them in a loop. A 1% threshold was applied to screen out images with less tumor-related images (1, 2, or 3), if too much background label (0) is found, it is removed from the images. The images are finally saved into a folder “glioma”.
# Process all images in the dataset
t1n_list = sorted(glob.glob(TRAIN_DATASET_PATH + '/*/*t1n.nii.gz'))
t1c_list = sorted(glob.glob(TRAIN_DATASET_PATH + '/*/*t1c.nii.gz'))
t2f_list = sorted(glob.glob(TRAIN_DATASET_PATH + '/*/*t2f.nii.gz'))
t2w_list = sorted(glob.glob(TRAIN_DATASET_PATH + '/*/*t2w.nii.gz'))
mask_list = sorted(glob.glob(TRAIN_DATASET_PATH + '/*/*seg.nii.gz'))
for img in range(len(t1n_list)): # Using t1n_list as all lists are of the same size
print("Now preparing image and masks number: ",
img)
temp_image_t1n = nib.load(t1n_list[img]).get_fdata()
temp_image_t1n = scaler.fit_transform(temp_image_t1n.reshape(-1,
temp_image_t1n.shape[-1])).reshape(temp_image_t1n.shape)
temp_image_t1c = nib.load(t1c_list[img]).get_fdata()
temp_image_t1c = scaler.fit_transform(temp_image_t1c.reshape(-1,
temp_image_t1c.shape[-1])).reshape(temp_image_t1c.shape)
temp_image_t2f = nib.load(t2f_list[img]).get_fdata()
temp_image_t2f =scaler.fit_transform(temp_image_t2f.reshape(-1,
temp_image_t2f.shape[-1])).reshape(temp_image_t2f.shape)
temp_image_t2w = nib.load(t2w_list[img]).get_fdata()
temp_image_t2w = scaler.fit_transform(temp_image_t2w.reshape(-1,
temp_image_t2w.shape[-1])).reshape(temp_image_t2w.shape)
temp_mask = nib.load(mask_list[img]).get_fdata()
temp_mask = temp_mask.astype(np.uint8)
temp_mask[temp_mask == 4] = 3 # Reassign mask values4 to 3
temp_combined_images = np.stack([temp_image_t1n,
temp_image_t1c, temp_image_t2f, temp_image_t2w], axis=3)
# Crop to a size divisible by 64 (or any desired size)
temp_combined_images = temp_combined_images[56:184,
56:184, 13:141]
temp_mask = temp_mask[56:184, 56:184, 13:141]
val, counts = np.unique(temp_mask, return_counts=True)
if (1 - (counts[0] / counts.sum())) > 0.01: #
At least 1% useful volume with labels that are not 0
print("Save Me")
temp_mask = to_categorical(temp_mask, num_classes=4)
np.save('glioma/images/image_' + str(img) + '.npy', temp_combined_images)
np.save('glioma/masks/mask_' + str(img) + '.npy', temp_mask)
else:
print("I am not a good addition to the model")
Data Splitting
The dataset is split into training, validation, and test sets using the `splitfolders` library. This ensures a balanced distribution of data across the sets. The split datasets were saved in a folder “glioma split data”.
import splitfolders
input_folder = 'glioma/'
output_folder = 'glioma split data/'splitfolders.ratio(input_folder, output=output_folder, seed=42, ratio=(.75, .15, .10), group_prefix=None)
This concludes Part 1: Data Preparation and Preprocessing. The dataset is now ready for training a lightweight brain tumour segmentation model. In the next part, we will build and train a 3D U-Net model using the pre-processed data. Here is the link to the complete code for Part 1.
PART 2: BUILDING THE 3D U-NET MODEL
Objective: Design a lightweight 3D U-Net model suitable for low-resource systems.
A. Introduction to 3D U-Net
The 3D U-Net architecture is a powerful model for medical image segmentation tasks, particularly for volumetric data like brain MRI scans. It extends the traditional 2D U-Net by incorporating a third spatial dimension, enabling the model to capture 3D contextual information. This is especially important for
tasks like brain tumour segmentation, where tumours can span multiple slices in a 3D volume.
However, 3D U-Net models can be computationally expensive, making them challenging to run on low-resource systems (e.g., normal CPUs). To address this, we make several modifications to create a lightweight 3D U-Net:
- Fewer Layers: Reducing the number of convolutional layers to decrease
model complexity.
- Smaller Filters: Using fewer filters in each convolutional layer to
reduce memory usage.
- Efficient Patch Extraction: Processing smaller 3D patches instead of the
entire volume to fit within memory constraints.
B. Model Implementation
The code for the implementation of the 3D U-Net model was referenced from Sreenivas [Bhattiprolu, n.d.] who converted his 2D U-Net to a simple 3D U-Net model (https://github.com/bnsreenu/python_for_microscopists/blob/master/231_234_BraTa2020_Unet_segmentation/simple_3d_unet.py). The 3D U-Net uses 3D convolutions to process volumetric data, includes dropout layers to prevent overfitting, uses a contracting path (encoder) to extract features and an expanding path (decoder) to reconstruct segmentation maps, and outputs a multi-class segmentation mask.
kernel_initializer = 'he_uniform' #Try others if you want
################################################################
def simple_unet_model(IMG_HEIGHT, IMG_WIDTH, IMG_DEPTH, IMG_CHANNELS,
num_classes):
#Build the model
inputs = Input((IMG_HEIGHT, IMG_WIDTH, IMG_DEPTH,
IMG_CHANNELS))
#s = Lambda(lambda x: x / 255)(inputs) #No need
for this if we normalize our inputs beforehand
s = inputs
#Contraction path
c1 = Conv3D(16, (3, 3, 3), activation='relu',kernel_initializer=kernel_initializer, padding='same')(s)
c1 = Dropout(0.1)(c1)
c1 = Conv3D(16, (3, 3, 3), activation='relu',
kernel_initializer=kernel_initializer, padding='same')(c1)
p1 = MaxPooling3D((2, 2, 2))(c1)
c2 = Conv3D(32, (3, 3, 3), activation='relu',
kernel_initializer=kernel_initializer, padding='same')(p1)
c2 = Dropout(0.1)(c2)
c2 = Conv3D(32, (3, 3, 3), activation='relu',
kernel_initializer=kernel_initializer, padding='same')(c2)
p2 = MaxPooling3D((2, 2, 2))(c2)
c3 = Conv3D(64, (3, 3, 3), activation='relu',
kernel_initializer=kernel_initializer, padding='same')(p2)
c3 = Dropout(0.2)(c3)
c3 = Conv3D(64, (3, 3, 3), activation='relu',
kernel_initializer=kernel_initializer, padding='same')(c3)
p3 = MaxPooling3D((2, 2, 2))(c3)
c4 = Conv3D(128, (3, 3, 3), activation='relu',
kernel_initializer=kernel_initializer, padding='same')(p3)
c4 = Dropout(0.2)(c4)
c4 = Conv3D(128, (3, 3, 3), activation='relu',
kernel_initializer=kernel_initializer, padding='same')(c4)
p4 = MaxPooling3D(pool_size=(2, 2, 2))(c4)
c5 = Conv3D(256, (3, 3, 3), activation='relu',
kernel_initializer=kernel_initializer, padding='same')(p4)
c5 = Dropout(0.3)(c5)
c5 = Conv3D(256, (3, 3, 3), activation='relu',
kernel_initializer=kernel_initializer, padding='same')(c5)
#Expansive path
u6 = Conv3DTranspose(128, (2, 2, 2), strides=(2, 2, 2),
padding='same')(c5)
u6 = concatenate([u6, c4])
c6 = Conv3D(128, (3, 3, 3), activation='relu',
kernel_initializer=kernel_initializer, padding='same')(u6)
c6 = Dropout(0.2)(c6)
c6 = Conv3D(128, (3, 3, 3), activation='relu',
kernel_initializer=kernel_initializer, padding='same')(c6)
u7 = Conv3DTranspose(64, (2, 2, 2), strides=(2, 2, 2),
padding='same')(c6)
u7 = concatenate([u7, c3])
c7 = Conv3D(64, (3, 3, 3), activation='relu',
kernel_initializer=kernel_initializer, padding='same')(u7)
c7 = Dropout(0.2)(c7)
c7 = Conv3D(64, (3, 3, 3), activation='relu',
kernel_initializer=kernel_initializer, padding='same')(c7)
u8 = Conv3DTranspose(32, (2, 2, 2), strides=(2, 2, 2),
padding='same')(c7)
u8 = concatenate([u8, c2])
c8 = Conv3D(32, (3, 3, 3), activation='relu',
kernel_initializer=kernel_initializer, padding='same')(u8)
c8 = Dropout(0.1)(c8)
c8 = Conv3D(32, (3, 3, 3), activation='relu',
kernel_initializer=kernel_initializer, padding='same')(c8)
u9 = Conv3DTranspose(16, (2, 2, 2), strides=(2, 2, 2),
padding='same')(c8)
u9 = concatenate([u9, c1])
c9 = Conv3D(16, (3, 3, 3), activation='relu',
kernel_initializer=kernel_initializer, padding='same')(u9)
c9 = Dropout(0.1)(c9)
c9 = Conv3D(16, (3, 3, 3), activation='relu',
kernel_initializer=kernel_initializer, padding='same')(c9)
outputs = Conv3D(num_classes, (1, 1, 1),
activation='softmax')(c9)
model = Model(inputs=[inputs], outputs=[outputs])
#compile model outside of this function to make it flexible.
model.summary()
return model
#Test if everything is working ok.
model = simple_unet_model(128, 128, 128, 4, 4)
print(model.input_shape)
print(model.output_shape)
C. Patch Extraction and Data Augmentation
Patch Extraction
Processing entire 3D volumes can be memory-intensive, especially on low-resource systems. To address this, we extract smaller 3D patches from the input volumes. This reduces memory usage while still providing sufficient context for segmentation.
def extract_patch(image, mask, patch_size):
img_shape = image.shape[:3]
patch_x = np.random.randint(0, max(img_shape[0] - patch_size[0], 1))
patch_y = np.random.randint(0, max(img_shape[1] - patch_size[1], 1))
patch_z = np.random.randint(0, max(img_shape[2] - patch_size[2], 1))
return (
image[patch_x:patch_x + patch_size[0], patch_y:patch_y + patch_size[1], patch_z:patch_z + patch_size[2], :],
mask[patch_x:patch_x + patch_size[0], patch_y:patch_y + patch_size[1], patch_z:patch_z + patch_size[2]] )
The `extract_patch` function randomly selects a patch of a specified size from the input volume and mask.
Data Augmentation Techniques
Data augmentation improves model robustness by introducing variability into the training data. The `augment_image` function applies random rotations, flips, brightness adjustments, noise addition, and gamma correction.
def gamma_correction(image, gamma):
return np.clip(image gamma, 0, 1)
def augment_image(image, mask):
# Rotation
angle = np.random.uniform(-15, 15)
image = rotate(image, angle, axes=(0, 1), reshape=False, mode='reflect')
mask = rotate(mask, angle, axes=(0, 1), reshape=False, mode='reflect')
# Flipping
if np.random.rand() > 0.5:
image, mask = np.flip(image, axis=0),
np.flip(mask, axis=0)
if np.random.rand() > 0.5:
image, mask = np.flip(image, axis=1),
np.flip(mask, axis=1)
# Brightness Adjustment
brightness = np.random.uniform(0.9, 1.1)
image = np.clip(image * brightness, 0, 1)
# Noise Addition (Gaussian noise)
if np.random.rand() > 0.5:
noise = np.random.normal(0, 0.02,
image.shape)
image = np.clip(image + noise, 0, 1)
# Gamma Correction
if np.random.rand() > 0.5:
gamma = np.random.uniform(0.8, 1.2)
image = gamma_correction(image, gamma)
return image, mask
Image Loader Function
The `imageLoader` function generates batches of augmented patches for training. The code below combines the augment and extract patch function into a function that loads the images for training and evaluation “imageLoader”.
import numpy as np
import os
from scipy.ndimage import rotate
def load_img(img_dir, img_list):
images = []
for image_name in img_list:
if image_name.endswith('.npy'):
try:
image =
np.load(os.path.join(img_dir, image_name),
allow_pickle=True).astype(np.float32)
images.append(image)
except Exception as e:
print(f"Error loading file {image_name}: {e}")
return np.array(images) if images else np.array([])
def extract_patch(image, mask, patch_size):
img_shape = image.shape[:3]
patch_x = np.random.randint(0, max(img_shape[0] -
patch_size[0], 1))
patch_y = np.random.randint(0, max(img_shape[1] -
patch_size[1], 1))
patch_z = np.random.randint(0, max(img_shape[2] -
patch_size[2], 1))
return (
image[patch_x:patch_x + patch_size[0],
patch_y:patch_y + patch_size[1], patch_z:patch_z + patch_size[2], :],
mask[patch_x:patch_x + patch_size[0],
patch_y:patch_y + patch_size[1], patch_z:patch_z + patch_size[2]]
)
def gamma_correction(image, gamma):
return np.clip(image gamma,
0, 1)
def augment_image(image, mask):
# Rotation
angle = np.random.uniform(-15, 15)
image = rotate(image, angle, axes=(0, 1), reshape=False,
mode='reflect')
mask = rotate(mask, angle, axes=(0, 1), reshape=False,
mode='reflect')
# Flipping
if np.random.rand() > 0.5:
image, mask = np.flip(image, axis=0),
np.flip(mask, axis=0)
if np.random.rand() > 0.5:
image, mask = np.flip(image, axis=1),
np.flip(mask, axis=1)
# Brightness Adjustment
brightness = np.random.uniform(0.9, 1.1)
image = np.clip(image * brightness, 0, 1)
# Noise Addition (Gaussian noise)
if np.random.rand() > 0.5:
noise = np.random.normal(0, 0.02,
image.shape)
image = np.clip(image + noise, 0, 1)
# Gamma Correction
if np.random.rand() > 0.5:
gamma = np.random.uniform(0.8, 1.2)
image = gamma_correction(image, gamma)
return image, mask
def imageLoader(img_dir, img_list, mask_dir, mask_list, batch_size,
patch_size):
L = len(img_list)
while True:
for batch_start in range(0, L, batch_size):
limit = min(batch_start +
batch_size, L)
X = load_img(img_dir,
img_list[batch_start:limit])
Y = load_img(mask_dir,
mask_list[batch_start:limit])
if X is None or Y is None:
continue
X_patches, Y_patches =
zip(*[augment_image(*extract_patch(img, mask, patch_size)) for img, mask in
zip(X, Y)])
yield np.stack(X_patches,
axis=0), np.stack(Y_patches, axis=0)
This concludes Part 2: Building the 3D U-Net Model. The model is now ready for training on the preprocessed dataset. In the next part, we will cover the training process and evaluation metrics.
PART 3: TRAINING AND EVALUATION
Objective:
Train the model and evaluate its performance.
A. Training Setup
Defining the Dice Loss Function
The Dice Loss function was custom defined as a class. The class implements the Dice Score metric, commonly used in segmentation tasks to evaluate the overlap between predicted and ground truth masks. The Dice Score is particularly useful in medical image segmentation, such as brain image segmentation, where precise boundary delineation is crucial.
class DiceScore(tf.keras.metrics.Metric):
def __init__(self, num_classes, class_weights=None,
smooth=1e-6, kwargs):
super(DiceScore, self).__init__( kwargs)
self.num_classes = num_classes
self.smooth = smooth
self.class_weights = class_weights if
class_weights is not None else tf.ones(num_classes) # Default to equal
weights
self.dice_scores =
self.add_weight(name='dice_scores', shape=(self.num_classes,),
initializer='zeros')
def update_state(self, y_true, y_pred,
sample_weight=None):
# Flatten the tensors to ensure
computations are class-wise
y_true = tf.reshape(y_true, [-1]) #
Flatten ground truth
y_pred = tf.reshape(y_pred, [-1]) #
Flatten predictions
# Initialize the dice scores for each
class
dice_scores = []
for i in range(self.num_classes):
# Create binary masks for
class i
y_true_class =
tf.cast(tf.equal(y_true, i), 'float32')
y_pred_class =
tf.cast(tf.equal(tf.round(y_pred), i), 'float32')
# Calculate intersection
and union
intersection =
tf.reduce_sum(y_true_class * y_pred_class)
union =
tf.reduce_sum(y_true_class) + tf.reduce_sum(y_pred_class)
# Compute Dice score for
the current class
dice_class = (2. *
intersection + self.smooth) / (union + self.smooth)
# Apply class weight to the
Dice score
weighted_dice_class =
dice_class * self.class_weights[i]
dice_scores.append(weighted_dice_class)
# Update state by averaging weighted dice
scores across all classes
dice_scores = tf.stack(dice_scores)
self.dice_scores.assign(dice_scores)
def result(self):
# Return the mean of the weighted dice
scores
return tf.reduce_mean(self.dice_scores)
def reset_states(self):
# Reset the dice scores at the start of
each batch
self.dice_scores.assign(tf.zeros(self.num_classes))
Model Training
The model is trained using a loss function that combines Dice Loss and Categorical Focal Loss, both of which are well-suited for segmentation tasks. Dice Loss emphasizes the overlap between predictions and ground truth, improving segmentation accuracy, while Focal Loss addresses class imbalance by down-weighting easy examples and focusing on hard-to-classify pixels.
Model evaluation is conducted using accuracy, Intersection over Union (IoU), and a custom Dice Score metric. To enhance convergence, a learning rate scheduler is implemented to reduce the learning rate when validation loss plateaus. The model is optimized using the Nadam optimizer with a
fixed learning rate of 0.001, and gradient clipping is applied to prevent exploding gradients.
For efficient training, data generators are used to load and preprocesstraining and validation data in batches. Additionally, callbacks are incorporated to save the best model and implement early stopping, terminating training if validation loss ceases to improve.
import time
start =time.time()
# Set up loss function with class weights
wt0, wt1,wt2, wt3 = 0.25, 0.25, 0.25, 0.25
class_weights= np.array([wt0, wt1, wt2, wt3], dtype=np.float32)
dice_loss =sm.losses.DiceLoss(class_weights=class_weights)
focal_loss= sm.losses.CategoricalFocalLoss()
total_loss = 0.2 * dice_loss + 0.2 * focal_loss # focal loss of 0.2 (perf best
now) becomes 0.1
# Define metrics
metrics =['accuracy', sm.metrics.IOUScore(threshold=0.5), DiceScore(num_classes=4)]
# Learning rate scheduler (Cosine Decay)
lr_scheduler = tf.keras.callbacks.ReduceLROnPlateau(
monitor='val_loss', factor=0.5, patience=5, min_lr=1e-5 # from 6 to
4(perf better) to 3)
# Optimizer (use a fixed learning rate)
optimizer = keras.optimizers.Nadam(learning_rate=0.001, clipnorm=1.0) #increase from 0.001(perf best) to 0.01
# Data paths
DATA_ROOT = "glioma split data"
train_img_dir, train_mask_dir = f"{DATA_ROOT}/train/images/", f"{DATA_ROOT}/train/masks/"
val_img_dir, val_mask_dir = f"{DATA_ROOT}/val/images/",
f"{DATA_ROOT}/val/masks/"
# Load data
train_img_list, train_mask_list = os.listdir(train_img_dir), os.listdir(train_mask_dir)
val_img_list, val_mask_list = os.listdir(val_img_dir), os.listdir(val_mask_dir)
# Data
generators
batch_size = 2
patch_size = (64, 64, 64)
train_data = imageLoader(train_img_dir, train_img_list, train_mask_dir, train_mask_list,
batch_size, patch_size)
val_data = imageLoader(val_img_dir, val_img_list, val_mask_dir, val_mask_list, batch_size,
patch_size)
# Training parameters
steps_per_epoch = len(train_img_list) // batch_size
val_steps_per_epoch = len(val_img_list) // batch_size
epochs = 100
# Initialize and compile model
model = simple_unet_model(IMG_HEIGHT=64, IMG_WIDTH=64, IMG_DEPTH=64, IMG_CHANNELS=4,
num_classes=4)
model.compile(optimizer=optimizer,
loss=total_loss, metrics=metrics)
model.summary()
# Define callbacks
checkpoint_path = f"saved_model/3D_unet_{epochs}_epochs_{batch_size}_batch_patch_training.keras"
callbacks =[
EarlyStopping(monitor='val_loss', patience=10,
restore_best_weights=True),
ModelCheckpoint(checkpoint_path, monitor='val_loss',
save_best_only=True, mode='min'),
lr_scheduler # Learning rate decay callback
]
# Train model
history = model.fit(
train_data,
steps_per_epoch=steps_per_epoch,
epochs=epochs,
validation_data=val_data,
validation_steps=val_steps_per_epoch,
callbacks=callbacks,
verbose=1)
# Save training history
history_df = pd.DataFrame(history.history)
os.makedirs("model_history",exist_ok=True)
history_df.to_csv("model_history/training_history.csv", index=False)
print("Training history and best model saved successfully.")
end =time.time()
exec_time = (end - start)/60
print(f'execution time is - {exec_time}')
The model training ran on my system for an average of 1 hr and 5 minutes.
B. Model Evaluation
Loading the Trained Model
The trained model is loaded for inference without recompiling.
import os
import numpy as np
import keras from keras.models import load_model
from keras.metrics import MeanIoU
from matplotlib import pyplot as plt
# --------------------MODEL LOADING --------------------
# Path to trained model
model_path = 'saved_model/3D_unet_100_epochs_2_batch_patch_training.keras'
# Load the trained model without recompiling (for inference only)
my_model = load_model(model_path, compile=False)
#-------------------- DATA LOADING --------------------
# Path to dataset (Modify as needed)
DATA_PATH = "glioma split data"
# Train and Validation Directories
train_img_dir = os.path.join(DATA_PATH, "train/images/")
train_mask_dir = os.path.join(DATA_PATH, "train/masks/")
val_img_dir = os.path.join(DATA_PATH, "val/images/")
val_mask_dir = os.path.join(DATA_PATH, "val/masks/")
test_img_dir = os.path.join(DATA_PATH, "test/images/")
test_mask_dir = os.path.join(DATA_PATH, "test/masks/")
# Get list of images and masks
train_img_list, train_mask_list = os.listdir(train_img_dir), os.listdir(train_mask_dir)
val_img_list, val_mask_list = os.listdir(val_img_dir), os.listdir(val_mask_dir)
test_img_list, test_mask_list = os.listdir(test_img_dir), os.listdir(test_mask_dir)
# Define patch and batch size
patch_size = (64, 64, 64)
batch_size = 2
#-------------------- VALIDATION DATA EVALUATION --------------------
# Create Data Generator for validation set
val_img_loader = imageLoader(val_img_dir, val_img_list, val_mask_dir, val_mask_list,
batch_size, patch_size)
# Fetch a batch for evaluation
val_img_batch, val_mask_batch = val_img_loader.__next__()
# Convert masks to argmax format
val_mask_argmax = np.argmax(val_mask_batch, axis=4)
# Model prediction
val_pred_batch = my_model.predict(val_img_batch)
val_pred_argmax = np.argmax(val_pred_batch, axis=4)
# Compute Mean IoU
#sm.metrics.IOUScore(threshold=0.5)
n_classes = 4
iou_metric = MeanIoU(num_classes=n_classes)
iou_metric.update_state(val_pred_argmax, val_mask_argmax)
val_iou_score = iou_metric.result().numpy()
# Compute Dice Score
dice_metric = DiceScore(num_classes=n_classes)
dice_metric.update_state(val_mask_argmax, val_pred_argmax)
val_dice_score = dice_metric.result().numpy()
print(f"Validation Mean IoU: {val_iou_score}")
print(f"Validation Dice Score: {val_dice_score}")
# -------------------- TEST DATA EVALUATION --------------------
# Create Data Generator for test set
test_img_loader = imageLoader(test_img_dir, test_img_list, test_mask_dir, test_mask_list,
batch_size, patch_size)
# Fetch a batch for evaluation
test_img_batch, test_mask_batch = test_img_loader.__next__()
# Convert masks to argmax format
test_mask_argmax = np.argmax(test_mask_batch, axis=4)
# Model prediction
test_pred_batch = my_model.predict(test_img_batch)
test_pred_argmax = np.argmax(test_pred_batch, axis=4)
# Compute Mean IoU
iou_metric.update_state(test_pred_argmax, test_mask_argmax)
test_iou_score= iou_metric.result().numpy()
# Compute
Dice Score
dice_metric.update_state(test_mask_argmax,test_pred_argmax)
test_dice_score = dice_metric.result().numpy()
print(f"Test Mean IoU: {test_iou_score}")
print(f"Test Dice Score: {test_dice_score}")
Visualizing Results
A random test image is selected, and the model's prediction is visualized alongside the ground truth mask.
#
-------------------- VISUALIZATION --------------------
# Select a random test image for visualization
img_num = 21 # Change index as needed
test_img = np.load(os.path.join(test_img_dir, f"image_{img_num}.npy"))
test_mask = np.load(os.path.join(test_mask_dir, f"mask_{img_num}.npy"))
test_mask_argmax = np.argmax(test_mask, axis=3)
# Expand dimensions for model prediction
test_img_input = np.expand_dims(test_img, axis=0)
test_prediction = my_model.predict(test_img_input)
test_prediction_argmax = np.argmax(test_prediction, axis=4)[0, :, :, :]
# Select slice indices for visualization
slice_indices = [75, 90, 100] # Change slice indices as needed
# Plotting Results
plt.figure(figsize=(18,12))
for i, n_slice in enumerate(slice_indices):
# Rotate images to correct orientation
test_img_rotated = np.rot90(test_img[:, :, n_slice, 1]) #
Rotating 90 degrees
test_mask_rotated = np.rot90(test_mask_argmax[:, :, n_slice])
test_prediction_rotated = np.rot90(test_prediction_argmax[:, :,
n_slice])
# Plotting Results
plt.subplot(3, 4, i*4 + 1)
plt.title(f'Testing Image - Slice {n_slice}')
plt.imshow(test_img_rotated, cmap='gray')
plt.subplot(3, 4, i*4 + 2)
plt.title(f'Ground Truth - Slice {n_slice}')
plt.imshow(test_mask_rotated)
plt.subplot(3, 4, i*4 + 3)
plt.title(f'Prediction - Slice {n_slice}')
plt.imshow(test_prediction_rotated)
plt.subplot(3, 4, i*4 + 4)
plt.title(f'Overlay - Slice {n_slice}')
plt.imshow(test_img_rotated, cmap='gray')
plt.imshow(test_prediction_rotated, alpha=0.5) # Overlay
prediction mask
plt.tight_layout()
plt.show()
Visual Output:
This concludes Part 3: Training and Evaluation. The model has been trained, evaluated, and visualized, demonstrating its performance on brain tumour segmentation tasks. In the next part, we will cover deployment and practical usage of the model. Here is the link to the complete code for Part 3.
PART 4: DEPLOYMENT AND PRACTICAL USAGE
Objective: We will make the model usable for real-world
segmentation tasks.
A. Script Creation for Local Usage
import os
import numpy as np
import nibabel as nib
from sklearn.preprocessing import MinMaxScaler
# Function to preprocess a new NIfTI file
def preprocess_nifti(file_path, patch_size=(64, 64, 64)):
# Load the NIfTI file
image = nib.load(file_path).get_fdata()
# Normalize the image
scaler = MinMaxScaler()
image = scaler.fit_transform(image.reshape(-1, image.shape[-1])).reshape(image.shape)
# Pad or crop the image to match the patch size
if image.shape != patch_size:
# Example: Center cropping
start_x = (image.shape[0] - patch_size[0]) // 2
start_y = (image.shape[1] - patch_size[1]) // 2
start_z = (image.shape[2] - patch_size[2]) // 2
image = image[start_x:start_x + patch_size[0], start_y:start_y +
patch_size[1], start_z:start_z + patch_size[2]]
# Expand dimensions for model input
image = np.expand_dims(image, axis=0) # Add batch dimension
image = np.expand_dims(image, axis=-1) # Add channel dimension
return image
# Function to run segmentation and save results
def run_segmentation(model, input_file, output_folder):
# Preprocess the input file
input_image = preprocess_nifti(input_file)
# Run segmentation
prediction = model.predict(input_image)
prediction_argmax = np.argmax(prediction, axis=4)[0, :, :, :]
# Save the segmentation result
output_file = os.path.join(output_folder, os.path.basename(input_file).replace('.nii.gz', '_segmentation.nii.gz'))
nib.save(nib.Nifti1Image(prediction_argmax, np.eye(4)), output_file)
print(f"Segmentation saved to {output_file}")
# Example usage
input_file = "path/to/new_image.nii.gz"
output_folder = "path/to/output_folder"
run_segmentation(my_model, input_file, output_folder)
B. Deployment on Streamlit
First, we create a python script called “app.py” to host the code to load the model
and the scan image. This script will show the images and the predictions, and
also allow users to download the segmentation output.
Creation of a app.py
import streamlit as st
import numpy as np
import nibabel as nib
from tensorflow.keras.models import load_model
from sklearn.preprocessing import MinMaxScaler
import os
import matplotlib.pyplot as plt
import gdown # For downloading files from Google Drive
import zipfile # To handle folder uploads
import tempfile # To handle temporary files
from tensorflow.keras.utils import to_categorical
# Title of the app
st.title("Brain Tumor Segmentation using 3D U-Net - (Lightweight Architecture on Normal CPUs)")
# Function to download the default model from Google Drive
def download_default_model():
# Google Drive file ID for the
default model
file_id = "1lV1SgafomQKwgv1NW2cjlpyb4LwZXFwX"
# Replace with your file ID
output_path = "default_model.keras"
# Download the file if it doesn't already exist
if not
os.path.exists(output_path):
url =
f"https://drive.google.com/uc?id={file_id}"
gdown.download(url,
output_path, quiet=False)
return output_path
# Load the default model
@st.cache_resource # Cache the
model to avoid reloading on every interaction
def load_default_model():
# Download the model from
Google Drive
model_path = download_default_model()
model = load_model(model_path,
compile=False)
return model
default_model = load_default_model()
# Function to preprocess a NIfTI file
def preprocess_nifti(file_path):
# Load the NIfTI file
image = nib.load(file_path).get_fdata()
# Normalize the image
scaler = MinMaxScaler()
image = scaler.fit_transform(image.reshape(-1, image.shape[-1])).reshape(image.shape)
return image
# Function to combine 4 channels into a single 4-channel numpy array
def combine_channels(t1n, t1c, t2f, t2w):
# Stack the 4 channels along the last axis
combined_image = np.stack([t1n, t1c, t2f, t2w], axis=3)
combined_image = combined_image[56:184, 56:184, 13:141]
return combined_image
# Function to run segmentation
def run_segmentation(model, input_image):
# Add batch and channel dimensions
input_image =np.expand_dims(input_image, axis=0) #Add batch dimension
#input_image = np.expand_dims(input_image, axis=-1) #Add channel dimension
# Ensure the input_image has the correct shape
if len(input_image.shape) != 5:
st.error(f"Unexpected
shape for input_image: {input_image.shape}. Expected shape: (batch_size,
height, width, depth, channels).")
return None
# Run prediction
prediction =
model.predict(input_image)
prediction_argmax =
np.argmax(prediction, axis=4)[0, :, :, :]
return prediction_argmax
# Sidebar for model upload
st.sidebar.header("Upload Your Own Model")
uploaded_model = st.sidebar.file_uploader("Upload a Keras model
(.keras)", type=["keras"])
# Load the model (default or uploaded)
if uploaded_model is not None:
# Save the uploaded model
temporarily
with
open("temp_model.keras", "wb") as f:
f.write(uploaded_model.getbuffer())
# Load the uploaded model
try:
model =
load_model("temp_model.keras", compile=False)
st.sidebar.success("Custom model loaded successfully!")
except Exception as e:
st.sidebar.error(f"Error loading custom model: {e}")
st.sidebar.info("Using the default model instead.")
model = default_model
else:
model = default_model
st.sidebar.info("Using
the default model.")
# Main app: Upload a folder containing NIfTI files
st.header("Upload a Folder Containing NIfTI Files")
uploaded_folder = st.file_uploader("Upload a folder (as a zip file)
containing T1n, T1c, T2f, T2w NIfTI files", type=["zip"])
if uploaded_folder is not None:
# Create a temporary directory to extract the zip file
with
tempfile.TemporaryDirectory() as temp_dir:
# Save the uploaded zip
file
zip_path =
os.path.join(temp_dir, "uploaded_folder.zip")
with open(zip_path,
"wb") as f:
f.write(uploaded_folder.getbuffer())
# Extract the zip file
with
zipfile.ZipFile(zip_path, "r") as zip_ref:
zip_ref.extractall(temp_dir)
# Find the NIfTI files in
the extracted folder
t1n_path = None
t1c_path = None
t2f_path = None
t2w_path = None
mask_path = None
for root, _, files in
os.walk(temp_dir):
for file in files:
if
file.endswith("t1n.nii.gz"):
t1n_path =
os.path.join(root, file)
elif
file.endswith("t1c.nii.gz"):
t1c_path =
os.path.join(root, file)
elif
file.endswith("t2f.nii.gz"):
t2f_path =
os.path.join(root, file)
elif
file.endswith("t2w.nii.gz"):
t2w_path =
os.path.join(root, file)
elif
file.endswith("seg.nii.gz"):
mask_path =
os.path.join(root, file)
# Check if all required
files are found
if t1n_path and t1c_path
and t2f_path and t2w_path:
# Preprocess each
channel
t1n =
preprocess_nifti(t1n_path)
t1c =
preprocess_nifti(t1c_path)
t2f =
preprocess_nifti(t2f_path)
t2w =
preprocess_nifti(t2w_path)
# Combine the 4
channels
combined_image =
combine_channels(t1n, t1c, t2f, t2w)
# Print the shape of
combined_image for debugging
st.write(f"Shape
of combined_image: {combined_image.shape}")
# Ensure the
combined_image has the correct shape
if
len(combined_image.shape) != 4:
st.error(f"Unexpected shape for combined_image:
{combined_image.shape}. Expected shape: (height, width, depth,
channels).")
else:
# Run segmentation
st.write("Running segmentation...")
segmentation_result = run_segmentation(model, combined_image)
# Display the
segmentation result
st.write("Segmentation completed! Displaying results...")
# Load the ground
truth mask if available
if mask_path:
mask =
nib.load(mask_path).get_fdata()
mask =
mask.astype(np.uint8)
mask[mask ==
4] = 3 # Reassign mask values 4 to 3
mask_argmax =
np.argmax(to_categorical(mask, num_classes=4), axis=3)
else:
mask_argmax =
None
# Select slice
indices for visualization
slice_indices =
[75, 90, 100] # Example: 25%, 50%, 75%
of depth
# Plotting Results
fig, ax =
plt.subplots(3, 4, figsize=(18, 12))
for i, n_slice in
enumerate(slice_indices):
# Rotate
images to correct orientation
test_img_rotated = np.rot90(combined_image[:, :, n_slice, 0]) # Rotating 90 degrees
test_prediction_rotated = np.rot90(segmentation_result[:, :, n_slice])
# Plotting
Results
ax[i,
0].imshow(test_img_rotated, cmap='gray')
ax[i,
0].set_title(f'Testing Image - Slice {n_slice}')
if mask_path:
test_mask_rotated = np.rot90(mask_argmax[:, :, n_slice])
ax[i,
1].imshow(test_mask_rotated)
ax[i,
1].set_title(f'Ground Truth - Slice {n_slice}')
else:
ax[i,
1].axis('off') # Hide the ground truth
subplot if no mask is available
ax[i,
2].imshow(test_prediction_rotated)
ax[i,
2].set_title(f'Prediction - Slice {n_slice}')
ax[i,
3].imshow(test_img_rotated, cmap='gray')
ax[i,
3].imshow(test_prediction_rotated, alpha=0.5)
# Overlay prediction mask
ax[i,
3].set_title(f'Overlay - Slice {n_slice}')
plt.tight_layout()
st.pyplot(fig)
# Save the
segmentation result
output_file =
"segmentation_result.nii.gz"
nib.save(nib.Nifti1Image(segmentation_result.astype(np.float32),
np.eye(4)), output_file)
# Provide a
download link for the segmentation result
with
open(output_file, "rb") as f:
st.download_button(
label="Download Segmentation Result",
data=f,
file_name=output_file,
mime="application/octet-stream"
)
# Clean up
temporary files
os.remove(output_file)
else:
st.error("The
uploaded folder does not contain all required NIfTI files (T1n, T1c, T2f,
T2w).")
Save
the script in a “app.py” file.
Creation of a requirements.txt
streamlit
numpy
nibabel
tensorflow
scikit-learn
matplotlib
scipy
gdown
Deploy on Streamlit
The application was successfully deployed on Streamlit by hosting the necessary files on a GitHub repository. The GitHub account wasseamlessly integrated with Streamlit through their website, enabling a smooth deployment process. The deployed model can be accessed at the following link: https://brain-tumor-segmenatation-on-cpu-fnjvqbcqsjrxd9juxxfvgn.streamlit.app/
Thank you !!. That brings us to the end of this tutorial.
Additional Resources
- Links to the full code repository on GitHub. – “https://github.com/Heartz00/Brain-Tumor-segmenatation-on-CPU/tree/main”
- References to relevant papers and datasets:
Adewole, M., Rudie, J.D., Gbadamosi, A.,Zhang, D., Raymond, C., Ajigbotoshso, J., Toyobo, O., Aguh, K., Omidiji, O., Akinola R., Suwaid, M.A., Emegoakor, A., Ojo, N., Kalaiwo, C., Babatunde, G., Ogunleye, A., Gbadamosi, Y., Iorpagher, K., Onuwaje M., Betiku B., Saluja, R., Menze, B., Baid, U., Bakas, S., Dako, F., Fatade A., Anazodo, U.C. (2024) Expanding the Brain Tumor Segmentation (BraTS) data to include African
Populations (BraTS-Africa) (version 1) [Dataset]. The Cancer Imaging Archive. https://doi.org/10.7937/v8h6-8×67
Bhattiprolu, S. (n.d.). BraTa2020 Unet segmentation . GitHub. Retrieved [January, 2025], from https://github.com/bnsreenu/python_for_microscopists/tree/master/231_234_BraTa2020_Unet_segmentation
Oladele, A.B. et al. (2026). Lightweight 3D U-Net for Brain Tumor Segmentation on CPUs: Enabling Deep Learning in Low-Resource Environments. In: Anazodo, U., Zhang, D., Raymond, C., Kurt, M., Lekadir, K., Crimi, A. (eds) Medical Image Computing in Resource Constrained Settings. MIRASOL 2025. Lecture Notes in Computer Science, vol 16398. Springer, Cham. https://doi.org/10.1007/978-3-032-13654-1_6
Protocol references
Adewole, M., Rudie, J.D., Gbadamosi, A.,Zhang, D., Raymond, C., Ajigbotoshso, J., Toyobo, O., Aguh, K., Omidiji, O., Akinola R., Suwaid, M.A., Emegoakor, A., Ojo, N., Kalaiwo, C., Babatunde, G., Ogunleye, A., Gbadamosi, Y., Iorpagher, K., Onuwaje M., Betiku B., Saluja, R., Menze, B., Baid, U., Bakas, S., Dako, F., Fatade A., Anazodo, U.C. (2024) Expanding the Brain Tumor Segmentation (BraTS) data to include African
Populations (BraTS-Africa) (version 1) [Dataset]. The Cancer Imaging Archive. https://doi.org/10.7937/v8h6-8×67
Bhattiprolu, S. (n.d.). BraTa2020 Unet segmentation . GitHub. Retrieved [January, 2025], from https://github.com/bnsreenu/python_for_microscopists/tree/master/231_234_BraTa2020_Unet_segmentation
Oladele, A.B. et al. (2026). Lightweight 3D U-Net for Brain Tumor Segmentation on CPUs: Enabling Deep Learning in Low-Resource Environments. In: Anazodo, U., Zhang, D., Raymond, C., Kurt, M., Lekadir, K., Crimi, A. (eds) Medical Image Computing in Resource Constrained Settings. MIRASOL 2025. Lecture Notes in Computer Science, vol 16398. Springer, Cham. https://doi.org/10.1007/978-3-032-13654-1_6