JeeKay/brain-tumor-segmentation

1import os2 3import h5py4import matplotlib.pyplot as plt5import numpy as np6import streamlit as st7import tensorflow as tf8from PIL import Image9 10from models.unet_multitask import build_unet_multioutput11 12st.set_page_config(page_title="Brain Tumor Segmentation - BraTS20", layout="wide")13 14# --------- Load Model ---------15@st.cache_resource16def load_model():17    model_path = "models/unet_multihead_brats.keras"18    if not os.path.exists(model_path):19        st.error(f"Model weights file not found at {model_path}. Please check the path.")20        return None21    model = build_unet_multioutput(input_shape=(240, 240, 4))22    model.load_weights(model_path)23    return model24 25model = load_model()26if model is None:27    st.stop()28 29# --------- Utility Functions ---------30 31def overlay_mask(flair, wt, tc, et):32    flair_norm = (flair - flair.min()) / (flair.max() - flair.min() + 1e-8)33    if flair_norm.ndim == 2:34        flair_norm = np.expand_dims(flair_norm, axis=-1)35    overlay = np.repeat(flair_norm, 3, axis=-1)36 37    overlay[wt[..., 0].astype(bool)] = [1, 0, 0]  # Red: WT38    overlay[tc[..., 0].astype(bool)] = [0, 1, 0]  # Green: TC39    overlay[et[..., 0].astype(bool)] = [0, 0, 1]  # Blue: ET40 41    return overlay42 43def overlay_errors(flair, gt_masks, pred_masks):44    flair_norm = (flair - flair.min()) / (flair.max() - flair.min() + 1e-8)45    overlay = np.repeat(np.expand_dims(flair_norm, axis=-1), 3, axis=-1)46 47    for mask_name, color in zip(['wt', 'tc', 'et'], [[1, 0, 0], [0, 1, 0], [0, 0, 1]]):48        gt = gt_masks[mask_name][..., 0].astype(bool)49        pred = pred_masks[mask_name][..., 0].astype(bool)50 51        tp = np.logical_and(gt, pred)52        fn = np.logical_and(gt, ~pred)53        fp = np.logical_and(~gt, pred)54 55        overlay[tp] = color            # Correct56        overlay[fn] = [1, 1, 0]        # Yellow for FN57        overlay[fp] = [1, 0, 1]        # Magenta for FP58 59    return overlay60 61def load_h5_slice(file):62    try:63        with h5py.File(file, 'r') as f:64            if 'image' not in f or 'mask' not in f:65                raise ValueError("H5 file must contain 'image' and 'mask' datasets.")66            image = f['image'][()]67            mask = f['mask'][()]68 69        # Normalize image70        image = (image - np.mean(image, axis=(0, 1), keepdims=True)) / \71                (np.std(image, axis=(0, 1), keepdims=True) + 1e-6)72 73        # Handle different mask formats74        if mask.ndim == 3 and mask.shape[-1] == 3:75            ncr = mask[..., 0]76            ed = mask[..., 1]77            et = mask[..., 2]78            wt = ((ncr + ed + et) > 0).astype(np.float32)[..., np.newaxis]79            tc = ((ncr + et) > 0).astype(np.float32)[..., np.newaxis]80            et = (et > 0).astype(np.float32)[..., np.newaxis]81        else:82            mask = mask.astype(np.uint8)83            wt = (mask > 0).astype(np.float32)[..., np.newaxis]84            tc = np.isin(mask, [1, 4]).astype(np.float32)[..., np.newaxis]85            et = (mask == 4).astype(np.float32)[..., np.newaxis]86 87        flair = image[:, :, 3]  # FLAIR channel88 89        return image, flair, wt, tc, et90 91    except Exception as e:92        st.error(f"Error loading H5 file: {e}")93        return None, None, None, None, None94 95def preprocess_image(image):96    target_shape = (240, 240, 4)97    if image.shape != target_shape:98        st.warning(f"Resizing input image from {image.shape} to {target_shape}.")99        image = tf.image.resize(image, target_shape[:2], method=tf.image.ResizeMethod.BILINEAR)100        image = tf.cast(image, tf.float32)101    return image102 103# --------- Default Demo File Handling ---------104 105DEMO_FILE_PATH = "volume_101_slice_63.h5"106 107@st.cache_data108def get_demo_data():109    return load_h5_slice(DEMO_FILE_PATH)110 111# --------- Streamlit UI ---------112 113st.title("🧠 Brain Tumor Segmentation - BraTS2020")114 115st.markdown("### 🧪 Try it out:")116st.markdown("Use the demo below to see how the segmentation looks, or upload your own `.h5` file.")117 118# Upload section119uploaded_file = st.file_uploader("Upload a BraTS slice (.h5 file)", type=['h5'])120 121# Use uploaded file or fallback to demo122if uploaded_file is not None:123    file_to_use = uploaded_file124    source = "Uploaded File"125else:126    file_to_use = DEMO_FILE_PATH127    source = "Demo File"128 129# Load data130image, flair, wt, tc, et = load_h5_slice(file_to_use)131if image is None:132    st.error("Could not load selected file.")133else:134    image = preprocess_image(image)135    pred = model.predict(image[np.newaxis, ...])136 137    # Threshold predictions138    pred_masks = {139        'wt': (pred[0][0] > 0.5).astype(np.float32),140        'tc': (pred[1][0] > 0.5).astype(np.float32),141        'et': (pred[2][0] > 0.5).astype(np.float32)142    }143 144    gt_masks = {145        'wt': wt,146        'tc': tc,147        'et': et148    }149 150    # Overlays151    gt_overlay = overlay_mask(flair, wt, tc, et)152    pred_overlay = overlay_mask(flair, pred_masks['wt'], pred_masks['tc'], pred_masks['et'])153    error_overlay = overlay_errors(flair, gt_masks, pred_masks)154 155    # Display width156    IMAGE_DISPLAY_WIDTH = 300157 158    col1, col2, col3, col4 = st.columns(4)159 160    with col1:161        st.markdown("**Original FLAIR Image**")162        if flair.ndim == 3 and flair.shape[-1] == 1:163            flair = flair.squeeze(axis=-1)164        flair_norm = (flair - flair.min()) / (flair.max() - flair.min() + 1e-8)165        flair_display = np.repeat(flair_norm[..., np.newaxis], 3, axis=-1)166        st.image(flair_display, clamp=True, width=IMAGE_DISPLAY_WIDTH)167 168    with col2:169        st.markdown(f"**Ground Truth Overlay ({source})**")170        st.image(gt_overlay, clamp=True, width=IMAGE_DISPLAY_WIDTH)171 172    with col3:173        st.markdown("**Predicted Mask Overlay**")174        st.image(pred_overlay, clamp=True, width=IMAGE_DISPLAY_WIDTH)175 176    with col4:177        st.markdown("**Error Mask Overlay**\nFN=Yellow | FP=Magenta")178        st.image(error_overlay, clamp=True, width=IMAGE_DISPLAY_WIDTH)