Update app.py

app.py

@@ -1,30 +1,47 @@
 import gradio as gr
-import subprocess
-import os
-import shutil
-from huggingface_hub import hf_hub_download
 import torch
-import spaces  # Import spaces for GPU decoration
 import numpy as np
-import pandas as pd
+import os
+import nibabel as nib
 import torchio
 import torch.nn as nn
+import subprocess
+from scipy.ndimage.measurements import center_of_mass
 from huggingface_hub import hf_hub_download
 from monai.transforms import Compose, LoadImaged, Spacingd, CropForegroundd, SpatialPadd, CenterSpatialCropd
-from monai.data import Dataset
+from monai.data import Dataset, DataLoader
 from nnunet_mednext import create_mednext_encoder_v1
 
-# Debugging GPU environment
-if torch.cuda.is_available():
-    print(f"GPU is available: {torch.cuda.get_device_name(0)}")
-else:
-    print("No GPU available. Falling back to CPU.")
-    os.system("nvidia-smi")
+# Model and data directory setup
+MODEL_DIR = "/root/.cache/huggingface/hub"
+DATASET_DIR = os.path.join(MODEL_DIR, "BrainAgeNeXt")
+REPO_ID = "FrancescoLR/BrainAgeNeXt"
+
+# Ensure model directory exists
+os.makedirs(MODEL_DIR, exist_ok=True)
 
-# Device selection
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# 🔹 Function to Download Model Weights from Hugging Face
+def download_model():
+    if not os.path.exists(DATASET_DIR):
+        os.makedirs(DATASET_DIR, exist_ok=True)
+        print("Downloading BrainAgeNeXt model weights...")
+        for i in range(1, 6):
+            hf_hub_download(repo_id=REPO_ID, filename=f"BrainAge_{i}.pth", cache_dir=MODEL_DIR)
+        print("✅ BrainAgeNeXt model downloaded successfully.")
+
+# 🔹 Function to Load Model
+def initialize_model():
+    model_paths = [hf_hub_download(repo_id=REPO_ID, filename=f"BrainAge_{i}.pth", cache_dir=MODEL_DIR) for i in range(1, 6)]
+    
+    models = []
+    for model_path in model_paths:
+        model = MedNeXtEncReg().to(device)
+        model.load_state_dict(torch.load(model_path, map_location=device))
+        model.eval()
+        models.append(model)
+    return models
 
-# Model definition
+# 🔹 Define Model
 class MedNeXtEncReg(nn.Module):
     def __init__(self):
         super(MedNeXtEncReg, self).__init__()
@@ -46,21 +63,7 @@ class MedNeXtEncReg(nn.Module):
         age_estimate = self.regression_fc(x)
         return age_estimate.squeeze()
 
-# Download the model from Hugging Face Hub
-def initialize_model():
-    model_paths = [
-        hf_hub_download(repo_id="FrancescoLR/BrainAgeNeXt", filename=f"BrainAge_{i}.pth") for i in range(1, 6)
-    ]
-    
-    models = []
-    for model_path in model_paths:
-        model = MedNeXtEncReg().to(device)
-        model.load_state_dict(torch.load(model_path, map_location=device))
-        model.eval()
-        models.append(model)
-    return models
-
-# Define preprocessing transforms
+# 🔹 Preprocessing Pipeline
 def prepare_transforms():
     return Compose([
         LoadImaged(keys=["image"], ensure_channel_first=True),
@@ -71,63 +74,86 @@ def prepare_transforms():
         torchio.transforms.ZNormalization(masking_method=lambda x: x > 0, keys=["image"])
     ])
 
-# Process uploaded MRI scan
-def preprocess_mri(mri_path):
+# 🔹 Process MRI File
+def preprocess_mri(nifti_path):
     transforms = prepare_transforms()
