import os
import torch
import nibabel as nib
import gradio as gr
import tempfile
import yaml
import traceback # For detailed error printing
import zipfile
import dicom2nifti
import shutil
import subprocess # To run unzip command
import SimpleITK as sitk
import itk
import numpy as np
from scipy.signal import medfilt
import skimage.filters
import cv2 # For Gaussian Blur
import io # For saving plots to memory
import base64 # For encoding plots
import uuid # For unique IDs
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg') # Use non-interactive backend
# --- Custom CSS to hide buttons in logo columns ---
# REVERTING CSS CHANGES AS THEY HID THE LOGOS
# custom_css = """
# .logo-column button {
# /* display: none !important; */ /* This seemed to hide the whole component */
# visibility: hidden !important; /* Try making it invisible instead */
# }
# """
# --- Potential Import Issues (Check Paths in Docker/Local) ---
try:
# Assumes HD_BET is now at /app/BrainIAC/HD_BET or adjacent in src
from HD_BET.run import run_hd_bet
from monai.visualize.gradient_based import GuidedBackpropSmoothGrad
except ImportError as e:
print(f"Warning: Could not import HD_BET or MONAI visualize: {e}. Saliency/Preprocessing might fail.")
run_hd_bet = None
GuidedBackpropSmoothGrad = None
# Import necessary components from your existing modules
from model import Backbone, SingleScanModel, Classifier
from monai.transforms import Resized, ScaleIntensityd
# --- Constants ---
APP_DIR = os.path.dirname(__file__)
TEMPLATE_DIR = os.path.join(APP_DIR, "golden_image", "mni_templates")
PARAMS_RIGID_PATH = os.path.join(APP_DIR, "golden_image", "mni_templates", "Parameters_Rigid.txt")
DEFAULT_TEMPLATE_PATH = os.path.join(TEMPLATE_DIR, "nihpd_asym_13.0-18.5_t1w.nii")
HD_BET_CONFIG_PATH = os.path.join(APP_DIR, "HD_BET", "config.py") # May need adjustment based on actual HD_BET location
HD_BET_MODEL_DIR = os.path.join(APP_DIR, "hdbet_model") # Path to copied models
# --- Configuration Loading ---
def load_config():
config_path = os.path.join(APP_DIR, 'config.yml')
try:
with open(config_path, 'r') as file:
config = yaml.safe_load(file)
if 'data' not in config: config['data'] = {}
if 'image_size' not in config['data']: config['data']['image_size'] = [128, 128, 128]
except FileNotFoundError:
print(f"Warning: Configuration file not found at {config_path}. Using defaults.")
config = {
'gpu': {'device': 'cpu'},
'infer': {'checkpoints': 'checkpoints/mci_model.pt'}, # Updated for new MCI model filename
'data': {'image_size': [128, 128, 128]}
}
return config
config = load_config()
DEFAULT_IMAGE_SIZE = (128, 128, 128)
image_size_cfg = config.get('data', {}).get('image_size', DEFAULT_IMAGE_SIZE)
if not isinstance(image_size_cfg, (list, tuple)) or len(image_size_cfg) != 3:
print(f"Warning: Invalid image_size in config ({image_size_cfg}). Using default {DEFAULT_IMAGE_SIZE}.")
image_size = DEFAULT_IMAGE_SIZE
else:
image_size = tuple(image_size_cfg)
# --- Model Loading ---
def load_model(cfg):
device = torch.device(cfg.get('gpu', {}).get('device', 'cpu'))
backbone = Backbone()
classifier = Classifier(d_model=2048, num_classes=1) # Binary classification for MCI
model = SingleScanModel(backbone, classifier) # Using BP model for MCI classification
relative_path = cfg.get('infer', {}).get('checkpoints', 'checkpoints/mci_model.pt')
checkpoint_path_abs = os.path.join(APP_DIR, relative_path)
try:
print(f"Loading MCI classification model from: {checkpoint_path_abs}")
checkpoint = torch.load(checkpoint_path_abs, map_location=device, weights_only=False)
state_dict = checkpoint.get('model_state_dict', checkpoint)
model.load_state_dict(state_dict, strict=False)
model.to(device)
model.eval()
print(f"MCI classification model loaded successfully onto {device}")
return model, device
except FileNotFoundError:
print(f"Error: Checkpoint file not found at {checkpoint_path_abs}")
return None, device
except Exception as e:
print(f"Error loading model checkpoint: {e}")
traceback.print_exc()
return None, device
model, device = load_model(config)
# --- Preprocessing Functions (Copied/Adapted from app.py) ---
def bias_field_correction(img_array):
print(" Running N4 Bias Field Correction...")
image = sitk.GetImageFromArray(img_array)
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)
print(" N4 Correction finished.")
return sitk.GetArrayFromImage(corrected_image)
def denoise(volume, kernel_size=3):
print(f" Applying median filter denoising (kernel={kernel_size})...")
return medfilt(volume, kernel_size)
def rescale_intensity(volume, percentils=[0.5, 99.5], bins_num=256):
print(" Rescaling intensity...")
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:
print(" Otsu failed, skipping background mask.")
obj_volume = volume_float.flatten()
if obj_volume.size == 0:
print(" Warning: No foreground voxels found. Scaling full volume.")
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])
denominator = max_value - min_value
if denominator < 1e-6: denominator = 1e-6
output_volume = np.copy(volume_float)
if bins_num == 0:
output_volume = (volume_float - min_value) / denominator
output_volume = np.clip(output_volume, 0.0, 1.0)
else:
output_volume = np.round((volume_float - min_value) / denominator * (bins_num - 1))
output_volume = np.clip(output_volume, 0, bins_num - 1)
return output_volume.astype(np.float32)
def equalize_hist(volume, bins_num=256):
print(" Performing histogram equalization...")
mask = volume > 1e-6
obj_volume = volume[mask]
if obj_volume.size == 0:
print(" Warning: No non-zero voxels. Skipping equalization.")
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 enhance(img_array, run_bias_correction=True, kernel_size=3, percentils=[0.5, 99.5], bins_num=256, run_equalize_hist=True):
print("Starting enhancement pipeline...")
volume = img_array.astype(np.float32)
try:
if run_bias_correction: volume = bias_field_correction(volume)
volume = denoise(volume, kernel_size)
volume = rescale_intensity(volume, percentils, bins_num)
if run_equalize_hist: volume = equalize_hist(volume, bins_num)
print("Enhancement pipeline finished.")
return volume
except Exception as e:
print(f"Error during enhancement: {e}")
traceback.print_exc()
raise RuntimeError(f"Failed enhancing image: {e}")
def register_image(input_nifti_path, output_nifti_path):
print(f"Registering {input_nifti_path} to {DEFAULT_TEMPLATE_PATH}")
if not all(os.path.exists(p) for p in [PARAMS_RIGID_PATH, DEFAULT_TEMPLATE_PATH]):
raise FileNotFoundError("Elastix parameter or template file not found.")
fixed_image = itk.imread(DEFAULT_TEMPLATE_PATH, itk.F)
moving_image = itk.imread(input_nifti_path, itk.F)
parameter_object = itk.ParameterObject.New()
parameter_object.AddParameterFile(PARAMS_RIGID_PATH)
result_image, _ = itk.elastix_registration_method(fixed_image, moving_image, parameter_object=parameter_object, log_to_console=False)
itk.imwrite(result_image, output_nifti_path)
print(f"Registration output saved to {output_nifti_path}")
def run_enhance_on_file(input_nifti_path, output_nifti_path):
print(f"Running full enhancement on {input_nifti_path}")
img_sitk = sitk.ReadImage(input_nifti_path)
img_array = sitk.GetArrayFromImage(img_sitk)
enhanced_array = enhance(img_array, run_bias_correction=True)
enhanced_img_sitk = sitk.GetImageFromArray(enhanced_array)
enhanced_img_sitk.CopyInformation(img_sitk)
sitk.WriteImage(enhanced_img_sitk, output_nifti_path)
print(f"Enhanced image saved to {output_nifti_path}")
def run_skull_stripping(input_nifti_path, output_dir):
print(f"Running HD-BET skull stripping on {input_nifti_path}")
if run_hd_bet is None: raise RuntimeError("HD-BET module not imported.")
if not os.path.exists(HD_BET_CONFIG_PATH): raise FileNotFoundError(f"HD-BET config not found at {HD_BET_CONFIG_PATH}")
if not os.path.isdir(HD_BET_MODEL_DIR): raise FileNotFoundError(f"HD-BET models not found at {HD_BET_MODEL_DIR}")
base_name = os.path.basename(input_nifti_path).replace(".nii.gz", "").replace(".nii", "")
output_file_path = os.path.join(output_dir, f"{base_name}_bet.nii.gz")
output_mask_path = os.path.join(output_dir, f"{base_name}_bet_mask.nii.gz")
os.makedirs(output_dir, exist_ok=True)
try:
run_hd_bet(input_nifti_path, output_file_path, mode="fast", device='cpu', config_file=HD_BET_CONFIG_PATH, postprocess=False, do_tta=False, keep_mask=True, overwrite=True)
finally:
pass
if not os.path.exists(output_file_path): raise RuntimeError("HD-BET did not produce output file.")
print(f"Skull stripping output saved to {output_file_path}")
return output_file_path, output_mask_path
# --- MONAI Transforms ---
resize_transform = Resized(keys=["image"], spatial_size=image_size)
scale_transform = ScaleIntensityd(keys=["image"], minv=0.0, maxv=1.0)
def preprocess_nifti_for_model(nifti_path):
print(f"Preprocessing NIfTI for model: {nifti_path}")
scan_data = nib.load(nifti_path).get_fdata()
scan_tensor = torch.tensor(scan_data, dtype=torch.float32).unsqueeze(0) # Add C dim
sample = {"image": scan_tensor}
sample_resized = resize_transform(sample)
sample_scaled = scale_transform(sample_resized)
input_tensor = sample_scaled["image"].unsqueeze(0).to(device) # Add B dim
if input_tensor.dim() != 5: raise ValueError(f"Preprocessing resulted in incorrect shape: {input_tensor.shape}")
print(f" Final shape for model: {input_tensor.shape}")
return input_tensor
# --- Saliency Generation ---
def generate_saliency(model_to_use, input_tensor_5d):
if GuidedBackpropSmoothGrad is None: raise ImportError("MONAI visualize components not imported.")
if model_to_use is None: raise ValueError("Model not loaded.")
print("Generating saliency map...")
input_tensor_5d.requires_grad_(True)
visualizer = GuidedBackpropSmoothGrad(model=model_to_use.backbone.to(device), stdev_spread=0.15, n_samples=10, magnitude=True)
try:
with torch.enable_grad():
saliency_map_5d = visualizer(input_tensor_5d.to(device))
input_3d = input_tensor_5d.squeeze().cpu().detach().numpy()
saliency_3d = saliency_map_5d.squeeze().cpu().detach().numpy()
print("Saliency map generated.")
return input_3d, saliency_3d
except Exception as e:
print(f"Error during saliency map generation: {e}")
traceback.print_exc()
return None, None
finally:
input_tensor_5d.requires_grad_(False)
# --- Plotting Function (Returns NumPy arrays for Gradio) ---
def create_slice_plots(mri_data_3d, saliency_data_3d, slice_index):
print(f" Generating plots for slice index: {slice_index}")
if mri_data_3d is None or saliency_data_3d is None: return None, None, None
# Change from the third dimension (axis 2, sagittal) to the first dimension (axis 0, axial)
if not (0 <= slice_index < mri_data_3d.shape[0]):
print(f" Error: Slice index {slice_index} out of bounds (0-{mri_data_3d.shape[0]-1}).")
return None, None, None
# Function to save plot to NumPy array
def save_plot_to_numpy(fig):
with io.BytesIO() as buf:
fig.savefig(buf, format='png', bbox_inches='tight', pad_inches=0, dpi=75) # Adjust DPI as needed
plt.close(fig)
buf.seek(0)
img_arr = plt.imread(buf, format='png')
# Return RGBA array, can be simplified if only grayscale needed for input
return (img_arr * 255).astype(np.uint8)
try:
# Extract axial slice instead of sagittal
mri_slice = mri_data_3d[slice_index, :, :]
saliency_slice_orig = saliency_data_3d[slice_index, :, :]
# Normalize MRI Slice (using volume stats)
p1_vol, p99_vol = np.percentile(mri_data_3d, (1, 99))
mri_norm_denom = max(p99_vol - p1_vol, 1e-6)
mri_slice_norm = np.clip((mri_slice - p1_vol) / mri_norm_denom, 0, 1)
# Process Saliency Slice
saliency_slice = np.copy(saliency_slice_orig)
saliency_slice[saliency_slice < 0] = 0
saliency_slice_blurred = cv2.GaussianBlur(saliency_slice, (15, 15), 0)
s_max_vol = max(np.max(saliency_data_3d[saliency_data_3d >= 0]), 1e-6)
saliency_slice_norm = saliency_slice_blurred / s_max_vol
saliency_slice_thresholded = np.where(saliency_slice_norm > 0.0, saliency_slice_norm, 0) # Threshold slightly > 0
# Plot 1: Input Slice
fig1, ax1 = plt.subplots(figsize=(6, 6))
ax1.imshow(mri_slice_norm, cmap='gray', interpolation='none', origin='lower')
ax1.axis('off')
input_plot_np = save_plot_to_numpy(fig1)
# Plot 2: Saliency Heatmap
fig2, ax2 = plt.subplots(figsize=(6, 6))
ax2.imshow(saliency_slice_thresholded, cmap='magma', interpolation='none', origin='lower', vmin=0) # Set vmin
ax2.axis('off')
heatmap_plot_np = save_plot_to_numpy(fig2)
# Plot 3: Overlay
fig3, ax3 = plt.subplots(figsize=(6, 6))
ax3.imshow(mri_slice_norm, cmap='gray', interpolation='none', origin='lower')
if np.max(saliency_slice_thresholded) > 0:
ax3.contour(saliency_slice_thresholded, cmap='magma', origin='lower', linewidths=1.0, levels=np.linspace(saliency_slice_thresholded.min(), saliency_slice_thresholded.max(), 5)) # Adjust levels
ax3.axis('off')
overlay_plot_np = save_plot_to_numpy(fig3)
print(f" Generated numpy plots successfully for slice {slice_index}.")
return input_plot_np, heatmap_plot_np, overlay_plot_np
except Exception as e:
print(f"Error generating numpy plots for slice {slice_index}: {e}")
traceback.print_exc()
return None, None, None
# Add this function after the create_slice_plots function
def create_probability_gauge(probability):
"""
Creates a gauge visualization for the MCI probability.
Args:
probability (float): A value between 0 and 1 representing the MCI probability.
Returns:
numpy.ndarray: A numpy array containing the gauge visualization image.
"""
# Create a figure with a polar projection
fig = plt.figure(figsize=(8, 4))
ax = fig.add_subplot(111, polar=True)
# Set the min and max angles for the gauge (in radians)
# -pi/2 to pi/2 creates a half-circle (180 degrees)
theta_min = -np.pi/2
theta_max = np.pi/2
# Calculate the angle for the needle based on probability (0 to 1)
needle_angle = theta_min + probability * (theta_max - theta_min)
# Create a color gradient for the gauge background
cmap = plt.cm.RdYlGn_r # Red-Yellow-Green colormap (reversed)
# Draw the gauge background
theta = np.linspace(theta_min, theta_max, 100)
radii = np.ones_like(theta)
# Create color array for the gauge segments
norm = plt.Normalize(0, 1)
colors = cmap(np.linspace(0, 1, len(theta)))
# Draw colored bars for the gauge
width = (theta_max - theta_min) / len(theta)
bars = ax.bar(theta, radii, width=width, bottom=0.0, alpha=0.8, linewidth=0)
# Set the color for each bar
for bar, color in zip(bars, colors):
bar.set_facecolor(color)
# Add the needle
needle_length = 0.9
ax.annotate('', xy=(needle_angle, needle_length), xytext=(needle_angle, 0),
arrowprops=dict(arrowstyle='wedge', color='black', lw=2))
# Add boundary markers and labels - move them lower by adjusting y position (1.1 -> 1.3)
ax.text(theta_min, 2.2, 'Healthy Control', ha='left', va='center', fontsize=12)
ax.text(theta_max, 1.3, 'MCI', ha='right', va='center', fontsize=12)
# Add the probability text below the gauge
prob_text = f"Probability: {probability:.2f}"
fig.text(0.5, 0.15, prob_text, ha='center', va='center', fontsize=14, fontweight='bold')
# Set the limits and remove unnecessary elements
ax.set_ylim(0, 1.4) # Increased the upper limit to accommodate the lower labels
ax.set_theta_zero_location('N') # 0 at the top
ax.set_theta_direction(-1) # clockwise
ax.set_thetagrids([]) # Remove angle labels
ax.grid(False) # Remove grid
ax.set_rgrids([]) # Remove radial labels
ax.spines['polar'].set_visible(False) # Remove the outer circle
# Convert figure to numpy array
with io.BytesIO() as buf:
fig.savefig(buf, format='png', bbox_inches='tight', pad_inches=0.1, dpi=100)
plt.close(fig)
buf.seek(0)
img_arr = plt.imread(buf, format='png')
return (img_arr * 255).astype(np.uint8)
# --- Gradio Processing Function ---
def process_scan(file_type, uploaded_file, run_preprocess, generate_saliency_flag):
if model is None:
raise gr.Error("Model is not loaded. Cannot perform prediction.")
if uploaded_file is None:
raise gr.Error("No file uploaded.")
temp_dir = tempfile.mkdtemp()
print(f"Created temp directory: {temp_dir}")
nifti_for_preprocessing_path = None
error_message = None
prediction_text = "Processing..."
# Initialize outputs to None or placeholder images/values
input_slice_img, heatmap_slice_img, overlay_slice_img = None, None, None
probability_gauge = None
saliency_state = {"input_path": None, "saliency_path": None, "num_slices": 0}
slider_update = gr.Slider(value=0, minimum=0, maximum=1, visible=False) # Initially hidden, use max=1 to avoid log(0) error
try:
# --- Handle Upload and DICOM Conversion ---
file_path = uploaded_file.name # Get path from Gradio file object
filename = os.path.basename(file_path)
print(f"Processing '{filename}' (type: {file_type})")
if file_type == 'NIfTI':
# Check if the filename ends with either .nii or .nii.gz
if not (filename.lower().endswith('.nii.gz') or filename.lower().endswith('.nii')):
raise gr.Error("Invalid NIfTI file. Please upload .nii or .nii.gz")
# Define the destination path (always .nii.gz for consistency)
dest_path = os.path.join(temp_dir, "uploaded_scan.nii.gz")
nifti_for_preprocessing_path = dest_path
# Check if the uploaded file is uncompressed .nii
if filename.lower().endswith('.nii') and not filename.lower().endswith('.nii.gz'):
print(f"Detected uncompressed .nii file: {filename}. Compressing to {dest_path}")
try:
# Load the uncompressed .nii file
img = nib.load(file_path)
# Save it as a compressed .nii.gz file
nib.save(img, dest_path)
print(f"Successfully compressed and saved to: {dest_path}")
except Exception as e:
raise gr.Error(f"Failed to load or compress .nii file: {e}")
else:
# If it's already .nii.gz, just copy it
print(f"Copying compressed NIfTI {filename} to: {dest_path}")
shutil.copy(file_path, dest_path)
# nifti_for_preprocessing_path is already set to dest_path
# print(f"NIfTI path for preprocessing: {nifti_for_preprocessing_path}") # Redundant logging
elif file_type == 'DICOM (zip)':
if not filename.endswith('.zip'):
raise gr.Error("Invalid DICOM file. Please upload a .zip archive.")
uploaded_zip_path = os.path.join(temp_dir, "dicom_files.zip")
shutil.copy(file_path, uploaded_zip_path)
print(f"Copied DICOM zip to: {uploaded_zip_path}")
dicom_input_dir = os.path.join(temp_dir, "dicom_input")
nifti_output_dir = os.path.join(temp_dir, "nifti_output")
os.makedirs(dicom_input_dir, exist_ok=True)
os.makedirs(nifti_output_dir, exist_ok=True)
try:
shutil.unpack_archive(uploaded_zip_path, dicom_input_dir)
print("Unzip successful.")
except Exception as e:
raise gr.Error(f"Error unzipping DICOM file: {e}")
try:
dicom2nifti.convert_directory(dicom_input_dir, nifti_output_dir, compression=True, reorient=True)
nifti_files = [f for f in os.listdir(nifti_output_dir) if f.endswith('.nii.gz')]
if not nifti_files: raise RuntimeError("dicom2nifti did not produce a .nii.gz file.")
nifti_for_preprocessing_path = os.path.join(nifti_output_dir, nifti_files[0])
print(f"DICOM conversion successful. NIfTI: {nifti_for_preprocessing_path}")
except Exception as e:
raise gr.Error(f"Error converting DICOM to NIfTI: {e}")
else:
raise gr.Error("Invalid file type selected.")
if not nifti_for_preprocessing_path or not os.path.exists(nifti_for_preprocessing_path):
raise gr.Error("Could not find the NIfTI file after initial processing.")
# --- Optional Preprocessing ---
nifti_to_predict_path = nifti_for_preprocessing_path
if run_preprocess:
print("--- Running Optional Preprocessing Pipeline ---")
try:
registered_path = os.path.join(temp_dir, "registered.nii.gz")
register_image(nifti_for_preprocessing_path, registered_path)
enhanced_path = os.path.join(temp_dir, "enhanced.nii.gz")
run_enhance_on_file(registered_path, enhanced_path)
skullstrip_output_dir = os.path.join(temp_dir, "skullstripped")
skullstripped_path, _ = run_skull_stripping(enhanced_path, skullstrip_output_dir)
nifti_to_predict_path = skullstripped_path
print("--- Optional Preprocessing Pipeline Complete ---")
except Exception as e:
raise gr.Error(f"Error during preprocessing: {e}")
else:
print("--- Skipping Optional Preprocessing Pipeline ---")
# --- Prediction (Changed for MCI Classification) ---
input_tensor_5d = preprocess_nifti_for_model(nifti_to_predict_path)
print("Performing MCI classification prediction...")
with torch.no_grad():
try:
output = model(input_tensor_5d)
# Convert output to probability
if isinstance(output, torch.Tensor):
logit = output.item()
else:
logit = output
# Apply sigmoid to get probability
probability = torch.sigmoid(torch.tensor(logit)).item()
predicted_class = 1 if probability > 0.5 else 0
class_label = "MCI" if predicted_class == 1 else "Healthy Control"
# Create the probability gauge visualization
probability_gauge = create_probability_gauge(probability)
# Format prediction text for classification
prediction_text = f"Prediction: {class_label} "
except Exception as pred_error:
print(f"Error during prediction: {pred_error}")
traceback.print_exc()
raise gr.Error(f"Failed to make prediction: {pred_error}")
print(prediction_text)
# --- Saliency Map Generation ---
if generate_saliency_flag:
print("--- Generating Saliency Data ---")
try:
input_3d, saliency_3d = generate_saliency(model, input_tensor_5d)
if input_3d is not None and saliency_3d is not None:
num_slices = input_3d.shape[0] # Using axial slices now (first dimension)
center_slice_index = num_slices // 2
# Save numpy arrays to the temp dir for the slider callback
unique_id = str(uuid.uuid4())
input_array_path = os.path.join(temp_dir, f"{unique_id}_input.npy")
saliency_array_path = os.path.join(temp_dir, f"{unique_id}_saliency.npy")
np.save(input_array_path, input_3d)
np.save(saliency_array_path, saliency_3d)
print(f"Saved input array: {input_array_path}")
print(f"Saved saliency array: {saliency_array_path}")
# Generate initial plots for the center slice
input_slice_img, heatmap_slice_img, overlay_slice_img = create_slice_plots(input_3d, saliency_3d, center_slice_index)
# Update state for the slider callback
saliency_state = {
"input_path": input_array_path,
"saliency_path": saliency_array_path,
"num_slices": num_slices
}
# Update and show the slider
slider_update = gr.Slider(value=center_slice_index, minimum=0, maximum=num_slices - 1, step=1, label="Select Slice", visible=True)
print("--- Saliency Generated and Initial Plot Created ---")
else:
error_message = "Saliency map generation failed."
print(f"Warning: {error_message}")
except ImportError as e:
error_message = f"Cannot generate saliency: {e}"
print(f"Warning: {error_message}")
except Exception as e:
error_message = f"Error during saliency processing: {e}"
traceback.print_exc()
print(f"Warning: {error_message}")
except Exception as e:
print(f"Error in process_scan: {e}")
traceback.print_exc()
# Use gr.Warning for non-fatal errors shown to user
if error_message: # Prepend specific error if available
gr.Warning(f"{error_message}. General error: {e}")
else:
gr.Warning(f"An error occurred: {e}")
# Return default/error states for outputs
return "Error during processing", None, None, None, None, gr.Slider(visible=False), {"input_path": None, "saliency_path": None, "num_slices": 0}
finally:
# Optional: Schedule cleanup of the temp_dir if files aren't needed long-term
# Be cautious if files ARE needed by slider state. Gradio might handle this?
# shutil.rmtree(temp_dir, ignore_errors=True)
# print(f"Cleaned up temp directory: {temp_dir}") # <--- Defer cleanup
pass
# Return results including the probability gauge
return prediction_text, input_slice_img, heatmap_slice_img, overlay_slice_img, probability_gauge, slider_update, saliency_state
# --- Gradio Slider Update Function ---
def update_slice_viewer(slice_index, current_state):
input_path = current_state.get("input_path")
saliency_path = current_state.get("saliency_path")
if not input_path or not saliency_path or not os.path.exists(input_path) or not os.path.exists(saliency_path):
print(f"Warning: Cannot update slice viewer. Missing or invalid numpy array paths in state: {current_state}")
# Return None or placeholder images to indicate error
return None, None, None
try:
input_3d = np.load(input_path)
saliency_3d = np.load(saliency_path)
num_slices = input_3d.shape[0] # Using axial slices (first dimension)
# Ensure slice_index is valid (Gradio slider should handle bounds, but double-check)
slice_index = int(slice_index)
if not (0 <= slice_index < num_slices):
print(f"Warning: Invalid slice index {slice_index} received by update function.")
return None, None, None # Or return previous plots?
# Generate new plots for the selected slice
input_slice_img, heatmap_slice_img, overlay_slice_img = create_slice_plots(input_3d, saliency_3d, slice_index)
return input_slice_img, heatmap_slice_img, overlay_slice_img
except Exception as e:
print(f"Error updating slice viewer for index {slice_index}: {e}")
traceback.print_exc()
# Return None or indicate error
return None, None, None
# --- Build Gradio Interface ---
with gr.Blocks(css="""
#header-row {
min-height: 150px;
align-items: center;
}
.logo-img img {
height: 150px;
object-fit: contain;
}
.probability-gauge {
display: flex;
justify-content: center;
margin-top: 1rem;
}
""") as demo:
# Header Row with Logos and Title
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: MCI Classification"
"
"
)
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: MCI Classification
**Model Description**
A 3D ResNet50 model trained to predict Mild Cognitive Impairment (MCI) from T1-weighted MRI scans.
**Training Dataset**
- **Subjects**: Trained on T1-weighted MRI scans from subjects with MCI and healthy controls
- **Imaging Modality**: T1-weighted MRI
- **Preprocessing**: Registration to MNI, N4 bias correction, histogram equalization, skull stripping
**Input**
- Format: NIfTI or zipped DICOM
- Required sequence: T1w (3D)
**Output**
- Binary classification: MCI or Healthy Control
- Probability score for MCI
**Intended Use**
- Research use only!
**NOTE**
- Not validated on T2, FLAIR, DWI or other sequences
- Not validated on pathological cases beyond MCI
- 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!
""")
# Use gr.State to store paths to numpy arrays for the slider callback
saliency_state = gr.State({"input_path": None, "saliency_path": None, "num_slices": 0})
# Main Content Row (Controls Left, Output Right)
with gr.Row():
with gr.Column(scale=1):
with gr.Group():
gr.Markdown("### Controls")
file_type = gr.Radio(["NIfTI", "DICOM (zip)"], label="Select Input File Type", value="NIfTI")
scan_file = gr.File(label="Upload Scan File")
run_preprocess = gr.Checkbox(label="Run Preprocessing Pipeline ", value=False)
generate_saliency_checkbox = gr.Checkbox(label="Generate Saliency Maps ", value=True)
submit_btn = gr.Button("Classify MCI", variant="primary")
with gr.Column(scale=3):
with gr.Group():
gr.Markdown("### Classification Result")
prediction_output = gr.Label(label="Classification Result")
# Add the probability gauge visualization
gr.Markdown("MCI Probability
")
probability_gauge = gr.Image(label="Probability Gauge", type="numpy", show_label=False, elem_classes=["probability-gauge"])
with gr.Group():
gr.Markdown("### Saliency Map 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 Slice
")
input_slice_img = gr.Image(label="Input Slice", type="numpy", show_label=False)
with gr.Column():
gr.Markdown("Saliency Heatmap
")
heatmap_slice_img = gr.Image(label="Saliency Heatmap", type="numpy", show_label=False)
with gr.Column():
gr.Markdown("Overlay
")
overlay_slice_img = gr.Image(label="Overlay", type="numpy", show_label=False)
# --- Wire Components ---
submit_btn.click(
fn=process_scan,
inputs=[file_type, scan_file, run_preprocess, generate_saliency_checkbox],
outputs=[prediction_output, input_slice_img, heatmap_slice_img, overlay_slice_img, probability_gauge, slice_slider, saliency_state]
)
slice_slider.change(
fn=update_slice_viewer,
inputs=[slice_slider, saliency_state],
outputs=[input_slice_img, heatmap_slice_img, overlay_slice_img]
)
# --- Launch the App ---
if __name__ == "__main__":
if model is None:
print("ERROR: Model failed to load. Gradio app cannot start.")
else:
print("Launching Gradio Interface...")
demo.launch(server_name="0.0.0.0", server_port=7860, debug=False, share=False)