import os
import cv2
import random
import glob
import PIL
import shutil
import numpy as np
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
from skimage import data
from skimage.util import montage
import skimage.transform as skTrans
from skimage.transform import rotate
from skimage.transform import resize
from PIL import Image, ImageOps
import nibabel as nib
import keras
import keras.backend as K
from keras.callbacks import CSVLogger
import tensorflow as tf
from tensorflow.keras.utils import plot_model
from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
from tensorflow.keras.models import *
from tensorflow.keras.layers import *
from tensorflow.keras.optimizers import *
from tensorflow.keras.callbacks import ModelCheckpoint, ReduceLROnPlateau, EarlyStopping, TensorBoard
from tensorflow.keras.layers.experimental import preprocessing
TRAIN_DATASET_PATH = "/kaggle/input/brats20-dataset-training-validation/BraTS2020_TrainingData/MICCAI_BraTS2020_TrainingData/"
# load .nii file as a numpy array
test_image_flair = nib.load(TRAIN_DATASET_PATH + "BraTS20_Training_355/BraTS20_Training_355_flair.nii").get_fdata()
print("Shape: ", test_image_flair.shape)
print("Dtype: ", test_image_flair.dtype)
print("Min: ", test_image_flair.min())
print("Max: ", test_image_flair.max())
As we can see the maximum pixel value of the image is 1854.6 which gives us a good idea to rescale these values. Rescaling pixel values is essential because it standardizes the intensity values across multimodal MRI scans (T1, T1ce, T2, FLAIR) to a uniform scale. This uniformity is crucial for consistent model training and analysis, as it ensures that variations in pixel intensity due to different imaging modalities and equipment do not bias the segmentation model. It facilitates the model's ability to learn meaningful features from the data for accurate tumor segmentation, regardless of the inherent differences in the raw MRI scans.
scaler = MinMaxScaler()
# Scale the test_image_flair array and then reshape it back to its original dimensions.
# This ensures the data is normalized/standardized for model input without altering its spatial structure.
test_image_flair = scaler.fit_transform(test_image_flair.reshape(-1, test_image_flair.shape[-1])).reshape(test_image_flair.shape)
print("Min: ", test_image_flair.min())
print("Max: ", test_image_flair.max())
# rescaling t1
test_image_t1 = nib.load(TRAIN_DATASET_PATH + 'BraTS20_Training_354/BraTS20_Training_354_t1.nii').get_fdata()
test_image_t1 = scaler.fit_transform(test_image_t1.reshape(-1, test_image_t1.shape[-1])).reshape(test_image_t1.shape)
# rescaling t1ce
test_image_t1ce = nib.load(TRAIN_DATASET_PATH + 'BraTS20_Training_354/BraTS20_Training_354_t1ce.nii').get_fdata()
test_image_t1ce = scaler.fit_transform(test_image_t1ce.reshape(-1, test_image_t1ce.shape[-1])).reshape(test_image_t1ce.shape)
# rescaling t2
test_image_t2 = nib.load(TRAIN_DATASET_PATH + 'BraTS20_Training_354/BraTS20_Training_354_t2.nii').get_fdata()
test_image_t2 = scaler.fit_transform(test_image_t2.reshape(-1, test_image_t2.shape[-1])).reshape(test_image_t2.shape)
# we will not rescale the mask
test_image_seg = nib.load(TRAIN_DATASET_PATH + 'BraTS20_Training_354/BraTS20_Training_354_seg.nii').get_fdata()
slice = 95
print("Slice Number: " + str(slice))
plt.figure(figsize=(12, 8))
# T1
plt.subplot(2, 3, 1)
plt.imshow(test_image_t1[:,:,slice], cmap='gray')
plt.title('T1')
# T1ce
plt.subplot(2, 3, 2)
plt.imshow(test_image_t1ce[:,:,slice], cmap='gray')
plt.title('T1ce')
# T2
plt.subplot(2, 3, 3)
plt.imshow(test_image_t2[:,:,slice], cmap='gray')
plt.title('T2')
# Flair
plt.subplot(2, 3, 4)
plt.imshow(test_image_flair[:,:,slice], cmap='gray')
plt.title('FLAIR')
# Mask
plt.subplot(2, 3, 5)
plt.imshow(test_image_seg[:,:,slice])
plt.title('Mask')
plt.show()
Images Format¶
You may have noticed that these images are in the .nii format. These scans are NIfTI files (Neuroimaging Informatics Technology Initiative). A NIfTI image provides a digital representation of a 3D object, such as a brain in our case. The NIfTI format is widely used in neuroimaging because it efficiently handles complex, multi-dimensional data and includes metadata for spatial orientation and scaling.
To gain a better understanding of the data structure, let's examine the shape of a modality scan and a segmentation mask. Note that in this dataset, all modalities and their corresponding segmentation masks share the same dimensions:
# Modality shape
print("Modality: ", test_image_t1.shape)
# Segmentation shape
print("Segmentation: ", test_image_seg.shape)
As we can see, the modalities and segmentations have 3 dimensions. Each dimension consists of a series of two-dimensional images, known as slices, which all contain the same number of pixels and are stacked together to create this 3D representation. For instance, in our previous example, we displayed the 95th slice of a certain dimension.
These 3 dimensions correspond to the three spatial dimensions of the image: width, height, and depth. In medical imaging, these dimensions are referred to as the axial (transverse), coronal, and sagittal planes, corresponding to the three main orientations of the human body and, therefore, the human brain.
Here is a quick presentation of these 3 planes:
- Axial (Transverse) Plane: This plane divides the body into upper and lower parts and is perpendicular to the long axis of the body. In brain imaging, an axial slice is a horizontal cut through the brain.
- Coronal (Frontal) Plane: This plane divides the body into front (anterior) and back (posterior) parts. A coronal slice in brain imaging is a vertical cut from one side of the head to the other, dividing the brain into front and back sections.
- Sagittal (Lateral) Plane: This plane divides the body into left and right parts. A sagittal slice in brain imaging is a vertical cut from front to back, dividing the brain into left and right sections.
Understanding these planes is crucial for accurately interpreting medical images and effectively utilizing them in segmentation tasks. Each plane provides a different perspective, helping to localize and identify anatomical structures and abnormalities.
slice = 95
print("Slice number: " + str(slice))
plt.figure(figsize=(12, 8))
# Apply a 90° rotation with an automatic resizing, otherwise the display is less obvious to analyze
# T1 - Transverse View
plt.subplot(1, 3, 1)
plt.imshow(test_image_t1ce[:,:,slice], cmap='gray')
plt.title('T1 - Transverse View')
# T1 - Frontal View
plt.subplot(1, 3, 2)
plt.imshow(rotate(test_image_t1ce[:,slice,:], 90, resize=True), cmap='gray')
plt.title('T1 - Frontal View')
# T1 - Sagittal View
plt.subplot(1, 3, 3)
plt.imshow(rotate(test_image_t1ce[slice,:,:], 90, resize=True), cmap='gray')
plt.title('T1 - Sagittal View')
plt.show()
This visualization helps illustrate how each plane cuts through the brain and highlights different anatomical features, aiding in comprehensive analysis and segmentation.
Now that we understand the three dimensions, let's explore why we chose to display a specific slice, setting the slice number to 95. By examining all the slices, we can identify the slice that best represents the region of interest, such as the area where a tumor is located.
To visualize this, we will display all the slices along one dimension:
plt.figure(figsize=(10, 10))
plt.subplot(1, 1, 1)
# montage allows us to concatenate multiple images of the same size horizontally and vertically
plt.imshow(rotate(montage(test_image_t1ce[:,:,:]), 90, resize=True), cmap ='gray');
As you can see, two black parts are present on each side of our montage. These black parts correspond to the first and last slices of the plane. This indicates that a large portion of the slices does not contain much information, which is expected as the slices progressively represent the brain from one end to the other.
This observation is consistent across all modalities, planes, and even the expert-segmented images. The experts were unable to segment the slices that lack significant information, as these slices mostly capture areas outside the brain or the outermost regions that do not contain relevant anatomical details.
To further illustrate this, let's visualize a range of slices:
# Skip 50:-50 slices since there is not much to see
plt.figure(figsize=(10, 10))
plt.subplot(1, 1, 1)
plt.imshow(rotate(montage(test_image_t1ce[50:-50,:,:]), 90, resize=True), cmap ='gray');
# Skip 50:-50 slices since there is not much to see
plt.figure(figsize=(10, 10))
plt.subplot(1, 1, 1)
plt.imshow(rotate(montage(test_image_seg[50:-50,:,:]), 90, resize=True), cmap ='gray');
By focusing on a range of slices that contain meaningful content, we can better understand the anatomical structures and the regions where abnormalities, such as tumors, may be present. This approach helps in identifying the slices that are most relevant for analysis and segmentation.
This analysis applies to all other modalities and planes as well. The segmented images by experts also reflect this pattern, as they concentrate on the slices containing significant information. Slices that do not capture much of the brain's anatomy are often left unsegmented, as they provide little to no valuable data for diagnosis or analysis.
Now, let's focus on the segmentations provided by the experts. These segmentations provide critical information about the tumor:
- Location: Identifies the exact location of the tumor in the brain.
- Size and Shape: Measures the tumor's dimensions and form.
- Type: Helps differentiate between various tumor types.
- Heterogeneity: Reveals different components within the tumor (e.g., solid tissue, necrotic areas).
- Surrounding Tissue Involvement: Shows how the tumor interacts with adjacent brain structures.
- Cross-Modality Comparison: Provides a comprehensive view by comparing different imaging modalities.
# Plotting the segmantation
cmap = matplotlib.colors.ListedColormap(['#440054', '#3b528b', '#18b880', '#e6d74f'])
norm = matplotlib.colors.BoundaryNorm([-0.5, 0.5, 1.5, 2.5, 3.5], cmap.N)
# plotting the 95th slice
plt.imshow(test_image_seg[:,:,95], cmap=cmap, norm=norm)
plt.colorbar()
plt.show()
# Isolation of class 0
seg_0 = test_image_seg.copy()
seg_0[seg_0 != 0] = np.nan
# Isolation of class 1
seg_1 = test_image_seg.copy()
seg_1[seg_1 != 1] = np.nan
# Isolation of class 2
seg_2 = test_image_seg.copy()
seg_2[seg_2 != 2] = np.nan
# Isolation of class 4
seg_4 = test_image_seg.copy()
seg_4[seg_4 != 4] = np.nan
# Define legend
class_names = ['class 0', 'class 1', 'class 2', 'class 4']
legend = [plt.Rectangle((0, 0), 1, 1, color=cmap(i), label=class_names[i]) for i in range(len(class_names))]
fig, ax = plt.subplots(1, 5, figsize=(20, 20))
ax[0].imshow(test_image_seg[:,:, slice], cmap=cmap, norm=norm)
ax[0].set_title('Original Segmentation')
ax[0].legend(handles=legend, loc='lower left')
ax[1].imshow(seg_0[:,:, slice], cmap=cmap, norm=norm)
ax[1].set_title('Not Tumor (class 0)')
ax[2].imshow(seg_1[:,:, slice], cmap=cmap, norm=norm)
ax[2].set_title('Non-Enhancing Tumor (class 1)')
ax[3].imshow(seg_2[:,:, slice], cmap=cmap, norm=norm)
ax[3].set_title('Edema (class 2)')
ax[4].imshow(seg_4[:,:, slice], cmap=cmap, norm=norm)
ax[4].set_title('Enhancing Tumor (class 4)')
plt.show()
These insights are essential for accurate diagnosis, treatment planning, and monitoring. The segmentations help visualize the tumor clearly, making it easier to understand its characteristics and interactions with surrounding tissues.
To summarize our data exploration:
- Modalities: Each sample/patient has 4 different modalities (T1, T1CE, T2, and FLAIR), accompanied by a segmentation indicating tumor areas.
- Selected Modalities: We will focus on T1CE and FLAIR as they provide complementary information about the brain's anatomy and tissue contrast.
- 3D Images: Each image is 3D and can be analyzed through 3 different planes composed of 2D slices.
- Informative Slices: Many slices contain little or no information. We will use slices in the (50:-50) interval. You can customize this range, but note that changing it may affect training time.
- Segmentation Classes: Segmentation images contain 1 to 4 classes. Class 4 will be reassigned to 3 since class 3 is missing.
- Background Class: Class 0 (background) is over-represented. Cropping might remove important information, so we will keep the images as they are.
Now that we understand our data better, it's time to prepare for model training.
Split the Dataset¶
To train and evaluate our model effectively, we need to split our dataset into three parts: the training set (70-80%) for training the model, the validation set (10-15%) for tuning hyperparameters and preventing overfitting, and the test set (10-15%) for evaluating the final model's performance. We can use random splitting to divide the data or stratified splitting to maintain the same class distribution across sets, which is particularly useful for imbalanced datasets. Properly splitting the dataset ensures that our model is robust and performs well on unseen data.
# lists of directories with studies
train_and_val_directories = [f.path for f in os.scandir(TRAIN_DATASET_PATH) if f.is_dir()]
def pathListIntoIds(dirList):
x = []
for i in range(0,len(dirList)):
x.append(dirList[i][dirList[i].rfind('/')+1:])
return x
train_and_test_ids = pathListIntoIds(train_and_val_directories);
train_test_ids, val_ids = train_test_split(train_and_test_ids,test_size=0.2)
train_ids, test_ids = train_test_split(train_test_ids,test_size=0.15)
# Print data distribution (Train: 68%, Test: 12%, Val: 20%)
print(f"Train length: {len(train_ids)}")
print(f"Validation length: {len(val_ids)}")
print(f"Test length: {len(test_ids)}")
plt.bar(["Train","Valid","Test"],
[len(train_ids), len(val_ids), len(test_ids)],
align='center',
color=[ 'green','red', 'blue'],
label=["Train", "Valid", "Test"]
)
plt.legend()
plt.ylabel('Number of Images')
plt.title('Data Distribution')
plt.show()
DataGenerator¶
To train a neural network for image segmentation, we need both raw image data (X) and ground truth segmentations (y). This allows the network to learn tumor patterns and make accurate predictions from patient scans. However, directly loading all 3D images can overload system memory and cause shape mismatch errors. Therefore, we use a Data Generator for image preprocessing, which includes several steps:
- Retrieve paths: Obtain the paths for the T1CE and FLAIR modalities (for complementary anatomical and tissue contrast information) and the ground truth segmentation.
- Load data: Load the selected slices (60-135) from these modalities and the corresponding segmentation.
- Create arrays: Form X arrays with the slices from T1CE and FLAIR, and y arrays with the segmentation slices.
- Class reassignment: Assign a value of 3 to all instances of 4 in the mask array to correct the missing class issue.
Additional preprocessing steps include:
- Axial plane: Use the axial plane for its square shape (240x240). This enables visualizing predictions across all planes without impact.
- One-Hot Encoding: Apply One-Hot Encoding to the y array to convert classes (0 to 3) into a numerical format suitable for neural networks, avoiding any implied hierarchy between classes. Here is what it consists of, for one slice:
- Resize images: Resize each slice from (240x240) to (128x128). This shape is chosen because it is a power of two, fitting well with pooling layers (MaxPooling2D) in CNNs, and balances computational efficiency and information preservation.
While resizing to (256x256) could retain more detail, it significantly increases training time and memory usage. We can choose to train with (256x256) images, but we’ll need to adjust your U-Net architecture accordingly.
Now that we have a clear understanding of our data and preprocessing steps, we can proceed to prepare our model for training.
# Define seg-areas
SEGMENT_CLASSES = {
0 : 'NOT tumor',
1 : 'NECROTIC/CORE', # or NON-ENHANCING tumor CORE
2 : 'EDEMA',
3 : 'ENHANCING' # original 4 -> converted into 3
}
# Select Slices and Image Size
VOLUME_SLICES = 100
VOLUME_START_AT = 22 # first slice of volume that we will include
IMG_SIZE=128
import os
# Function to filter out any patients missing their segmentation mask
def clean_id_list(id_list):
cleaned_list = []
for patient_id in id_list:
# Construct the exact path the generator will look for
case_path = os.path.join(TRAIN_DATASET_PATH, patient_id)
seg_path = os.path.join(case_path, f'{patient_id}_seg.nii')
# Only keep the ID if the file actually exists on the hard drive
if os.path.exists(seg_path):
cleaned_list.append(patient_id)
else:
print(f"🗑️ Removed corrupted/missing patient: {patient_id}")
return cleaned_list
# Clean all three lists
print("Scanning lists for corrupted files...")
train_ids = clean_id_list(train_ids)
val_ids = clean_id_list(val_ids)
test_ids = clean_id_list(test_ids)
# Re-initialize the generators with the cleaned lists
training_generator = DataGenerator(train_ids)
valid_generator = DataGenerator(val_ids)
test_generator = DataGenerator(test_ids)
class DataGenerator(keras.utils.Sequence):
'Generates data for Keras'
def __init__(self, list_IDs, dim=(IMG_SIZE,IMG_SIZE), batch_size=1, n_channels=2, shuffle=True):
'Initialization'
self.dim = dim
self.batch_size = batch_size
self.list_IDs = list_IDs
self.n_channels = n_channels
self.shuffle = shuffle
self.on_epoch_end()
# NEW: Initialize CLAHE and Morphological Kernel once per generator
self.clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
self.morph_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3,3))
def __len__(self):
'Denotes the number of batches per epoch'
return int(np.floor(len(self.list_IDs) / self.batch_size))
def __getitem__(self, index):
'Generate one batch of data'
indexes = self.indexes[index*self.batch_size:(index+1)*self.batch_size]
Batch_ids = [self.list_IDs[k] for k in indexes]
X, y = self.__data_generation(Batch_ids)
return X, y
def on_epoch_end(self):
'Updates indexes after each epoch'
self.indexes = np.arange(len(self.list_IDs))
if self.shuffle == True:
np.random.shuffle(self.indexes)
def __data_generation(self, Batch_ids):
'Generates data containing batch_size samples'
X = np.zeros((self.batch_size*VOLUME_SLICES, *self.dim, self.n_channels), dtype=np.float32)
y = np.zeros((self.batch_size*VOLUME_SLICES, 240, 240), dtype=np.float32)
# 1. Generate and Enhance MRI Data (CLAHE)
for c, i in enumerate(Batch_ids):
case_path = os.path.join(TRAIN_DATASET_PATH, i)
data_path = os.path.join(case_path, f'{i}_flair.nii')
flair = nib.load(data_path).get_fdata()
data_path = os.path.join(case_path, f'{i}_t1ce.nii')
t1ce = nib.load(data_path).get_fdata()
data_path = os.path.join(case_path, f'{i}_seg.nii')
seg = nib.load(data_path).get_fdata()
for j in range(VOLUME_SLICES):
# Extract slices
f_slice = flair[:,:,j+VOLUME_START_AT]
t_slice = t1ce[:,:,j+VOLUME_START_AT]
# Normalize to 0-255 for CLAHE processing slice-by-slice
f_max, t_max = np.max(f_slice), np.max(t_slice)
if f_max > 0: f_slice = (f_slice / f_max * 255).astype(np.uint8)
else: f_slice = f_slice.astype(np.uint8)
if t_max > 0: t_slice = (t_slice / t_max * 255).astype(np.uint8)
else: t_slice = t_slice.astype(np.uint8)
# Apply CLAHE
f_slice = self.clahe.apply(f_slice)
t_slice = self.clahe.apply(t_slice)
# Resize and scale back to 0.0 - 1.0 for the Neural Network
X[j +VOLUME_SLICES*c,:,:,0] = cv2.resize(f_slice, (IMG_SIZE, IMG_SIZE)) / 255.0
X[j +VOLUME_SLICES*c,:,:,1] = cv2.resize(t_slice, (IMG_SIZE, IMG_SIZE)) / 255.0
# Assign raw mask
y[j +VOLUME_SLICES*c] = seg[:,:,j+VOLUME_START_AT]
# 2. Process Masks
y[y==4] = 3
mask = tf.one_hot(y, 4)
Y = tf.image.resize(mask, (IMG_SIZE, IMG_SIZE)).numpy() # Convert to numpy for OpenCV
# 3. Enhance Masks (Morphological Smoothing)
for b in range(Y.shape[0]): # Loop through the batch
for channel in range(1, 4): # Only smooth tumor channels (1, 2, 3)
single_class_mask = Y[b, :, :, channel].astype(np.uint8)
# Apply Closing (fills holes), then Opening (removes jagged noise)
smoothed = cv2.morphologyEx(single_class_mask, cv2.MORPH_CLOSE, self.morph_kernel)
smoothed = cv2.morphologyEx(smoothed, cv2.MORPH_OPEN, self.morph_kernel)
Y[b, :, :, channel] = smoothed
# Recalculate the Background channel (0) so it perfectly wraps around the smoothed tumors
Y[b, :, :, 0] = 1.0 - np.clip(np.sum(Y[b, :, :, 1:], axis=-1), 0.0, 1.0)
return np.array(X, dtype=np.float32), np.array(Y, dtype=np.float32)
# Reinitialize generators with the newly enhanced pipeline
training_generator = DataGenerator(train_ids)
valid_generator = DataGenerator(val_ids)
test_generator = DataGenerator(test_ids)
Preprocessing Steps Summary¶
- Data Generator: Utilized to process and send data to the neural network without overloading memory.
- Epoch Handling: For each epoch, the model processes 250 samples from the training dataset.
- Sample Analysis: Each sample consists of 150 slices (100 slices each from two modalities) resized to (128, 128).
- Data Shapes:
- X Array: Shape (128, 128, 100, 2) for input images.
- Ground Truth (y): One-Hot encoded segmentation with shape (100, 128, 128, 4).
This preprocessing setup ensures efficient data handling and preparation for training the neural network.
# Define a function to display one slice and its segmentation
def display_slice_and_segmentation(flair, t1ce, segmentation):
fig, axes = plt.subplots(1, 3, figsize=(10, 5))
axes[0].imshow(flair, cmap='gray')
axes[0].set_title('Flair')
axes[0].axis('off')
axes[1].imshow(t1ce, cmap='gray')
axes[1].set_title('T1CE')
axes[1].axis('off')
axes[2].imshow(segmentation) # Displaying segmentation
axes[2].set_title('Segmentation')
axes[2].axis('off')
plt.tight_layout()
plt.show()
# Retrieve the batch from the training generator
X_batch, Y_batch = training_generator[8]
# Extract Flair, T1CE, and segmentation from the batch
flair_batch = X_batch[:, :, :, 0]
t1ce_batch = X_batch[:, :, :, 1]
segmentation_batch = np.argmax(Y_batch, axis=-1) # Convert one-hot encoded to categorical
# Extract the 50th slice from Flair, T1CE, and segmentation
slice_index = 60 # Indexing starts from 0
slice_flair = flair_batch[slice_index]
slice_t1ce = t1ce_batch[slice_index]
slice_segmentation = segmentation_batch[slice_index]
# Display the 50th slice and its segmentation
display_slice_and_segmentation(slice_flair, slice_t1ce, slice_segmentation)
Loss Function and Evaluation Metrics¶
Loss Function¶
When training a CNN, selecting an appropriate loss function is crucial for accurately reflecting the network's performance. This function compares the predicted pixels to the ground truth for each patient, guiding the model to update its weights to minimize the loss and improve prediction accuracy.
For multi-class classification problems, a commonly used loss function is categorical cross-entropy. This function measures the difference between the predicted probability distribution of each pixel and the one-hot encoded ground truth values. Additionally, segmentation models sometimes use the Dice loss function, which focuses on the overlap between the predicted and actual segments, further refining the accuracy of the segmentation.
Using these loss functions helps ensure that the model learns effectively and produces accurate segmentations of brain tumors.
import tensorflow as tf
# dice loss as defined above for 4 classes
def dice_coef(y_true, y_pred, smooth=1.0):
class_num = 4
total_loss = 0.0 # Initialize total_loss outside the loop for cleaner syntax
for i in range(class_num):
# Replace K.flatten with tf.reshape to flatten the arrays
y_true_f = tf.reshape(y_true[:,:,:,i], [-1])
y_pred_f = tf.reshape(y_pred[:,:,:,i], [-1])
# Replace K.sum with tf.reduce_sum
intersection = tf.reduce_sum(y_true_f * y_pred_f)
loss = ((2. * intersection + smooth) / (tf.reduce_sum(y_true_f) + tf.reduce_sum(y_pred_f) + smooth))
total_loss = total_loss + loss
total_loss = total_loss / class_num
return total_loss
Per Class Dice Coefficient Functions¶
These functions calculate the Dice coefficient for specific tumor classes in segmentation tasks, helping to measure the model's performance in predicting different tumor regions:
dice_coef_necrotic: Calculates the Dice coefficient for the necrotic (dead tissue) tumor region. It computes the intersection over the sum of squares of the true and predicted values for the necrotic class.
dice_coef_edema: Calculates the Dice coefficient for the edema (swelling) tumor region. It computes the intersection over the sum of squares of the true and predicted values for the edema class.
dice_coef_enhancing: Calculates the Dice coefficient for the enhancing tumor region. It computes the intersection over the sum of squares of the true and predicted values for the enhancing class.
import tensorflow as tf
# Define per class evaluation of dice coef
def dice_coef_necrotic(y_true, y_pred, epsilon=1e-6):
# Replace K.sum, K.abs, and K.square with their tf equivalents
intersection = tf.reduce_sum(tf.abs(y_true[:,:,:,1] * y_pred[:,:,:,1]))
return (2. * intersection) / (tf.reduce_sum(tf.square(y_true[:,:,:,1])) + tf.reduce_sum(tf.square(y_pred[:,:,:,1])) + epsilon)
def dice_coef_edema(y_true, y_pred, epsilon=1e-6):
intersection = tf.reduce_sum(tf.abs(y_true[:,:,:,2] * y_pred[:,:,:,2]))
return (2. * intersection) / (tf.reduce_sum(tf.square(y_true[:,:,:,2])) + tf.reduce_sum(tf.square(y_pred[:,:,:,2])) + epsilon)
def dice_coef_enhancing(y_true, y_pred, epsilon=1e-6):
intersection = tf.reduce_sum(tf.abs(y_true[:,:,:,3] * y_pred[:,:,:,3]))
return (2. * intersection) / (tf.reduce_sum(tf.square(y_true[:,:,:,3])) + tf.reduce_sum(tf.square(y_pred[:,:,:,3])) + epsilon)
Each function includes a small constant (epsilon) to avoid division by zero.
Evaluation Metrics¶
To effectively monitor the model's performance, we use various evaluation metrics:
- Accuracy: Measures the overall proportion of correctly classified pixels. However, it can be misleading with imbalanced datasets like BraTS2020, where the background class is overrepresented.
- Intersection over Union (IoU): Measures the overlap between the predicted and ground truth segmentations.
- Dice Coefficient: Evaluates the similarity between predicted and ground truth segmentations.
- Sensitivity (Recall or True Positive Rate): Measures the proportion of positive ground truth pixels correctly predicted as positive.
- Precision (Positive Predictive Value): Measures the proportion of predicted positive pixels that are actually positive.
- Specificity (True Negative Rate): Measures the proportion of negative ground truth pixels correctly predicted as negative.
Using these metrics provides a comprehensive assessment of the model’s performance, addressing the limitations of using accuracy alone on imbalanced datasets.
import tensorflow as tf
import keras
# Computing Precision
def precision(y_true, y_pred):
y_true_f = tf.cast(y_true, tf.float32)
y_pred_f = tf.cast(y_pred, tf.float32)
true_positives = tf.reduce_sum(tf.round(tf.clip_by_value(y_true_f * y_pred_f, 0.0, 1.0)))
predicted_positives = tf.reduce_sum(tf.round(tf.clip_by_value(y_pred_f, 0.0, 1.0)))
return true_positives / (predicted_positives + keras.backend.epsilon())
# Computing Sensitivity (Recall)
def sensitivity(y_true, y_pred):
y_true_f = tf.cast(y_true, tf.float32)
y_pred_f = tf.cast(y_pred, tf.float32)
true_positives = tf.reduce_sum(tf.round(tf.clip_by_value(y_true_f * y_pred_f, 0.0, 1.0)))
possible_positives = tf.reduce_sum(tf.round(tf.clip_by_value(y_true_f, 0.0, 1.0)))
return true_positives / (possible_positives + keras.backend.epsilon())
# Computing Specificity
def specificity(y_true, y_pred):
y_true_f = tf.cast(y_true, tf.float32)
y_pred_f = tf.cast(y_pred, tf.float32)
true_negatives = tf.reduce_sum(tf.round(tf.clip_by_value((1.0 - y_true_f) * (1.0 - y_pred_f), 0.0, 1.0)))
possible_negatives = tf.reduce_sum(tf.round(tf.clip_by_value(1.0 - y_true_f, 0.0, 1.0)))
return true_negatives / (possible_negatives + keras.backend.epsilon())
Define the Segmentation Model¶
U-Net¶
We will use the U-Net architecture, a convolutional neural network (CNN) specifically designed for biomedical image segmentation. U-Net is particularly effective for segmenting regions of interest that are small and have complex shapes, such as tumors in MRI scans. This architecture was introduced in 2015 by Olaf Ronneberger, Philipp Fischer, and Thomas Brox in their paper U-Net: Convolutional Networks for Biomedical Image Segmentation.
Since the BraTS2020 dataset consists of 3D images, with each image comprising multiple 2D slices in three orthogonal planes, we have two options: using a 2D U-Net or a 3D U-Net.
- 3D U-Net: More suitable for leveraging the 3D spatial context of the images, reducing the risk of false positives and false negatives from partial information in individual slices. However, it requires more computational resources and memory.
- 2D U-Net: Faster and requires less memory, advantageous when working with large datasets or limited computational resources.
In practice, it’s useful to try both architectures and compare their performance. We will opt for a 2D U-Net implementation due to its efficiency and lower resource requirements.
def build_unet(inputs, ker_init, dropout):
conv1 = Conv2D(32, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(inputs)
conv1 = Conv2D(32, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(conv1)
pool = MaxPooling2D(pool_size=(2, 2))(conv1)
conv = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(pool)
conv = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(conv)
pool1 = MaxPooling2D(pool_size=(2, 2))(conv)
conv2 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(pool1)
conv2 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(conv2)
pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)
conv3 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(pool2)
conv3 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(conv3)
pool4 = MaxPooling2D(pool_size=(2, 2))(conv3)
conv5 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(pool4)
conv5 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(conv5)
drop5 = Dropout(dropout)(conv5)
up7 = Conv2D(256, 2, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(UpSampling2D(size = (2,2))(drop5))
merge7 = concatenate([conv3,up7], axis = 3)
conv7 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(merge7)
conv7 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(conv7)
up8 = Conv2D(128, 2, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(UpSampling2D(size = (2,2))(conv7))
merge8 = concatenate([conv2,up8], axis = 3)
conv8 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(merge8)
conv8 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(conv8)
up9 = Conv2D(64, 2, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(UpSampling2D(size = (2,2))(conv8))
merge9 = concatenate([conv,up9], axis = 3)
conv9 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(merge9)
conv9 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(conv9)
up = Conv2D(32, 2, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(UpSampling2D(size = (2,2))(conv9))
merge = concatenate([conv1,up], axis = 3)
conv = Conv2D(32, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(merge)
conv = Conv2D(32, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(conv)
conv10 = Conv2D(4, (1,1), activation = 'softmax')(conv)
return Model(inputs = inputs, outputs = conv10)
Here is the visual representation of our model:
Build and Plot the Model¶
input_layer = Input((IMG_SIZE, IMG_SIZE, 2))
generator = build_unet(input_layer, 'he_normal', 0.2)
plot_model(generator,
show_shapes = True,
show_dtype=False,
show_layer_names = True,
rankdir = 'TB',
expand_nested = False,
dpi = 70)
Building the Discriminator (cGAN)¶
In a Conditional GAN (cGAN) for segmentation, the Discriminator acts as an anatomical critic. It takes two inputs: the original MRI slice and a segmentation mask. Its job is to classify whether the provided mask is the "Real" ground truth (drawn by a doctor) or a "Fake" generated by our U-Net.
We will use a Convolutional Neural Network (CNN) that compresses these inputs down to a single probability score (1 for Real, 0 for Fake).
from tensorflow.keras.layers import Input, Conv2D, Concatenate, LeakyReLU, Flatten, Dense, Dropout
from tensorflow.keras.models import Model
def build_discriminator(image_shape=(IMG_SIZE, IMG_SIZE, 2), mask_shape=(IMG_SIZE, IMG_SIZE, 4)):
# Inputs: The MRI scan and the Mask (either real or generated)
inp_img = Input(shape=image_shape, name='mri_image')
inp_mask = Input(shape=mask_shape, name='segmentation_mask')
# Concatenate the MRI and the Mask along the channel axis
merged = Concatenate()([inp_img, inp_mask])
# Downsampling layers
d = Conv2D(64, (4,4), strides=(2,2), padding='same')(merged)
d = LeakyReLU(alpha=0.2)(d)
d = Conv2D(128, (4,4), strides=(2,2), padding='same')(d)
d = LeakyReLU(alpha=0.2)(d)
d = Conv2D(256, (4,4), strides=(2,2), padding='same')(d)
d = LeakyReLU(alpha=0.2)(d)
# Flatten and output a single classification value
d = Flatten()(d)
d = Dropout(0.3)(d)
validity = Dense(1, activation='sigmoid')(d)
return Model([inp_img, inp_mask], validity, name='discriminator')
discriminator = build_discriminator()
plot_model(discriminator, show_shapes=True, show_layer_names=True, rankdir='TB', expand_nested=False, dpi=70)
Optimizers and Loss Functions¶
Unlike standard models, a GAN requires two separate optimizers. We use Adam with a lower learning rate and a beta_1 of 0.5, which is standard for stabilizing GAN training.
For the loss:
- Discriminator Loss: Binary Crossentropy (guessing 1 for Real, 0 for Fake).
- Generator Loss: A combination of Adversarial Loss (fooling the discriminator) AND standard Categorical Crossentropy (ensuring the mask still matches the correct pixel classes).
import tensorflow as tf
# Optimizers
generator_optimizer = tf.keras.optimizers.Adam(learning_rate=0.0002, beta_1=0.5)
discriminator_optimizer = tf.keras.optimizers.Adam(learning_rate=0.0002, beta_1=0.5)
# Loss functions
bce_loss = tf.keras.losses.BinaryCrossentropy()
cce_loss = tf.keras.losses.CategoricalCrossentropy()
# Weight for the segmentation loss vs adversarial loss
LAMBDA_SEG = 100 # Forces the generator to prioritize accurate pixel segmentation
Custom Training Step¶
We define a single training step using tf.GradientTape. This allows us to manually calculate the gradients and update the weights for both the Generator and the Discriminator simultaneously. The @tf.function decorator compiles this into a static graph for much faster execution.
@tf.function
def train_step(mri_images, real_masks):
with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
# 1. Generator creates a fake mask
fake_masks = generator(mri_images, training=True)
# 2. Discriminator evaluates both real and fake masks
real_output = discriminator([mri_images, real_masks], training=True)
fake_output = discriminator([mri_images, fake_masks], training=True)
# 3. Calculate Discriminator Loss
disc_loss_real = bce_loss(tf.ones_like(real_output), real_output)
disc_loss_fake = bce_loss(tf.zeros_like(fake_output), fake_output)
total_disc_loss = disc_loss_real + disc_loss_fake
# 4. Calculate Generator Loss
# Adversarial loss: generator wants discriminator to guess 1 (Real)
gen_adv_loss = bce_loss(tf.ones_like(fake_output), fake_output)
# Segmentation loss: generated mask must match the real mask shape/classes
gen_seg_loss = cce_loss(real_masks, fake_masks)
total_gen_loss = gen_adv_loss + (LAMBDA_SEG * gen_seg_loss)
# 5. Calculate gradients
generator_gradients = gen_tape.gradient(total_gen_loss, generator.trainable_variables)
discriminator_gradients = disc_tape.gradient(total_disc_loss, discriminator.trainable_variables)
# 6. Apply gradients
generator_optimizer.apply_gradients(zip(generator_gradients, generator.trainable_variables))
discriminator_optimizer.apply_gradients(zip(discriminator_gradients, discriminator.trainable_variables))
return total_gen_loss, total_disc_loss, gen_seg_loss
Execution: GAN Training Loop¶
This cell executes the custom training loop over our BraTS dataset. It replaces the standard model.fit() and provides clean, batch-by-batch logging to monitor the adversarial network's progress.
import time
EPOCHS = 35
print("Starting cGAN Training...")
train_start_time = time.time()
for epoch in range(EPOCHS):
epoch_start_time = time.time()
print(f"\n--- Epoch {epoch + 1}/{EPOCHS} ---")
gen_loss_list = []
disc_loss_list = []
seg_loss_list = []
# Iterate over the training generator
for batch_idx in range(len(training_generator)):
mri_batch, mask_batch = training_generator[batch_idx]
g_loss, d_loss, seg_loss = train_step(mri_batch, mask_batch)
gen_loss_list.append(g_loss)
disc_loss_list.append(d_loss)
seg_loss_list.append(seg_loss)
# Clean logging
if batch_idx == 0:
print(f" ↳ Batch 1 compiled & completed!")
elif (batch_idx + 1) % 50 == 0:
print(f" ↳ Processing batch {batch_idx + 1}/{len(training_generator)}...")
# Calculate metrics for the epoch
avg_g_loss = tf.reduce_mean(gen_loss_list)
avg_d_loss = tf.reduce_mean(disc_loss_list)
avg_seg_loss = tf.reduce_mean(seg_loss_list)
epoch_duration = time.time() - epoch_start_time
print(f" Epoch {epoch + 1} Finished in {epoch_duration/60:.2f} mins")
print(f" Gen Loss: {avg_g_loss:.4f} | Disc Loss: {avg_d_loss:.4f} | Seg Loss: {avg_seg_loss:.4f}")
# Save the generator weights periodically to prevent data loss
if (epoch + 1) % 5 == 0:
generator.save_weights(f'cgan_generator_epoch_{epoch+1}.weights.h5')
total_duration = time.time() - train_start_time
print(f"\nTraining Complete! Total Time: {total_duration/60:.2f} minutes")
Inference and Visualization¶
Now that the Generator (U-Net) has been trained with adversarial feedback, we evaluate its performance. We will pull a batch from our validation/test generator and visualize the predictions against the original ground truth.
import matplotlib.pyplot as plt
import numpy as np
import random
print("🔍 Hunting through validation data for a visible tumor...")
found_tumor = False
attempts = 0
# Keep searching until we find a batch with at least one tumor
while not found_tumor and attempts < 50:
random_batch_idx = random.randint(0, len(valid_generator) - 1)
val_mri_batch, val_mask_batch = valid_generator[random_batch_idx]
# Check the ground truth for this batch
real_masks_argmax = np.argmax(val_mask_batch, axis=-1)
# Make a list of all slices that are NOT just purple background
tumor_indices = [i for i in range(len(real_masks_argmax)) if np.max(real_masks_argmax[i]) > 0]
if len(tumor_indices) > 0:
found_tumor = True
else:
attempts += 1
if not found_tumor:
print("⚠️ Looked through 50 random batches and couldn't find a tumor! Check your validation data.")
else:
print(f"✅ Found {len(tumor_indices)} tumor slices in batch {random_batch_idx}! Generating images...")
# Generate predictions ONLY for the winning batch
predicted_masks = generator.predict(val_mri_batch, verbose=0)
pred_masks_argmax = np.argmax(predicted_masks, axis=-1)
# RANDOMLY select up to 3 slices from the ones that actually have tumors
num_samples = min(3, len(tumor_indices))
random_tumor_slices = random.sample(tumor_indices, num_samples)
fig, axes = plt.subplots(num_samples, 3, figsize=(12, 4 * num_samples))
# Safety catch if there's only 1 tumor in the batch
if num_samples == 1: axes = [axes]
for i, idx in enumerate(random_tumor_slices):
# Original MRI
axes[i][0].imshow(val_mri_batch[idx, :, :, 0], cmap='gray')
axes[i][0].set_title(f"Original MRI (Slice {idx})")
axes[i][0].axis('off')
# Ground Truth Mask
axes[i][1].imshow(real_masks_argmax[idx], cmap='viridis', vmin=0, vmax=3)
axes[i][1].set_title("Ground Truth Mask")
axes[i][1].axis('off')
# cGAN Predicted Mask
axes[i][2].imshow(pred_masks_argmax[idx], cmap='viridis', vmin=0, vmax=3)
axes[i][2].set_title("cGAN Prediction")
axes[i][2].axis('off')
plt.tight_layout()
plt.show()
Save the Final Model¶
With the adversarial training complete, we secure our final U-Net. We only need to save the Generator because it is the network responsible for actual tumor segmentation. The Discriminator is no longer needed for inference.
generator.save('final_champion_model.keras')
print("done")
import tensorflow as tf
import numpy as np
# Load your ALREADY TRAINED 4x4 cGAN Generator
print("📥 Loading saved 4x4 cGAN model...")
model = tf.keras.models.load_model('/kaggle/input/models/dhruvgupta1402/4x4-enhanced/keras/default/1/final_champion_model.keras', compile=False)
print("📊 Evaluating 4x4 cGAN with enhancement(Strictly on slices containing tumors)...")
total_psnr = 0.0
total_ssim = 0.0
tumor_slice_count = 0
num_batches = len(valid_generator)
for i in range(num_batches):
mri_batch, mask_batch = valid_generator[i]
predictions = model.predict(mri_batch, verbose=0)
real_argmax = np.argmax(mask_batch, axis=-1)
pred_argmax = np.argmax(predictions, axis=-1)
for j in range(len(real_argmax)):
if np.max(real_argmax[j]) > 0: # Only evaluate actual tumors
r_tensor = tf.convert_to_tensor(real_argmax[j], dtype=tf.float32)
p_tensor = tf.convert_to_tensor(pred_argmax[j], dtype=tf.float32)
r_tensor = tf.expand_dims(tf.expand_dims(r_tensor, axis=0), axis=-1)
p_tensor = tf.expand_dims(tf.expand_dims(p_tensor, axis=0), axis=-1)
total_psnr += tf.image.psnr(r_tensor, p_tensor, max_val=3.0)[0].numpy()
total_ssim += tf.image.ssim(r_tensor, p_tensor, max_val=3.0)[0].numpy()
tumor_slice_count += 1
if (i + 1) % 50 == 0:
print(f" ↳ Processed {i + 1}/{num_batches} batches...")
avg_psnr = total_psnr / tumor_slice_count if tumor_slice_count > 0 else 0
avg_ssim = total_ssim / tumor_slice_count if tumor_slice_count > 0 else 0
print("\n🏆 4x4 cGAN Structural Metrics 🏆")
print(f"Average PSNR: {avg_psnr:.4f} dB")
print(f"Average SSIM: {avg_ssim:.4f}")