Update app.py

app.py

@@ -6,8 +6,8 @@ import torch.nn as nn
 import numpy as np
 
 app = FastAPI(
-    title="NeuroHealth EEG API",
-    version="2.0"
+    title="NeuroHealth Unified EEG API",
+    version="3.0"
 )
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
@@ -20,14 +20,42 @@ AD_CLASSES = ["Alzheimer", "FTD", "Control"]
 PD_CLASSES = ["Parkinson", "Control"]
 
 # =========================================================
-# REQUEST SCHEMA
+# INPUT SCHEMA
 # =========================================================
 
 class EEGRequest(BaseModel):
     features: list
 
 # =========================================================
-# EEG CNN MODEL (MATCHES YOUR .pth FILES)
+# ✔ MLP MODEL (MATCHES YOUR .pth FILES)
+# =========================================================
+
+class EEG_MLP(nn.Module):
+    def __init__(self, input_dim, output_dim):
+        super().__init__()
+
+        self.net = nn.Sequential(
+            nn.Linear(input_dim, 512),
+            nn.BatchNorm1d(512),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+
+            nn.Linear(512, 256),
+            nn.BatchNorm1d(256),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+
+            nn.Linear(256, 128),
+            nn.ReLU(),
+
+            nn.Linear(128, output_dim)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+# =========================================================
+# OPTIONAL CNN (ONLY IF YOU TRAIN TRUE EEG CNN LATER)
 # =========================================================
 
 class EEG_CNN(nn.Module):
@@ -55,42 +83,48 @@ class EEG_CNN(nn.Module):
         return self.fc(x)
 
 # =========================================================
-# LOAD MODELS (CLEAN VERSION)
+# LOAD MODELS (BASED ON YOUR REAL FILES)
 # =========================================================
 
-ad_model = EEG_CNN(3).to(DEVICE)
+# ✔ THESE ARE MLP MODELS (CONFIRMED BY "net.*" KEYS)
+
+ad_model = EEG_MLP(95, 3).to(DEVICE)
 ad_model.load_state_dict(torch.load("AD_eeg_cnn.pth", map_location=DEVICE))
 ad_model.eval()
 
-pd_model = EEG_CNN(2).to(DEVICE)
+pd_model = EEG_MLP(95, 2).to(DEVICE)
 pd_model.load_state_dict(torch.load("PD_eeg_cnn.pth", map_location=DEVICE))
 pd_model.eval()
 
-print("All EEG CNN models loaded successfully")
+print("All models loaded successfully")
 
 # =========================================================
-# CORE PREDICTION FUNCTION
+# SAFE PREDICTION ENGINE
 # =========================================================
 
-def predict(model, features, classes):
+def predict(model, features, classes, model_type="mlp"):
     x = torch.tensor(features, dtype=torch.float32)
 
     if x.numel() != 95:
         raise ValueError("Expected exactly 95 features")
 
-    # reshape EEG features → (19 channels × 5 timesteps)
-    x = x.view(19, 5).unsqueeze(0).to(DEVICE)
+    # ✔ MLP MODE: flat input
+    if model_type == "mlp":
+        x = x.unsqueeze(0).to(DEVICE)
+
+    # ✔ CNN MODE (future use only)
+    elif model_type == "cnn":
+        x = x.view(19, 5).unsqueeze(0).to(DEVICE)
 
     with torch.no_grad():
-        outputs = model(x)
-        probs = torch.softmax(outputs, dim=1).cpu().numpy()[0]
+        out = model(x)
+        probs = torch.softmax(out, dim=1).cpu().numpy()[0]
 
     pred = int(np.argmax(probs))
-    confidence = float(probs[pred])
 
     return {
         "prediction": classes[pred],
-        "confidence": confidence,
+        "confidence": float(probs[pred]),
         "probabilities": {
             classes[i]: float(probs[i]) for i in range(len(classes))
         }
@@ -102,20 +136,20 @@ def predict(model, features, classes):
 
 @app.get("/")
 def home():
-    return {"message": "NeuroHealth EEG API Running"}
+    return {"message": "NeuroHealth Unified API Running"}
 
 # =========================================================
-# AD ROUTE
+# AD PREDICTION
 # =========================================================
 
 @app.post("/predict/ad")
-def predict_ad(request: EEGRequest):
-    return predict(ad_model, request.features, AD_CLASSES)
+def predict_ad(req: EEGRequest):
+    return predict(ad_model, req.features, AD_CLASSES, "mlp")
 
 # =========================================================
-# PD ROUTE
+# PD PREDICTION
 # =========================================================
 
 @app.post("/predict/pd")
-def predict_pd(request: EEGRequest):
-    return predict(pd_model, request.features, PD_CLASSES)
+def predict_pd(req: EEGRequest):
+    return predict(pd_model, req.features, PD_CLASSES, "mlp")