Update app.py

app.py

@@ -5,19 +5,31 @@ import torch
 import torch.nn as nn
 import numpy as np
 
+# =========================================================
+# APP
+# =========================================================
+
 app = FastAPI(
     title="NeuroHealth EEG API",
-    version="3.1"
+    version="4.0"
 )
 
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # =========================================================
 # LABELS
 # =========================================================
 
-AD_CLASSES = ["Alzheimer", "FTD", "Control"]
-PD_CLASSES = ["Parkinson", "Control"]
+AD_CLASSES = [
+    "Alzheimer",
+    "FTD",
+    "Control"
+]
+
+PD_CLASSES = [
+    "Parkinson",
+    "Control"
+]
 
 # =========================================================
 # REQUEST MODEL
@@ -27,7 +39,7 @@ class EEGRequest(BaseModel):
     features: list
 
 # =========================================================
-# ✔ TRUE MODEL (MATCHES YOUR TRAINED STATE_DICT)
+# MODEL ARCHITECTURE
 # =========================================================
 
 class EEG_MLP(nn.Module):
@@ -36,13 +48,15 @@ class EEG_MLP(nn.Module):
 
         self.net = nn.Sequential(
             nn.Linear(input_dim, 256),
-            nn.BatchNorm1d(256),
             nn.ReLU(),
 
+            nn.BatchNorm1d(256),
+
             nn.Linear(256, 128),
-            nn.BatchNorm1d(128),
             nn.ReLU(),
 
+            nn.BatchNorm1d(128),
+
             nn.Linear(128, output_dim)
         )
 
@@ -50,44 +64,81 @@ class EEG_MLP(nn.Module):
         return self.net(x)
 
 # =========================================================
-# LOAD MODELS (CRITICAL: INPUT DIM = 76)
+# SETTINGS
 # =========================================================
 
 INPUT_DIM = 76
 
+AD_MODEL_PATH = "AD_eeg_cnn_ad_ftd_cn.pt"
+PD_MODEL_PATH = "PD_eeg_cnn.pth"
+
+# =========================================================
+# LOAD MODELS
+# =========================================================
+
+print("Loading AD model...")
+
 ad_model = EEG_MLP(INPUT_DIM, 3).to(DEVICE)
-ad_model.load_state_dict(torch.load("AD_eeg_cnn.pth", map_location=DEVICE))
+
+ad_model.load_state_dict(
+    torch.load(
+        AD_MODEL_PATH,
+        map_location=DEVICE
+    )
+)
+
 ad_model.eval()
 
+print("AD model loaded successfully")
+
+print("Loading PD model...")
+
 pd_model = EEG_MLP(INPUT_DIM, 2).to(DEVICE)
-pd_model.load_state_dict(torch.load("PD_eeg_cnn.pth", map_location=DEVICE))
+
+pd_model.load_state_dict(
+    torch.load(
+        PD_MODEL_PATH,
+        map_location=DEVICE
+    )
+)
+
 pd_model.eval()
 
-print("Models loaded successfully with correct architecture")
+print("PD model loaded successfully")
+
+print("All models loaded")
 
 # =========================================================
-# PREDICTION ENGINE
+# PREDICTION FUNCTION
 # =========================================================
 
-def predict(model, features, classes):
-    x = torch.tensor(features, dtype=torch.float32)
+def predict(model, features, class_names):
+
+    x = torch.tensor(
+        features,
+        dtype=torch.float32
+    )
 
     if x.numel() != INPUT_DIM:
-        raise ValueError(f"Expected {INPUT_DIM} features, got {x.numel()}")
+        raise ValueError(
+            f"Expected {INPUT_DIM} features but received {x.numel()}"
+        )
 
     x = x.unsqueeze(0).to(DEVICE)
 
     with torch.no_grad():
         outputs = model(x)
-        probs = torch.softmax(outputs, dim=1).cpu().numpy()[0]
+        probs = torch.softmax(outputs, dim=1)
+        probs = probs.cpu().numpy()[0]
 
-    pred = int(np.argmax(probs))
+    pred_idx = int(np.argmax(probs))
 
     return {
-        "prediction": classes[pred],
-        "confidence": float(probs[pred]),
+        "prediction": class_names[pred_idx],
+        "confidence": float(probs[pred_idx]),
         "probabilities": {
-            classes[i]: float(probs[i]) for i in range(len(classes))
+            class_names[i]: float(probs[i])
+            for i in range(len(class_names))
         }
     }
 
@@ -98,14 +149,31 @@ def predict(model, features, classes):
 @app.get("/")
 def home():
     return {
-        "message": "NeuroHealth EEG API Running",
-        "input_dim": INPUT_DIM
+        "status": "running",
+        "device": str(DEVICE),
+        "input_features": INPUT_DIM,
+        "ad_classes": AD_CLASSES,
+        "pd_classes": PD_CLASSES
+    }
+
+@app.get("/health")
+def health():
+    return {
+        "status": "healthy"
     }
 
 @app.post("/predict/ad")
 def predict_ad(req: EEGRequest):
-    return predict(ad_model, req.features, AD_CLASSES)
+    return predict(
+        ad_model,
+        req.features,
+        AD_CLASSES
+    )
 
 @app.post("/predict/pd")
 def predict_pd(req: EEGRequest):
-    return predict(pd_model, req.features, PD_CLASSES)
+    return predict(
+        pd_model,
+        req.features,
+        PD_CLASSES
+    )