JeeKay/brain-tumor-segmentation
0
1#!/usr/bin/env python2# coding: utf-83 4# In[1]:5 6 7import os8import numpy as np9import tensorflow as tf10from tensorflow.keras.callbacks import ModelCheckpoint, ReduceLROnPlateau, EarlyStopping11from dataprep_multitask import get_train_val_datasets12from utils.losses import dice_loss, dice_coefficient, focal_tversky_loss # custom metrics13from models.unet_multitask import build_unet_multioutput14from glob import glob15from visualize import visualize_batch16import h5py17 18os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' # 3=errors only19os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'20os.environ['TF_GPU_ALLOCATOR'] = 'cuda_malloc_async'21 22# Set seeds for reproducibility23tf.random.set_seed(42)24np.random.seed(42)25 26# Paths27DATA_DIR = "BraTS20/BraTS2020_training_data/content/data"28MODEL_SAVE_PATH = "models/unet_brats.keras"29 30# Hyperparameters31BATCH_SIZE = 832EPOCHS = 3033 34# Load datasets35train_dataset, val_dataset = get_train_val_datasets(DATA_DIR, batch_size=BATCH_SIZE)36 37 38# ------------------- Model Setup -------------------39model = build_unet_multioutput(input_shape=(240, 240, 4))40model.compile(41 optimizer=tf.keras.optimizers.Adam(learning_rate=1e-4),42 loss={43 'wt_head': focal_tversky_loss,44 'tc_head': focal_tversky_loss,45 'et_head': focal_tversky_loss,46 },47 loss_weights={48 'wt_head': 1.0,49 'tc_head': 1.0,50 'et_head': 2.0,51 },52 metrics={53 'wt_head': dice_coefficient,54 'tc_head': dice_coefficient,55 'et_head': dice_coefficient,56 }57)58 59# ------------------- Callbacks -------------------60callbacks = [61 ModelCheckpoint("models/unet_multihead_brats.keras", save_best_only=True, monitor='val_et_head_dice_coefficient', mode='max'),62 ReduceLROnPlateau(monitor='val_et_head_dice_coefficient', factor=0.5, patience=4, min_lr=1e-6, verbose=1),63 EarlyStopping(monitor='val_et_head_dice_coefficient', mode='max', patience=8, restore_best_weights=True)64]65 66# ------------------- Training -------------------67history = model.fit(68 train_dataset,69 validation_data=val_dataset,70 steps_per_epoch=len(train_dataset) // BATCH_SIZE,71 validation_steps=len(val_dataset) // BATCH_SIZE,72 epochs=EPOCHS,73 callbacks=callbacks74)75 76 77# In[4]:78 79 80import matplotlib.pyplot as plt81 82print("\nGenerating predictions and visualizing...\n")83 84# Get one batch from validation set85for x_batch, y_batch in val_dataset.take(1):86 preds = model.predict(x_batch)87 pred_masks = {88 'wt_head': (preds[0] > 0.5).astype(np.float32),89 'tc_head': (preds[1] > 0.5).astype(np.float32),90 'et_head': (preds[2] > 0.5).astype(np.float32),91 }92 93 def overlay_prediction(img, pred_masks):94 flair = img[:, :, 3]95 flair_norm = (flair - flair.min()) / (flair.max() - flair.min() + 1e-8)96 overlay = np.stack([flair_norm]*3, axis=-1)97 overlay[pred_masks['wt_head'][..., 0] == 1] = [1, 0, 0] # Red98 overlay[pred_masks['tc_head'][..., 0] == 1] = [0, 1, 0] # Green99 overlay[pred_masks['et_head'][..., 0] == 1] = [0, 0, 1] # Blue100 return overlay101 102 for i in range(min(3, x_batch.shape[0])):103 image = x_batch[i].numpy()104 flair = image[:, :, 3]105 pred_overlay = overlay_prediction(image, {106 'wt_head': pred_masks['wt_head'][i],107 'tc_head': pred_masks['tc_head'][i],108 'et_head': pred_masks['et_head'][i],109 })110 111 plt.figure(figsize=(12, 4))112 plt.subplot(1, 2, 1)113 plt.imshow(flair, cmap='gray')114 plt.title("FLAIR")115 plt.axis('off')116 117 plt.subplot(1, 2, 2)118 plt.imshow(pred_overlay)119 plt.title("Predicted Mask Overlay")120 plt.axis('off')121 plt.show()122 123 break124 125 126# In[5]:127 128 129import matplotlib.pyplot as plt130 131def overlay_mask(flair, wt, tc, et):132 """133 Build RGB overlay on FLAIR using WT (red), TC (green), ET (blue) masks.134 """135 flair_norm = (flair - flair.min()) / (flair.max() - flair.min() + 1e-8)136 overlay = np.stack([flair_norm]*3, axis=-1)137 overlay[wt[..., 0] == 1] = [1, 0, 0] # Red138 overlay[tc[..., 0] == 1] = [0, 1, 0] # Green139 overlay[et[..., 0] == 1] = [0, 0, 1] # Blue140 return overlay141 142 143# Predict one batch144for x_batch, y_batch in val_dataset.take(1):145 preds = model.predict(x_batch)146 147 for i in range(3): # Show first 3 samples148 img = x_batch[i].numpy()149 flair = img[:, :, 3]150 151 # Predicted masks (thresholded)152 wt_pred = (preds[0][i] > 0.5).astype(np.float32)153 tc_pred = (preds[1][i] > 0.5).astype(np.float32)154 et_pred = (preds[2][i] > 0.5).astype(np.float32)155 156 # Ground truth masks157 wt_true = y_batch['wt_head'][i].numpy()158 tc_true = y_batch['tc_head'][i].numpy()159 et_true = y_batch['et_head'][i].numpy()160 161 # Overlays162 gt_overlay = overlay_mask(flair, wt_true, tc_true, et_true)163 pred_overlay = overlay_mask(flair, wt_pred, tc_pred, et_pred)164 165 # Plot side-by-side166 plt.figure(figsize=(10, 5))167 plt.subplot(1, 2, 1)168 plt.imshow(gt_overlay)169 plt.title("Ground Truth Overlay")170 plt.axis('off')171 172 plt.subplot(1, 2, 2)173 plt.imshow(pred_overlay)174 plt.title("Predicted Overlay")175 plt.axis('off')176 177 plt.tight_layout()178 plt.show()179 180 break181 182 183# In[6]:184 185 186import matplotlib.pyplot as plt187import numpy as np188 189def overlay_mask(flair, wt, tc, et):190 """191 Build RGB overlay on FLAIR using WT (red), TC (green), ET (blue).192 """193 flair_norm = (flair - flair.min()) / (flair.max() - flair.min() + 1e-8)194 overlay = np.stack([flair_norm]*3, axis=-1)195 overlay[wt[..., 0] == 1] = [1, 0, 0] # Red: WT196 overlay[tc[..., 0] == 1] = [0, 1, 0] # Green: TC197 overlay[et[..., 0] == 1] = [0, 0, 1] # Blue: ET198 return overlay199 200def overlay_errors(flair, gt, pred):201 """202 Compare prediction vs ground truth:203 - TP (correct) regions keep their color204 - FP (predicted but not GT): Magenta205 - FN (GT but not predicted): Yellow206 """207 flair_norm = (flair - flair.min()) / (flair.max() - flair.min() + 1e-8)208 overlay = np.stack([flair_norm]*3, axis=-1)209 210 for mask_name, color, idx in zip(['wt_head', 'tc_head', 'et_head'],211 [[1, 0, 0], [0, 1, 0], [0, 0, 1]],212 range(3)):213 gt_mask = gt[mask_name][..., 0]214 pred_mask = pred[mask_name][..., 0]215 216 tp = np.logical_and(gt_mask == 1, pred_mask == 1)217 fn = np.logical_and(gt_mask == 1, pred_mask == 0) # Missed218 fp = np.logical_and(gt_mask == 0, pred_mask == 1) # Extra219 220 overlay[tp] = color # Correct221 overlay[fn] = [1, 1, 0] # Yellow for FN222 overlay[fp] = [1, 0, 1] # Magenta for FP223 224 return overlay225 226# Run on one batch227for x_batch, y_batch in val_dataset.take(1):228 preds = model.predict(x_batch)229 230 for i in range(3):231 img = x_batch[i].numpy()232 flair = img[:, :, 3]233 234 # Threshold predictions235 pred_masks = {236 'wt_head': (preds[0][i] > 0.5).astype(np.float32),237 'tc_head': (preds[1][i] > 0.5).astype(np.float32),238 'et_head': (preds[2][i] > 0.5).astype(np.float32)239 }240 241 gt_masks = {242 'wt_head': y_batch['wt_head'][i].numpy(),243 'tc_head': y_batch['tc_head'][i].numpy(),244 'et_head': y_batch['et_head'][i].numpy()245 }246 247 # Overlays248 gt_overlay = overlay_mask(flair, gt_masks['wt_head'], gt_masks['tc_head'], gt_masks['et_head'])249 pred_overlay = overlay_mask(flair, pred_masks['wt_head'], pred_masks['tc_head'], pred_masks['et_head'])250 error_overlay = overlay_errors(flair, gt_masks, pred_masks)251 252 # Plot all 3253 plt.figure(figsize=(18, 5))254 plt.subplot(1, 3, 1)255 plt.imshow(gt_overlay)256 plt.title("Ground Truth Overlay")257 plt.axis('off')258 259 plt.subplot(1, 3, 2)260 plt.imshow(pred_overlay)261 plt.title("Predicted Overlay")262 plt.axis('off')263 264 plt.subplot(1, 3, 3)265 plt.imshow(error_overlay)266 plt.title("Error Overlay\nFN=Yellow | FP=Magenta")267 plt.axis('off')268 269 plt.tight_layout()270 plt.show()271 272 break273 274 275# In[ ]:276 277 278 279 280