Update app.py

app.py

@@ -17,6 +17,9 @@ model = Conv5_FC3(input_size= [
 model.load_state_dict(checkpoint_state["model"])
 model.eval()
 
+# Class labels
+CLASSES = ["CN", "AD"]
+
 def preprocess_nii(nii_file):
     # Load NIfTI file
     img = nib.load(nii_file)
@@ -39,14 +42,16 @@ def predict(input_image):
     x = preprocess_nii(input_image)
     with torch.no_grad():
         output = model(x)
-    output = output.squeeze(0).cpu().float()
-    return output[0]
+        
+    results = {cls: float(prob) for cls, prob in zip(CLASSES, probs)}
+    
+    return results
 
 # Gradio app: file upload instead of image
 demo = gr.Interface(
     fn=predict,
     inputs=gr.File(type="filepath", label=".nii.gz MRI upload"),
-    outputs="label",
+    outputs="json",
     title="ClinicaDL MRI Classifier",
     description="Upload a .nii.gz file to get the model's prediction."
 )