Update app.py

app.py

@@ -1,21 +1,17 @@
 import torch
 import gradio as gr
+import plotly.graph_objects as go
 from nilearn import datasets
 from nilearn.connectome import ConnectivityMeasure
 from nilearn.maskers import MultiNiftiMapsMasker
 import numpy as np
 
-#device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# Force torch to use CPU only
-device = torch.device("cpu")
-
-
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+# Load the model
 try:
-    scripted_model = torch.jit.load("fmri_encoder_commercial.pt", map_location="cpu")
+    scripted_model = torch.jit.load("fmri_encoder_commercial.pt", map_location=device)
     
-    # If the model is wrapped in DataParallel, unwrap it
     if isinstance(scripted_model, torch.nn.DataParallel):
         scripted_model = scripted_model.module
     
@@ -26,7 +22,7 @@ except Exception as e:
     exit(1)
 
 # Fetch atlas (e.g., DiFuMo)
-dim = 64  # Number of ROIs
+dim = 64  
 try:
     difumo = datasets.fetch_atlas_difumo(dimension=dim, resolution_mm=2, legacy_format=False)
     atlas_filename = difumo.maps
@@ -34,7 +30,7 @@ except Exception as e:
     print(f"Error fetching atlas: {str(e)}")
     exit(1)
 
-# Create masker to extract features
+# Create masker
 masker = MultiNiftiMapsMasker(
     maps_img=atlas_filename,
     standardize=True,
@@ -45,13 +41,12 @@ masker = MultiNiftiMapsMasker(
 # Connectivity measure
 connectome_measure = ConnectivityMeasure(kind='correlation', vectorize=True, discard_diagonal=True)
 
-# Modified feature extraction function
+# Feature extraction function
 def extract_features_multiple(func_preproc_files):
     all_features = []
     if not func_preproc_files:
         return all_features
     
-    # Fit the masker on the first subject
     print("Fitting masker on the first subject...")
     masker.fit(func_preproc_files[0])
     
@@ -64,42 +59,66 @@ def extract_features_multiple(func_preproc_files):
     print("All subjects processed.")
     return all_features
 
-# Prediction function with error handling
-def validate_inputs(features_tensor):
-    if features_tensor.shape[1] != 2016:
-        raise ValueError(f"Expected 2016 features but got {features_tensor.shape[1]}")
-    if features_tensor.dim() != 2:
-        raise ValueError(f"Expected 2D tensor but got {features_tensor.dim()}D")
+# Function to generate a Plotly probability plot
+def plot_probability(probability):
+    labels = ["No Autism", "Autism"]
+    probs = [1 - probability, probability]
+    colors = ["#6a0dad", "#d896ff"]  # Dark purple and light purple
+
+    fig = go.Figure()
+
+    fig.add_trace(go.Bar(
+        x=labels,
+        y=probs,
+        marker=dict(color=colors),
+        text=[f"{(1-probability)*100:.1f}%", f"{probability*100:.1f}%"],
+        textposition="auto",
+    ))
+
+    fig.update_layout(
+        title="Autism Prediction Probability",
+        paper_bgcolor="black",
+        plot_bgcolor="black",
+        font=dict(color="white"),
+        xaxis=dict(title="Diagnosis", showgrid=False),
+        yaxis=dict(title="Probability", showgrid=True, gridcolor="gray"),
+    )
+
+    return fig
 
+# Prediction function
 def predict_autism(fmri_files, age, gender):
     try:
         if not fmri_files:
-            return "Please upload at least one valid .nii.gz file."
+            return "Please upload at least one valid .nii.gz file.", None
         
         features_list = extract_features_multiple(fmri_files)
         if not features_list:
-            return "Error: Failed to extract features from the fMRI files."
+            return "Error: Failed to extract features from the fMRI files.", None
         
-        # CORRECTED TENSOR SHAPES
-        age_tensor = torch.tensor([float(age)], dtype=torch.float32).to(device)  # Shape: [1]
-        gender_tensor = torch.tensor([int(gender)], dtype=torch.long).to(device)  # Shape: [1]
+        age_tensor = torch.tensor([float(age)], dtype=torch.float32).to(device)
+        gender_tensor = torch.tensor([int(gender)], dtype=torch.long).to(device)
 
         predictions = []
+        plots = []
+
         for features in features_list:
             features_tensor = torch.tensor(features, dtype=torch.float32).unsqueeze(0).to(device)
-            validate_inputs(features_tensor)
-            
+
             with torch.no_grad():
                 prediction = scripted_model(features_tensor, age_tensor, gender_tensor)
                 probability = torch.sigmoid(prediction).item()
             
             result = f"Prediction: {'Autism' if probability > 0.5 else 'No Autism'} (Confidence: {probability:.2%})"
             predictions.append(result)
+
+            # Generate Plotly probability plot
+            plots.append(plot_probability(probability))
         
-        return "\n".join(predictions)
-    
+        return "\n".join(predictions), plots[0]  # Return text and Plotly figure
+
     except Exception as e:
-        return f"Error: {str(e)}"
+        return f"Error: {str(e)}", None
 
 # Gradio interface
 iface = gr.Interface(
@@ -109,11 +128,14 @@ iface = gr.Interface(
         gr.Number(label="Age", minimum=0, maximum=120),
         gr.Radio(["0", "1"], label="Gender (0: Female, 1: Male)"),
     ],
-    outputs=gr.Text(label="Prediction Result"),
+    outputs=[
+        gr.Text(label="Prediction Result"),
+        gr.Plot(label="Prediction Probability Plot"),
+    ],
     title="Autism Prediction from fMRI Data",
     description="Upload one or more preprocessed fMRI files (.nii.gz) and enter the subject's age and gender to predict autism.",
     theme="default",
     flagging_mode="never"
 )
 
-iface.launch()
+iface.launch()