import os
import yaml
import torch
import nibabel as nib
import numpy as np
import gradio as gr
from typing import Tuple
import tempfile
import shutil
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg') # Use non-interactive backend
import cv2 # For Gaussian Blur
import io # For saving plots to memory
import base64 # For encoding plots
import uuid # For unique IDs
import traceback # For detailed error printing
import SimpleITK as sitk
import itk
from scipy.signal import medfilt
import skimage.filters
from monai.transforms import Compose, LoadImaged, EnsureChannelFirstd, Resized, NormalizeIntensityd, ToTensord, EnsureTyped
from monai.inferers import sliding_window_inference
from model import ViTUNETRSegmentationModel
# Optional HD-BET import (packaged locally like in MCI app)
try:
from HD_BET.run import run_hd_bet
from HD_BET.hd_bet import hd_bet
except Exception as e:
print(f"Warning: HD_BET not available: {e}")
run_hd_bet = None
hd_bet = None
APP_DIR = os.path.dirname(__file__)
TEMPLATE_DIR = os.path.join(APP_DIR, "golden_image", "mni_templates")
PARAMS_RIGID_PATH = os.path.join(TEMPLATE_DIR, "Parameters_Rigid.txt")
DEFAULT_TEMPLATE_PATH = os.path.join(TEMPLATE_DIR, "temp_head.nii.gz")
FLAIR_TEMPLATE_PATH = os.path.join(TEMPLATE_DIR, "nihpd_asym_04.5-18.5_t2w.nii.gz")
HD_BET_CONFIG_PATH = os.path.join(APP_DIR, "HD_BET", "config.py")
HD_BET_MODEL_DIR = os.path.join(APP_DIR, "hdbet_model")
def load_config() -> dict:
cfg_path = os.path.join(APP_DIR, "config.yml")
if os.path.exists(cfg_path):
with open(cfg_path, "r") as f:
return yaml.safe_load(f)
# Defaults
return {
"gpu": {"device": "cpu"},
"infer": {
"checkpoints": "./checkpoints/idh_model.ckpt",
"simclr_checkpoint": None,
"threshold": 0.5,
"image_size": [96, 96, 96],
},
}
def build_model(cfg: dict):
device = torch.device(cfg.get("gpu", {}).get("device", "cpu"))
infer_cfg = cfg.get("infer", {})
model_cfg = cfg.get("model", {})
simclr_path = None#os.path.join(APP_DIR, infer_cfg.get("simclr_checkpoint", ""))
ckpt_path = os.path.join(APP_DIR, infer_cfg.get("checkpoints", ""))
model = ViTUNETRSegmentationModel(
simclr_ckpt_path=None,
img_size=tuple(model_cfg.get("img_size", [96, 96, 96])),
in_channels=model_cfg.get("in_channels", 1),
out_channels=model_cfg.get("out_channels", 1)
)
# Load finetuned checkpoint (Lightning or plain state_dict)
if os.path.exists(ckpt_path):
checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=False)
if "state_dict" in checkpoint:
state_dict = checkpoint["state_dict"]
new_state_dict = {}
for key, value in state_dict.items():
if key.startswith("model."):
new_state_dict[key[len("model."):]] = value
else:
new_state_dict[key] = value
else:
new_state_dict = checkpoint
model.load_state_dict(new_state_dict, strict=False)
else:
print(f"Warning: Segmentation checkpoint not found at {ckpt_path}. Model will use backbone-only weights.")
model.to(device)
model.eval()
return model, device
# ---------------- Preprocessing (Registration + Enhancement + Skull Stripping) ----------------
def bias_field_correction(img_array: np.ndarray) -> np.ndarray:
image = sitk.GetImageFromArray(img_array.astype(np.float32))
if image.GetPixelID() != sitk.sitkFloat32:
image = sitk.Cast(image, sitk.sitkFloat32)
maskImage = sitk.OtsuThreshold(image, 0, 1, 200)
corrector = sitk.N4BiasFieldCorrectionImageFilter()
numberFittingLevels = 4
max_iters = [min(50 * (2 ** i), 200) for i in range(numberFittingLevels)]
corrector.SetMaximumNumberOfIterations(max_iters)
corrected_image = corrector.Execute(image, maskImage)
return sitk.GetArrayFromImage(corrected_image)
def denoise(volume: np.ndarray, kernel_size: int = 3) -> np.ndarray:
return medfilt(volume, kernel_size)
def rescale_intensity(volume: np.ndarray, percentils=[0.5, 99.5], bins_num=256) -> np.ndarray:
volume_float = volume.astype(np.float32)
try:
t = skimage.filters.threshold_otsu(volume_float, nbins=256)
volume_masked = np.copy(volume_float)
volume_masked[volume_masked < t] = 0
obj_volume = volume_masked[np.where(volume_masked > 0)]
except ValueError:
obj_volume = volume_float.flatten()
if obj_volume.size == 0:
obj_volume = volume_float.flatten()
min_value = np.min(obj_volume)
max_value = np.max(obj_volume)
else:
min_value = np.percentile(obj_volume, percentils[0])
max_value = np.percentile(obj_volume, percentils[1])
denom = max_value - min_value
if denom < 1e-6:
denom = 1e-6
if bins_num == 0:
output_volume = (volume_float - min_value) / denom
output_volume = np.clip(output_volume, 0.0, 1.0)
else:
output_volume = np.round((volume_float - min_value) / denom * (bins_num - 1))
output_volume = np.clip(output_volume, 0, bins_num - 1)
return output_volume.astype(np.float32)
def equalize_hist(volume: np.ndarray, bins_num=256) -> np.ndarray:
mask = volume > 1e-6
obj_volume = volume[mask]
if obj_volume.size == 0:
return volume
hist, bins = np.histogram(obj_volume, bins_num, range=(obj_volume.min(), obj_volume.max()))
cdf = hist.cumsum()
cdf_normalized = (bins_num - 1) * cdf / float(cdf[-1])
equalized_obj_volume = np.interp(obj_volume, bins[:-1], cdf_normalized)
equalized_volume = np.copy(volume)
equalized_volume[mask] = equalized_obj_volume
return equalized_volume.astype(np.float32)
def run_enhance_on_file(input_nifti_path: str, output_nifti_path: str):
"""
Simplified enhancement - just copy the file since N4 is now done in registration.
This maintains compatibility with the existing preprocessing pipeline.
"""
print(f"Enhancement step (N4 already applied during registration): {input_nifti_path}")
# Since N4 bias correction is now handled in registration, just copy the file
import shutil
shutil.copy2(input_nifti_path, output_nifti_path)
print(f"Enhancement complete (passthrough): {output_nifti_path}")
def register_image_sitk(input_nifti_path: str, output_nifti_path: str, template_path: str, interp_type='linear'):
"""
MRI registration with SimpleITK matching the provided script approach.
Args:
input_nifti_path: Path to input NIfTI file
output_nifti_path: Path to save registered output
template_path: Path to template image
interp_type: Interpolation type ('linear', 'bspline', 'nearest_neighbor')
"""
print(f"Registering {input_nifti_path} to template {template_path}")
# Read template and moving images
fixed_img = sitk.ReadImage(template_path, sitk.sitkFloat32)
moving_img = sitk.ReadImage(input_nifti_path, sitk.sitkFloat32)
# Apply N4 bias correction to moving image
moving_img = sitk.N4BiasFieldCorrection(moving_img)
# Resample fixed image to 1mm isotropic
old_size = fixed_img.GetSize()
old_spacing = fixed_img.GetSpacing()
new_spacing = (1, 1, 1)
new_size = [
int(round((old_size[0] * old_spacing[0]) / float(new_spacing[0]))),
int(round((old_size[1] * old_spacing[1]) / float(new_spacing[1]))),
int(round((old_size[2] * old_spacing[2]) / float(new_spacing[2])))
]
# Set interpolation type
if interp_type == 'linear':
interp_type = sitk.sitkLinear
elif interp_type == 'bspline':
interp_type = sitk.sitkBSpline
elif interp_type == 'nearest_neighbor':
interp_type = sitk.sitkNearestNeighbor
else:
interp_type = sitk.sitkLinear
# Resample fixed image
resample = sitk.ResampleImageFilter()
resample.SetOutputSpacing(new_spacing)
resample.SetSize(new_size)
resample.SetOutputOrigin(fixed_img.GetOrigin())
resample.SetOutputDirection(fixed_img.GetDirection())
resample.SetInterpolator(interp_type)
resample.SetDefaultPixelValue(fixed_img.GetPixelIDValue())
resample.SetOutputPixelType(sitk.sitkFloat32)
fixed_img = resample.Execute(fixed_img)
# Initialize transform
transform = sitk.CenteredTransformInitializer(
fixed_img,
moving_img,
sitk.Euler3DTransform(),
sitk.CenteredTransformInitializerFilter.GEOMETRY)
# Set up registration method
registration_method = sitk.ImageRegistrationMethod()
registration_method.SetMetricAsMattesMutualInformation(numberOfHistogramBins=50)
registration_method.SetMetricSamplingStrategy(registration_method.RANDOM)
registration_method.SetMetricSamplingPercentage(0.01)
registration_method.SetInterpolator(sitk.sitkLinear)
registration_method.SetOptimizerAsGradientDescent(
learningRate=1.0,
numberOfIterations=100,
convergenceMinimumValue=1e-6,
convergenceWindowSize=10)
registration_method.SetOptimizerScalesFromPhysicalShift()
registration_method.SetShrinkFactorsPerLevel(shrinkFactors=[4, 2, 1])
registration_method.SetSmoothingSigmasPerLevel(smoothingSigmas=[2, 1, 0])
registration_method.SmoothingSigmasAreSpecifiedInPhysicalUnitsOn()
registration_method.SetInitialTransform(transform)
# Execute registration
final_transform = registration_method.Execute(fixed_img, moving_img)
# Apply transform and save registered image
moving_img_resampled = sitk.Resample(
moving_img,
fixed_img,
final_transform,
sitk.sitkLinear,
0.0,
moving_img.GetPixelID())
sitk.WriteImage(moving_img_resampled, output_nifti_path)
print(f"Registration complete. Saved to: {output_nifti_path}")
def register_image(input_nifti_path: str, output_nifti_path: str):
"""Wrapper to maintain compatibility - now uses SimpleITK registration."""
if not os.path.exists(DEFAULT_TEMPLATE_PATH):
raise FileNotFoundError(f"Template file missing: {DEFAULT_TEMPLATE_PATH}")
register_image_sitk(input_nifti_path, output_nifti_path, DEFAULT_TEMPLATE_PATH)
def run_skull_stripping(input_nifti_path: str, output_dir: str):
"""
Brain extraction using HD-BET direct integration matching the script approach.
Args:
input_nifti_path: Path to input NIfTI file
output_dir: Directory to save skull-stripped output
Returns:
tuple: (output_file_path, output_mask_path)
"""
print(f"Running HD-BET skull stripping on {input_nifti_path}")
if hd_bet is None:
raise RuntimeError("HD-BET not available. Please include HD_BET and hdbet_model in src/IDH.")
if not os.path.exists(HD_BET_MODEL_DIR):
raise FileNotFoundError(f"HD-BET models not found at {HD_BET_MODEL_DIR}")
os.makedirs(output_dir, exist_ok=True)
# Get base filename and prepare HD-BET compatible naming
base_name = os.path.basename(input_nifti_path).replace('.nii.gz', '').replace('.nii', '')
# HD-BET expects files with _0000 suffix - create temporary file if needed
temp_input_dir = os.path.join(output_dir, "temp_input")
os.makedirs(temp_input_dir, exist_ok=True)
# Copy input file with _0000 suffix for HD-BET
temp_input_path = os.path.join(temp_input_dir, f"{base_name}_0000.nii.gz")
shutil.copy2(input_nifti_path, temp_input_path)
# Set device
device = "0" if torch.cuda.is_available() else "cpu"
try:
# Also try setting the specific model file path
model_file = os.path.join(HD_BET_MODEL_DIR, '0.model')
if os.path.exists(model_file):
print(f"Local model file exists at: {model_file}")
else:
print(f"Warning: Model file not found at: {model_file}")
# List directory contents for debugging
if os.path.exists(HD_BET_MODEL_DIR):
print(f"Contents of {HD_BET_MODEL_DIR}: {os.listdir(HD_BET_MODEL_DIR)}")
else:
print(f"Directory {HD_BET_MODEL_DIR} does not exist")
# Run HD-BET directly on the temporary directory
print(f"Running hd_bet with input_dir: {temp_input_dir}, output_dir: {output_dir}")
hd_bet(temp_input_dir, output_dir, device=device, mode='fast', tta=0)
# HD-BET outputs files with original naming convention
output_file_path = os.path.join(output_dir, f"{base_name}_0000.nii.gz")
output_mask_path = os.path.join(output_dir, f"{base_name}_0000_mask.nii.gz")
# Rename to expected format for compatibility
final_output_path = os.path.join(output_dir, f"{base_name}_bet.nii.gz")
final_mask_path = os.path.join(output_dir, f"{base_name}_bet_mask.nii.gz")
if os.path.exists(output_file_path):
shutil.move(output_file_path, final_output_path)
if os.path.exists(output_mask_path):
shutil.move(output_mask_path, final_mask_path)
# Clean up temporary directory
shutil.rmtree(temp_input_dir, ignore_errors=True)
if not os.path.exists(final_output_path):
raise RuntimeError(f"HD-BET did not produce output file: {final_output_path}")
print(f"Skull stripping complete. Output saved to: {final_output_path}")
return final_output_path, final_mask_path
except Exception as e:
# Clean up on error
shutil.rmtree(temp_input_dir, ignore_errors=True)
raise RuntimeError(f"HD-BET skull stripping failed: {str(e)}")
# ---------------- Visualization Functions ----------------
def create_segmentation_plots(input_data_3d, seg_mask_3d, slice_index):
"""Create segmentation visualization plots: Input, Mask, and Overlay."""
print(f"Generating segmentation plots for slice index: {slice_index}")
if any(data is None for data in [input_data_3d, seg_mask_3d]):
return None, None, None
# Check bounds - using axis 2 for axial slices
if not (0 <= slice_index < input_data_3d.shape[2]):
print(f"Error: Slice index {slice_index} out of bounds (0-{input_data_3d.shape[2]-1}).")
return None, None, None
def save_plot_to_numpy(fig):
with io.BytesIO() as buf:
fig.savefig(buf, format='png', bbox_inches='tight', pad_inches=0, dpi=75)
plt.close(fig)
buf.seek(0)
img_arr = plt.imread(buf, format='png')
return (img_arr * 255).astype(np.uint8)
try:
# Extract axial slices - using axis 2 (last dimension)
input_slice = input_data_3d[:, :, slice_index]
mask_slice = seg_mask_3d[:, :, slice_index]
# Normalize input slice
def normalize_slice(slice_data, volume_data):
p1, p99 = np.percentile(volume_data, (1, 99))
denom = max(p99 - p1, 1e-6)
return np.clip((slice_data - p1) / denom, 0, 1)
input_slice_norm = normalize_slice(input_slice, input_data_3d)
# Create plots
plots = []
# Input Image
fig1, ax1 = plt.subplots(figsize=(6, 6))
ax1.imshow(input_slice_norm, cmap='gray', interpolation='none', origin='lower')
ax1.axis('off')
ax1.set_title('Input Image', fontsize=14, color='white', pad=10)
plots.append(save_plot_to_numpy(fig1))
# Segmentation Mask
fig2, ax2 = plt.subplots(figsize=(6, 6))
ax2.imshow(mask_slice, cmap='hot', interpolation='none', origin='lower', vmin=0, vmax=1)
ax2.axis('off')
ax2.set_title('Segmentation Mask', fontsize=14, color='white', pad=10)
plots.append(save_plot_to_numpy(fig2))
# Overlay
fig3, ax3 = plt.subplots(figsize=(6, 6))
ax3.imshow(input_slice_norm, cmap='gray', interpolation='none', origin='lower')
# Create mask overlay with transparency - using red colormap
mask_overlay = np.ma.masked_where(mask_slice < 0.5, mask_slice)
ax3.imshow(mask_overlay, cmap='Reds', interpolation='none', origin='lower', alpha=0.7, vmin=0, vmax=1)
ax3.axis('off')
ax3.set_title('Overlay', fontsize=14, color='white', pad=10)
plots.append(save_plot_to_numpy(fig3))
print(f"Generated 3 segmentation plots successfully for axial slice {slice_index}.")
return tuple(plots)
except Exception as e:
print(f"Error generating segmentation plots for slice {slice_index}: {e}")
traceback.print_exc()
return tuple([None] * 3)
# ---------------- Saliency Generation (Legacy - keeping for reference) ----------------
def extract_attention_map(vit_model, image, layer_idx=-1, img_size=(96, 96, 96), patch_size=16):
"""
Extracts the attention map from a Vision Transformer (ViT) model.
This function wraps the attention blocks of the ViT to capture the attention
weights during a forward pass. It then processes these weights to generate
a 3D saliency map corresponding to the model's focus on the input image.
"""
attention_maps = {}
original_attns = {}
# A wrapper class to intercept and store attention weights from a ViT block.
class AttentionWithWeights(torch.nn.Module):
def __init__(self, original_attn_module):
super().__init__()
self.original_attn_module = original_attn_module
self.attn_weights = None
def forward(self, x):
# The original implementation of the attention module may not return
# the attention weights. This wrapper recalculates them to ensure they
# are captured. This is based on the standard ViT attention mechanism.
output = self.original_attn_module(x)
if hasattr(self.original_attn_module, 'qkv'):
qkv = self.original_attn_module.qkv(x)
batch_size, seq_len, _ = x.shape
# Assuming qkv has been fused and has shape (batch_size, seq_len, 3 * num_heads * head_dim)
qkv = qkv.reshape(batch_size, seq_len, 3, self.original_attn_module.num_heads, -1)
qkv = qkv.permute(2, 0, 3, 1, 4)
q, k, v = qkv[0], qkv[1], qkv[2]
attn = (q @ k.transpose(-2, -1)) * self.original_attn_module.scale
self.attn_weights = attn.softmax(dim=-1)
return output
# Store original attention modules and replace with wrappers
for i, block in enumerate(vit_model.blocks):
if hasattr(block, 'attn'):
original_attns[i] = block.attn
block.attn = AttentionWithWeights(block.attn)
try:
# Perform a forward pass to execute the wrapped modules and capture weights
with torch.no_grad():
_ = vit_model(image)
# Collect the captured attention weights from each block
for i, block in enumerate(vit_model.blocks):
if hasattr(block.attn, 'attn_weights') and block.attn.attn_weights is not None:
attention_maps[f"layer_{i}"] = block.attn.attn_weights.detach()
finally:
# Restore original attention modules
for i, original_attn in original_attns.items():
vit_model.blocks[i].attn = original_attn
if not attention_maps:
raise RuntimeError("Could not extract any attention maps. Please check the ViT model structure.")
# Select the attention map from the specified layer
if layer_idx < 0:
layer_idx = len(attention_maps) + layer_idx
layer_name = f"layer_{layer_idx}"
if layer_name not in attention_maps:
raise ValueError(f"Layer {layer_idx} not found. Available layers: {list(attention_maps.keys())}")
layer_attn = attention_maps[layer_name]
# Average attention across all heads
head_attn = layer_attn[0].mean(dim=0)
# Get attention from the [CLS] token to all other image patches
cls_attn = head_attn[0, 1:]
# Reshape the 1D attention vector into a 3D volume
patches_per_dim = img_size[0] // patch_size
total_patches = patches_per_dim ** 3
# Pad or truncate if the number of patches doesn't align
if cls_attn.shape[0] != total_patches:
if cls_attn.shape[0] > total_patches:
cls_attn = cls_attn[:total_patches]
else:
padded = torch.zeros(total_patches, device=cls_attn.device)
padded[:cls_attn.shape[0]] = cls_attn
cls_attn = padded
cls_attn_3d = cls_attn.reshape(patches_per_dim, patches_per_dim, patches_per_dim)
cls_attn_3d = cls_attn_3d.unsqueeze(0).unsqueeze(0) # Add batch and channel dims
# Upsample the attention map to the full image resolution
upsampled_attn = torch.nn.functional.interpolate(
cls_attn_3d,
size=img_size,
mode='trilinear',
align_corners=False
).squeeze()
# Normalize the map to [0, 1] for visualization
upsampled_attn = upsampled_attn.cpu().numpy()
upsampled_attn = (upsampled_attn - upsampled_attn.min()) / (upsampled_attn.max() - upsampled_attn.min())
return upsampled_attn
def generate_saliency_dual(model, input_tensor, layer_idx=-1):
"""
Generate saliency maps for dual-input IDH model.
Args:
model: The complete IDH model
input_tensor: Dual input tensor (batch_size, 2, C, D, H, W)
layer_idx: ViT layer to visualize
Returns:
tuple: (flair_input_3d, t1c_input_3d, flair_saliency_3d)
"""
print("Generating saliency maps for dual input...")
try:
# Extract individual images from dual input
# input_tensor shape: [batch_size, 2, C, D, H, W]
flair_tensor = input_tensor[:, 0] # [batch, C, D, H, W]
t1c_tensor = input_tensor[:, 1] # [batch, C, D, H, W]
# Get the ViT backbone
vit_model = model.backbone.backbone
# Generate attention map only for FLAIR
flair_attn = extract_attention_map(vit_model, flair_tensor, layer_idx)
# Convert input tensors to numpy for visualization
flair_input_3d = flair_tensor.squeeze().cpu().detach().numpy()
t1c_input_3d = t1c_tensor.squeeze().cpu().detach().numpy()
print("Saliency maps generated successfully.")
return flair_input_3d, t1c_input_3d, flair_attn
except Exception as e:
print(f"Error during saliency generation: {e}")
traceback.print_exc()
return None, None, None
# ---------------- Visualization Functions ----------------
def create_slice_plots_dual(flair_data_3d, t1c_data_3d, flair_saliency_3d, slice_index):
"""Create slice plots for simplified dual input visualization: T1c, FLAIR, FLAIR attention."""
print(f"Generating plots for slice index: {slice_index}")
if any(data is None for data in [flair_data_3d, t1c_data_3d, flair_saliency_3d]):
return None, None, None
# Check bounds - using axis 2 for axial slices
if not (0 <= slice_index < flair_data_3d.shape[2]):
print(f"Error: Slice index {slice_index} out of bounds (0-{flair_data_3d.shape[2]-1}).")
return None, None, None
def save_plot_to_numpy(fig):
with io.BytesIO() as buf:
fig.savefig(buf, format='png', bbox_inches='tight', pad_inches=0, dpi=75)
plt.close(fig)
buf.seek(0)
img_arr = plt.imread(buf, format='png')
return (img_arr * 255).astype(np.uint8)
try:
# Extract axial slices - using axis 2 (last dimension)
flair_slice = flair_data_3d[:, :, slice_index]
t1c_slice = t1c_data_3d[:, :, slice_index]
flair_saliency_slice = flair_saliency_3d[:, :, slice_index]
# Normalize input slices
def normalize_slice(slice_data, volume_data):
p1, p99 = np.percentile(volume_data, (1, 99))
denom = max(p99 - p1, 1e-6)
return np.clip((slice_data - p1) / denom, 0, 1)
flair_slice_norm = normalize_slice(flair_slice, flair_data_3d)
t1c_slice_norm = normalize_slice(t1c_slice, t1c_data_3d)
# Process saliency slice
def process_saliency_slice(saliency_slice, saliency_volume):
saliency_slice = np.copy(saliency_slice)
saliency_slice[saliency_slice < 0] = 0
saliency_slice_blurred = cv2.GaussianBlur(saliency_slice, (15, 15), 0)
s_max = max(np.max(saliency_volume[saliency_volume >= 0]), 1e-6)
saliency_slice_norm = saliency_slice_blurred / s_max
return np.where(saliency_slice_norm > 0.0, saliency_slice_norm, 0)
flair_sal_processed = process_saliency_slice(flair_saliency_slice, flair_saliency_3d)
# Create plots
plots = []
# T1c Input
fig1, ax1 = plt.subplots(figsize=(6, 6))
ax1.imshow(t1c_slice_norm, cmap='gray', interpolation='none', origin='lower')
ax1.axis('off')
ax1.set_title('T1c Input', fontsize=14, color='white', pad=10)
plots.append(save_plot_to_numpy(fig1))
# FLAIR Input
fig2, ax2 = plt.subplots(figsize=(6, 6))
ax2.imshow(flair_slice_norm, cmap='gray', interpolation='none', origin='lower')
ax2.axis('off')
ax2.set_title('FLAIR Input', fontsize=14, color='white', pad=10)
plots.append(save_plot_to_numpy(fig2))
# FLAIR Attention
fig3, ax3 = plt.subplots(figsize=(6, 6))
ax3.imshow(flair_sal_processed, cmap='magma', interpolation='none', origin='lower', vmin=0)
ax3.axis('off')
ax3.set_title('FLAIR Attention', fontsize=14, color='white', pad=10)
plots.append(save_plot_to_numpy(fig3))
print(f"Generated 3 plots successfully for axial slice {slice_index}.")
return tuple(plots)
except Exception as e:
print(f"Error generating plots for slice {slice_index}: {e}")
traceback.print_exc()
return tuple([None] * 3)
# ---------------- Inference ----------------
def get_validation_transform(image_size: Tuple[int, int, int]):
return Compose([
LoadImaged(keys=["image"]),
EnsureChannelFirstd(keys=["image"]),
Resized(keys=["image"], spatial_size=tuple(image_size), mode="trilinear"),
NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
EnsureTyped(keys=["image"]),
ToTensord(keys=["image"]),
])
def preprocess_nifti(image_path: str, image_size: Tuple[int, int, int], device: torch.device) -> torch.Tensor:
transform = get_validation_transform(image_size)
sample = {"image": image_path}
sample = transform(sample)
image = sample["image"].unsqueeze(0).to(device) # Add batch dimension
return image
def save_nifti_for_download(data_array: np.ndarray, reference_path: str, output_path: str, affine=None):
"""
Save a numpy array as NIfTI file for download, preserving spatial information from reference.
Args:
data_array: 3D numpy array to save
reference_path: Path to reference NIfTI file for header info
output_path: Path where to save the output file
affine: Optional affine matrix, if None will use reference
"""
try:
# Load reference image to get header and affine
ref_img = nib.load(reference_path)
if affine is None:
affine = ref_img.affine
# Create new NIfTI image
new_img = nib.Nifti1Image(data_array, affine, ref_img.header)
# Save the file
nib.save(new_img, output_path)
print(f"Saved NIfTI file: {output_path}")
return output_path
except Exception as e:
print(f"Error saving NIfTI file: {e}")
return None
def predict_segmentation(input_file, threshold: float, do_preprocess: bool, cfg: dict, model, device):
try:
if input_file is None:
return {"error": "Please upload a NIfTI file (.nii.gz)."}, None, None, None, gr.Slider(visible=False), {"input_paths": None, "mask_paths": None, "num_slices": 0}, None, None
input_path = input_file.name if hasattr(input_file, 'name') else input_file
if not (input_path.endswith(".nii") or input_path.endswith(".nii.gz")):
return {"error": "Input must be a NIfTI file (.nii or .nii.gz)."}, None, None, None, gr.Slider(visible=False), {"input_paths": None, "mask_paths": None, "num_slices": 0}, None, None
work_dir = tempfile.mkdtemp()
final_input_path = input_path
try:
# Optional preprocessing pipeline for FLAIR
if do_preprocess:
# Registration to FLAIR template
reg_path = os.path.join(work_dir, "flair_registered.nii.gz")
register_image_sitk(input_path, reg_path, FLAIR_TEMPLATE_PATH)
# Enhancement
enh_path = os.path.join(work_dir, "flair_enhanced.nii.gz")
run_enhance_on_file(reg_path, enh_path)
# Skull stripping
skullstrip_dir = os.path.join(work_dir, "skullstripped")
bet_path, _ = run_skull_stripping(enh_path, skullstrip_dir)
final_input_path = bet_path
# Inference
image_size = cfg.get("infer", {}).get("image_size", [96, 96, 96])
training_cfg = cfg.get("training", {})
input_tensor = preprocess_nifti(final_input_path, image_size, device)
with torch.no_grad():
# Use sliding window inference for better results
seg_logits = sliding_window_inference(
inputs=input_tensor,
roi_size=tuple(image_size),
sw_batch_size=training_cfg.get("sw_batch_size", 2),
predictor=model,
overlap=0.5
)
# Apply sigmoid and threshold to get binary mask
seg_prob = torch.sigmoid(seg_logits)
seg_mask = (seg_prob > threshold).float()
# Convert to numpy for visualization
input_3d = input_tensor.squeeze().cpu().detach().numpy()
seg_prob_3d = seg_prob.squeeze().cpu().detach().numpy()
seg_mask_3d = seg_mask.squeeze().cpu().detach().numpy()
# Calculate statistics
total_voxels = np.prod(seg_mask_3d.shape)
segmented_voxels = int(np.sum(seg_mask_3d))
segmentation_percentage = (segmented_voxels / total_voxels) * 100
prediction_result = {
"segmented_voxels": segmented_voxels,
"total_voxels": total_voxels,
"segmentation_percentage": float(segmentation_percentage),
"threshold": float(threshold),
"preprocessing": bool(do_preprocess),
"max_probability": float(np.max(seg_prob_3d)),
"mean_probability": float(np.mean(seg_prob_3d))
}
# Initialize visualization outputs
input_img = seg_mask_img = overlay_img = None
slider_update = gr.Slider(visible=False)
viz_state = {"input_paths": None, "mask_paths": None, "num_slices": 0}
# Initialize download files
download_preprocessed = None
download_mask = None
# Generate visualizations
print("--- Generating Visualizations ---")
try:
num_slices = input_3d.shape[2] # Use axis 2 for axial slices
center_slice_index = num_slices // 2
# Save numpy arrays for slider callback
unique_id = str(uuid.uuid4())
temp_paths = []
for name, data in [("input", input_3d), ("seg_prob", seg_prob_3d), ("seg_mask", seg_mask_3d)]:
path = os.path.join(work_dir, f"{unique_id}_{name}.npy")
np.save(path, data)
temp_paths.append(path)
# Generate initial plots for center slice
plots = create_segmentation_plots(input_3d, seg_mask_3d, center_slice_index)
if plots and all(p is not None for p in plots):
input_img, seg_mask_img, overlay_img = plots
# Update state and slider
viz_state = {
"input_paths": [temp_paths[0]], # [input]
"mask_paths": temp_paths[1:], # [seg_prob, seg_mask]
"num_slices": num_slices
}
slider_update = gr.Slider(value=center_slice_index, minimum=0, maximum=num_slices-1, step=1, label="Select Slice", visible=True)
print("--- Visualization Generation Complete ---")
except Exception as e:
print(f"Error during visualization generation: {e}")
traceback.print_exc()
# Generate downloadable files
print("--- Generating Download Files ---")
try:
# Create download filenames
base_name = os.path.splitext(os.path.basename(input_path))[0]
if base_name.endswith('.nii'):
base_name = os.path.splitext(base_name)[0]
# Save preprocessed image (the actual array that was fed to the model)
preprocessed_download_path = os.path.join(work_dir, f"{base_name}_preprocessed.nii.gz")
# Save the preprocessed numpy array that was actually used for inference
saved_preprocessed_path = save_nifti_for_download(
input_3d, # This is the preprocessed array that was visualized
input_path, # Use original input as reference for header/affine
preprocessed_download_path
)
if saved_preprocessed_path:
download_preprocessed = gr.File(value=saved_preprocessed_path, visible=True, label="Download Preprocessed Image")
# Save segmentation mask
mask_download_path = os.path.join(work_dir, f"{base_name}_segmentation_mask.nii.gz")
saved_mask_path = save_nifti_for_download(
seg_mask_3d,
final_input_path,
mask_download_path
)
if saved_mask_path:
download_mask = gr.File(value=saved_mask_path, visible=True, label="Download Segmentation Mask")
print("--- Download Files Generated ---")
except Exception as e:
print(f"Error generating download files: {e}")
traceback.print_exc()
return (prediction_result, input_img, seg_mask_img, overlay_img, slider_update, viz_state, download_preprocessed, download_mask)
except Exception as e:
shutil.rmtree(work_dir, ignore_errors=True)
return {"error": f"Processing failed: {str(e)}"}, None, None, None, gr.Slider(visible=False), {"input_paths": None, "mask_paths": None, "num_slices": 0}, None, None
except Exception as e:
return {"error": str(e)}, None, None, None, gr.Slider(visible=False), {"input_paths": None, "mask_paths": None, "num_slices": 0}, None, None
def update_slice_viewer_segmentation(slice_index, current_state):
"""Update slice viewer for segmentation visualization."""
input_paths = current_state.get("input_paths", [])
mask_paths = current_state.get("mask_paths", [])
if not input_paths or not mask_paths or len(input_paths) != 1 or len(mask_paths) != 2:
print(f"Warning: Invalid state for slice viewer update: {current_state}")
return None, None, None
try:
# Load numpy arrays
input_3d = np.load(input_paths[0])
seg_mask_3d = np.load(mask_paths[1]) # Use the binary mask, not probabilities
# Validate slice index
slice_index = int(slice_index)
if not (0 <= slice_index < input_3d.shape[2]): # Use axis 2 for axial slices
print(f"Warning: Invalid slice index {slice_index}")
return None, None, None
# Generate new plots
plots = create_segmentation_plots(input_3d, seg_mask_3d, slice_index)
return plots if plots else tuple([None] * 3)
except Exception as e:
print(f"Error updating slice viewer for index {slice_index}: {e}")
traceback.print_exc()
return tuple([None] * 3)
def build_interface():
cfg = load_config()
model, device = build_model(cfg)
default_threshold = float(cfg.get("infer", {}).get("threshold", 0.5))
with gr.Blocks(title="BrainIAC: Glioma Segmentation", css="""
#header-row {
min-height: 150px;
align-items: center;
}
.logo-img img {
height: 150px;
object-fit: contain;
}
""") as demo:
# --- Header with Logos ---
with gr.Row(elem_id="header-row"):
with gr.Column(scale=1):
gr.Image(os.path.join(APP_DIR, "static/images/kannlab.png"),
show_label=False, interactive=False,
show_download_button=False,
container=False,
elem_classes=["logo-img"])
with gr.Column(scale=3):
gr.Markdown(
""
"BrainIAC: Glioma Segmentation"
"
"
)
with gr.Column(scale=1):
gr.Image(os.path.join(APP_DIR, "static/images/brainiac.jpeg"),
show_label=False, interactive=False,
show_download_button=False,
container=False,
elem_classes=["logo-img"])
# --- Add model description section ---
with gr.Accordion("ℹ️ Model Details and Usage Guide", open=False):
gr.Markdown("""
### 🧠 BrainIAC: Glioma Segmentation
**Model Description**
A Vision Transformer UNETR (ViT-UNETR) model with BrainIAC as pre-trained backbone designed for glioma segmentation from MRI scans.
**Training Dataset**
- **Subjects**: Trained on MRI scans from glioma patients
- **Imaging Modalities**: Single modality MRI FLAIR, and binary mask
- **Preprocessing**: N4 bias correction, MNI registration, and skull stripping (HD-BET)
**Input**
- Format: NIfTI (.nii or .nii.gz)
- Single MRI FLAIR sequence
- Image size: Automatically resized to 96×96×96 voxels
**Output**
- Binary segmentation mask highlighting glioma regions
- Segmentation statistics (volume, percentage)
- Probability maps and overlay visualization
**Intended Use**
- Research use only!
**NOTE**
- Single modality input FLAIR
- Not validated on other MRI sequences
- Not validated for other brain pathologies beyond gliomas
- Upload PHI data at own risk!
- The model is hosted on a cloud-based CPU instance
- The data is not stored, shared or collected for any purpose!
**Visualization**
The interface shows three views for each slice:
- **Input Image**: The preprocessed MRI scan
- **Segmentation Mask**: The predicted binary mask
- **Overlay**: The mask overlaid on the input image
""")
# Use gr.State to store paths to numpy arrays for the slider callback
viz_state = gr.State({"input_paths": None, "mask_paths": None, "num_slices": 0})
# Main Content
gr.Markdown("**Upload MRI NIfTI volume** — Optional preprocessing performs registration to MNI, enhancement, and skull stripping.")
with gr.Row():
with gr.Column(scale=1):
with gr.Group():
gr.Markdown("### Controls")
input_file = gr.File(label="MRI Image (.nii or .nii.gz)")
preprocess_checkbox = gr.Checkbox(value=False, label="Preprocess NIfTI (debiasing + registration + skull stripping)")
threshold_input = gr.Slider(minimum=0.0, maximum=1.0, value=default_threshold, step=0.01, label="Segmentation Threshold")
predict_btn = gr.Button("Generate Segmentation", variant="primary")
with gr.Column(scale=2):
with gr.Group():
gr.Markdown("### Segmentation Result")
output_json = gr.JSON(label="Results")
# Download section
with gr.Row():
download_preprocessed_btn = gr.File(label="Download Preprocessed Image", visible=False)
download_mask_btn = gr.File(label="Download Segmentation Mask", visible=False)
# Segmentation visualization section
with gr.Group():
gr.Markdown("### Segmentation Viewer (Axial Slice)")
slice_slider = gr.Slider(label="Select Slice", minimum=0, maximum=0, step=1, value=0, visible=False)
with gr.Row():
with gr.Column():
gr.Markdown("Input Image
")
input_img = gr.Image(label="Input Image", type="numpy", show_label=False)
with gr.Column():
gr.Markdown("Segmentation Mask
")
seg_mask_img = gr.Image(label="Segmentation Mask", type="numpy", show_label=False)
with gr.Column():
gr.Markdown("Overlay
")
overlay_img = gr.Image(label="Overlay", type="numpy", show_label=False)
# Wire components
predict_btn.click(
fn=lambda f, prep, thr: predict_segmentation(f, thr, prep, cfg, model, device),
inputs=[input_file, preprocess_checkbox, threshold_input],
outputs=[output_json, input_img, seg_mask_img, overlay_img, slice_slider, viz_state, download_preprocessed_btn, download_mask_btn],
)
slice_slider.change(
fn=update_slice_viewer_segmentation,
inputs=[slice_slider, viz_state],
outputs=[input_img, seg_mask_img, overlay_img]
)
return demo
if __name__ == "__main__":
iface = build_interface()
iface.launch(server_name="0.0.0.0", server_port=7860)