Create app.py

app.py

@@ -0,0 +1,122 @@
+import gradio as gr
+import torch
+import numpy as np
+import pandas as pd
+import os
+import torchio
+import torch.nn as nn
+from huggingface_hub import hf_hub_download
+from monai.transforms import Compose, LoadImaged, Spacingd, CropForegroundd, SpatialPadd, CenterSpatialCropd
+from monai.data import Dataset
+from nnunet_mednext import create_mednext_encoder_v1
+
+# Device selection
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Model definition
+class MedNeXtEncReg(nn.Module):
+    def __init__(self):
+        super(MedNeXtEncReg, self).__init__()
+        self.mednextv1 = create_mednext_encoder_v1(
+            num_input_channels=1, num_classes=1, model_id='B', kernel_size=3, deep_supervision=True
+        )
+        self.global_avg_pool = nn.AdaptiveAvgPool3d((1, 1, 1))
+        self.regression_fc = nn.Sequential(
+            nn.Linear(512, 64),
+            nn.ReLU(),
+            nn.Dropout(0.0),
+            nn.Linear(64, 1)
+        )
+
+    def forward(self, x):
+        x = self.mednextv1(x)
+        x = self.global_avg_pool(x)
+        x = torch.flatten(x, start_dim=1)
+        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
+def prepare_transforms():
+    return Compose([
+        LoadImaged(keys=["image"], ensure_channel_first=True),
+        Spacingd(keys=["image"], pixdim=(1.0, 1.0, 1.0)),
+        CropForegroundd(keys=["image"], allow_smaller=True, source_key="image"),
+        SpatialPadd(keys=["image"], spatial_size=(160, 192, 160)),
+        CenterSpatialCropd(keys=["image"], roi_size=(160, 192, 160)),
+        torchio.transforms.ZNormalization(masking_method=lambda x: x > 0, keys=["image"])
+    ])
+
+# Process uploaded MRI scan
+def preprocess_mri(mri_path):
+    transforms = prepare_transforms()
+    data_dict = {"image": mri_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):
+        return "Error: MRI file not found"
+
+    # Load the model
+    models = initialize_model()
+
+    # Preprocess MRI
+    image = preprocess_mri(mri_path)
+
+    # Run predictions
+    predictions = []
+    with torch.no_grad():
+        for model in models:
+            pred = model(image)
+            predictions.append(pred.cpu().numpy())
+
+    # Compute median brain age prediction
+    predicted_brain_age = np.median(np.stack(predictions))
+
+    # 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
+    return f"Predicted Brain Age: {predicted_brain_age_corrected:.2f} years", \
+           f"Brain Age Difference (BAD): {brain_age_difference:.2f} years"
+
+# Gradio UI
+iface = 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
+if __name__ == "__main__":
+    iface.launch()