-    data_dict = {"image": mri_path}
+    data_dict = {"image": nifti_path}
     dataset = Dataset([data_dict], transform=transforms)
     dataloader = DataLoader(dataset, batch_size=1, num_workers=0)
     return next(iter(dataloader))["image"].to(device)
 
-# Predict brain age
-def predict_brain_age(mri_path, actual_age, sex):
-    if not os.path.exists(mri_path):
+# 🔹 Run Brain Age Prediction (Decorated for GPU Execution)
+@spaces.GPU(duration=90)
+def predict_brain_age(nifti_file, actual_age, sex):
+    if not os.path.exists(nifti_file.name):
         return "Error: MRI file not found"
 
-    # Load the model
+    # Load Model
     models = initialize_model()
 
     # Preprocess MRI
-    image = preprocess_mri(mri_path)
+    image = preprocess_mri(nifti_file.name)
 
-    # Run predictions
+    # Run Predictions
     predictions = []
     with torch.no_grad():
         for model in models:
             pred = model(image)
             predictions.append(pred.cpu().numpy())
 
-    # Compute median brain age prediction
+    # Compute Median Brain Age Prediction
     predicted_brain_age = np.median(np.stack(predictions))
 
-    # Apply correction based on actual age
+    # Apply Correction Based on Actual Age
     predicted_brain_age_corrected = (
         predicted_brain_age + (actual_age * 0.062) - 2.96 if actual_age > 18 else predicted_brain_age
     )
 
     brain_age_difference = predicted_brain_age_corrected - actual_age
 
-    # Output results
+    # Output Results
     return f"Predicted Brain Age: {predicted_brain_age_corrected:.2f} years", \
            f"Brain Age Difference (BAD): {brain_age_difference:.2f} years"
 
-# Gradio UI
-demo = gr.Interface(
-    fn=predict_brain_age,
-    inputs=[
-        gr.File(label="Upload MRI (NIfTI .nii.gz)"),
-        gr.Number(label="Enter Age"),
-        gr.Radio(["Male", "Female"], label="Select Sex")
-    ],
-    outputs=[
-        gr.Textbox(label="Predicted Brain Age"),
-        gr.Textbox(label="Brain Age Difference (BAD)")
-    ],
-    title="Brain Age Prediction with MedNeXt",
-    description="Upload an MRI scan (.nii.gz), enter your age and sex, and get a brain age prediction.",
-    theme="default"
-)
-
-# Launch the Gradio app
+# 🔹 Gradio Interface Setup
+with gr.Blocks() as demo:
+    gr.Markdown("""
+    # 🧠 Brain Age Prediction with MedNeXt
+    Upload an MRI scan (.nii.gz), enter your age and sex, and get a brain age prediction.
+    """)
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            mri_input = gr.File(label="Upload MRI (NIfTI .nii.gz)")
+            age_input = gr.Number(label="Enter Age", value=30)
+            sex_input = gr.Radio(["Male", "Female"], label="Select Sex")
+            submit_button = gr.Button("Predict")
+
+        with gr.Column(scale=2):
+            brain_age_output = gr.Textbox(label="Predicted Brain Age")
+            bad_output = gr.Textbox(label="Brain Age Difference (BAD)")
+
+    submit_button.click(
+        fn=predict_brain_age,
+        inputs=[mri_input, age_input, sex_input],
+        outputs=[brain_age_output, bad_output]
+    )
+
+    gr.Markdown("""
+    **Disclaimer:** This is a research tool and is not intended for clinical use.
+    """)
+# 🔹 Debugging GPU Environment
+if torch.cuda.is_available():
+    print(f"GPU available: {torch.cuda.get_device_name(0)}")
+    device = torch.device("cuda")
+else:
+    print("No GPU detected. Falling back to CPU.")
+    os.system("nvidia-smi")
+    device = torch.device("cpu")
+    
+# 🔹 Download Model Weights
+download_model()
+
+# 🔹 Run Gradio App
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch(share=True)