Update app.py

app.py

@@ -1,8 +1,10 @@
 import gradio as gr
 from huggingface_hub import hf_hub_download
 import torch
+import torch.nn.functional as F
 from clinicadl.utils.network.cnn.models import Conv5_FC3
 import nibabel as nib 
+import numpy as np
 
 # Download model from Hub
 model_path = hf_hub_download(repo_id="ARAMIS-LAB/CNN-AD-CN", filename="model.pth.tar")
@@ -27,7 +29,7 @@ def preprocess_nii(nii_file):
     data = img.get_fdata()  # numpy array (float64)
 
     # Normalize intensities
-    data = (data - np.mean(data)) / (np.std(data) + 1e-8)
+    data = (data - np.mean(data)) / np.std(data)
 
     # Convert to tensor
     tensor = torch.tensor(data, dtype=torch.float32).unsqueeze(0).unsqueeze(0)  
@@ -43,6 +45,7 @@ def predict(input_image):
     x = preprocess_nii(input_image)
     with torch.no_grad():
         output = model(x)
+        probs = F.softmax(logits, dim=1)  # convert to probabilities
         
     results = {cls: float(prob) for cls, prob in zip(CLASSES, probs)}