Update app.py

app.py

@@ -6,34 +6,72 @@ import torch.nn as nn
 import numpy as np
 
 app = FastAPI(
-    title="NeuroHealth EEG API",
-    version="1.0"
+    title="NeuroHealth Multi-Model EEG API",
+    version="2.0"
 )
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 # =========================================================
-# CLASS LABELS
+# LABELS
 # =========================================================
 
 AD_CLASSES = ["Alzheimer", "FTD", "Control"]
 PD_CLASSES = ["Parkinson", "Control"]
 
 # =========================================================
-# EXPECTED INPUT
+# INPUT
 # =========================================================
 
-INPUT_FEATURES = 95  # must reshape to (19, 5)
+class EEGRequest(BaseModel):
+    features: list
+
+# =========================================================
+# DEVICE
+# =========================================================
+
+def to_tensor(features):
+    return torch.tensor(features, dtype=torch.float32)
 
 # =========================================================
-# CNN MODEL (MATCHES YOUR CHECKPOINT EXACTLY)
+# ===================== MODELS ============================
 # =========================================================
 
+# ---------------- MLP (CSV MODELS) ----------------------
+
+INPUT_DIM = 95
+
+class MLP(nn.Module):
+    def __init__(self, output_dim):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(INPUT_DIM, 512),
+            nn.BatchNorm1d(512),
+            nn.GELU(),
+            nn.Dropout(0.3),
+
+            nn.Linear(512, 256),
+            nn.BatchNorm1d(256),
+            nn.GELU(),
+            nn.Dropout(0.3),
+
+            nn.Linear(256, 128),
+            nn.GELU(),
+
+            nn.Linear(128, output_dim)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+# ---------------- CNN (EEG MODELS) ----------------------
+
 class EEG_CNN(nn.Module):
     def __init__(self, output_dim):
         super().__init__()
 
-        self.conv1 = nn.Conv1d(in_channels=19, out_channels=32, kernel_size=7)
+        self.conv1 = nn.Conv1d(19, 32, kernel_size=7)
         self.bn1 = nn.BatchNorm1d(32)
 
         self.conv2 = nn.Conv1d(32, 64, kernel_size=5)
@@ -50,55 +88,52 @@ class EEG_CNN(nn.Module):
         x = torch.relu(self.bn2(self.conv2(x)))
         x = torch.relu(self.bn3(self.conv3(x)))
 
-        x = self.pool(x)
-        x = x.squeeze(-1)
+        x = self.pool(x).squeeze(-1)
         return self.fc(x)
 
 # =========================================================
-# LOAD MODELS
+# LOAD ALL MODELS
 # =========================================================
 
-ad_model = EEG_CNN(3).to(DEVICE)
-pd_model = EEG_CNN(2).to(DEVICE)
+# -------- AD MODELS --------
+ad_mlp = MLP(3).to(DEVICE)
+ad_mlp.load_state_dict(torch.load("AD_MLP.pt", map_location=DEVICE))
+ad_mlp.eval()
 
-ad_model.load_state_dict(torch.load("AD_MLP.pt", map_location=DEVICE))
-pd_model.load_state_dict(torch.load("PD_MLP.pt", map_location=DEVICE))
+ad_cnn = EEG_CNN(3).to(DEVICE)
+ad_cnn.load_state_dict(torch.load("AD_eeg_cnn.pth", map_location=DEVICE))
+ad_cnn.eval()
 
-ad_model.eval()
-pd_model.eval()
+# -------- PD MODELS --------
+pd_mlp = MLP(2).to(DEVICE)
+pd_mlp.load_state_dict(torch.load("PD_MLP.pt", map_location=DEVICE))
+pd_mlp.eval()
 
-print("Models Loaded Successfully")
-
-# =========================================================
-# REQUEST SCHEMA
-# =========================================================
+pd_cnn = EEG_CNN(2).to(DEVICE)
+pd_cnn.load_state_dict(torch.load("PD_eeg_cnn.pth", map_location=DEVICE))
+pd_cnn.eval()
 
-class EEGRequest(BaseModel):
-    features: list
+print("All models loaded successfully")
 
 # =========================================================
-# CORE PREDICTION FUNCTION
+# PREDICTION ENGINE
 # =========================================================
 
-def predict_model(model, features):
-    x = torch.tensor(features, dtype=torch.float32)
+def predict(model, features, mode="mlp"):
+    x = to_tensor(features)
 
-    # reshape 95 -> (19, 5)
-    if x.numel() != 95:
-        raise ValueError("Expected 95 features")
+    if mode == "mlp":
+        x = x.unsqueeze(0).to(DEVICE)
 
-    x = x.view(19, 5)       # (channels, time)
-    x = x.unsqueeze(0)      # (batch, 19, 5)
-    x = x.to(DEVICE)
+    elif mode == "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 pred, confidence, probs.tolist()
+    return pred, float(probs[pred]), probs.tolist()
 
 # =========================================================
 # ROUTES
@@ -106,40 +141,58 @@ def predict_model(model, features):
 
 @app.get("/")
 def home():
-    return {"message": "NeuroHealth EEG API Running"}
+    return {"message": "NeuroHealth Multi-Model API Running"}
 
 # =========================================================
-# ALZHEIMER
+# AD ENDPOINTS
 # =========================================================
 
-@app.post("/predict/alzheimer")
-def predict_alzheimer(request: EEGRequest):
+@app.post("/predict/ad/mlp")
+def ad_mlp_predict(req: EEGRequest):
+    p, c, prob = predict(ad_mlp, req.features, "mlp")
 
-    pred, confidence, probs = predict_model(ad_model, request.features)
+    return {
+        "model": "AD_MLP",
+        "prediction": AD_CLASSES[p],
+        "confidence": c,
+        "probabilities": dict(zip(AD_CLASSES, prob))
+    }
+
+
+@app.post("/predict/ad/eeg")
+def ad_cnn_predict(req: EEGRequest):
+    p, c, prob = predict(ad_cnn, req.features, "cnn")
 
     return {
-        "prediction": AD_CLASSES[pred],
-        "confidence": confidence,
-        "probabilities": {
-            AD_CLASSES[i]: float(probs[i])
-            for i in range(len(AD_CLASSES))
-        }
+        "model": "AD_EEG_CNN",
+        "prediction": AD_CLASSES[p],
+        "confidence": c,
+        "probabilities": dict(zip(AD_CLASSES, prob))
     }
 
 # =========================================================
-# PARKINSON
+# PD ENDPOINTS
 # =========================================================
 
-@app.post("/predict/parkinson")
-def predict_parkinson(request: EEGRequest):
+@app.post("/predict/pd/mlp")
+def pd_mlp_predict(req: EEGRequest):
+    p, c, prob = predict(pd_mlp, req.features, "mlp")
+
+    return {
+        "model": "PD_MLP",
+        "prediction": PD_CLASSES[p],
+        "confidence": c,
+        "probabilities": dict(zip(PD_CLASSES, prob))
+    }
+
 
-    pred, confidence, probs = predict_model(pd_model, request.features)
+@app.post("/predict/pd/eeg")
+def pd_cnn_predict(req: EEGRequest):
+    p, c, prob = predict(pd_cnn, req.features, "cnn")
 
     return {
-        "prediction": PD_CLASSES[pred],
-        "confidence": confidence,
-        "probabilities": {
-            PD_CLASSES[i]: float(probs[i])
-            for i in range(len(PD_CLASSES))
-        }
+        "model": "PD_EEG_CNN",
+        "prediction": PD_CLASSES[p],
+        "confidence": c,
+        "probabilities": dict(zip(PD_CLASSES, prob))
     }