Update app.py

app.py

@@ -6,7 +6,7 @@ import torch.nn as nn
 import numpy as np
 
 app = FastAPI(
-    title="NeuroHealth Multi-Model EEG API",
+    title="NeuroHealth EEG API",
     version="2.0"
 )
 
@@ -20,53 +20,16 @@ AD_CLASSES = ["Alzheimer", "FTD", "Control"]
 PD_CLASSES = ["Parkinson", "Control"]
 
 # =========================================================
-# INPUT
+# REQUEST SCHEMA
 # =========================================================
 
 class EEGRequest(BaseModel):
     features: list
 
 # =========================================================
-# DEVICE
+# EEG CNN MODEL (MATCHES YOUR .pth FILES)
 # =========================================================
 
-def to_tensor(features):
-    return torch.tensor(features, dtype=torch.float32)
-
-# =========================================================
-# ===================== 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__()
@@ -92,48 +55,46 @@ class EEG_CNN(nn.Module):
         return self.fc(x)
 
 # =========================================================
-# LOAD ALL MODELS
+# LOAD MODELS (CLEAN VERSION)
 # =========================================================
 
-# -------- 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_cnn = EEG_CNN(3).to(DEVICE)
-ad_cnn.load_state_dict(torch.load("AD_eeg_cnn.pth", map_location=DEVICE))
-ad_cnn.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()
+ad_model = EEG_CNN(3).to(DEVICE)
+ad_model.load_state_dict(torch.load("AD_eeg_cnn.pth", map_location=DEVICE))
+ad_model.eval()
 
-pd_cnn = EEG_CNN(2).to(DEVICE)
-pd_cnn.load_state_dict(torch.load("PD_eeg_cnn.pth", map_location=DEVICE))
-pd_cnn.eval()
+pd_model = EEG_CNN(2).to(DEVICE)
+pd_model.load_state_dict(torch.load("PD_eeg_cnn.pth", map_location=DEVICE))
+pd_model.eval()
 
-print("All models loaded successfully")
+print("All EEG CNN models loaded successfully")
 
 # =========================================================
-# PREDICTION ENGINE
+# CORE PREDICTION FUNCTION
 # =========================================================
 
-def predict(model, features, mode="mlp"):
-    x = to_tensor(features)
+def predict(model, features, classes):
+    x = torch.tensor(features, dtype=torch.float32)
 
-    if mode == "mlp":
-        x = x.unsqueeze(0).to(DEVICE)
+    if x.numel() != 95:
+        raise ValueError("Expected exactly 95 features")
 
-    elif mode == "cnn":
-        x = x.view(19, 5).unsqueeze(0).to(DEVICE)
+    # reshape EEG features → (19 channels × 5 timesteps)
+    x = x.view(19, 5).unsqueeze(0).to(DEVICE)
 
     with torch.no_grad():
-        out = model(x)
-        probs = torch.softmax(out, dim=1).cpu().numpy()[0]
+        outputs = model(x)
+        probs = torch.softmax(outputs, dim=1).cpu().numpy()[0]
 
     pred = int(np.argmax(probs))
-    return pred, float(probs[pred]), probs.tolist()
+    confidence = float(probs[pred])
+
+    return {
+        "prediction": classes[pred],
+        "confidence": confidence,
+        "probabilities": {
+            classes[i]: float(probs[i]) for i in range(len(classes))
+        }
+    }
 
 # =========================================================
 # ROUTES
@@ -141,58 +102,20 @@ def predict(model, features, mode="mlp"):
 
 @app.get("/")
 def home():
-    return {"message": "NeuroHealth Multi-Model API Running"}
+    return {"message": "NeuroHealth EEG API Running"}
 
 # =========================================================
-# AD ENDPOINTS
+# AD ROUTE
 # =========================================================
 
-@app.post("/predict/ad/mlp")
-def ad_mlp_predict(req: EEGRequest):
-    p, c, prob = predict(ad_mlp, req.features, "mlp")
-
-    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 {
-        "model": "AD_EEG_CNN",
-        "prediction": AD_CLASSES[p],
-        "confidence": c,
-        "probabilities": dict(zip(AD_CLASSES, prob))
-    }
+@app.post("/predict/ad")
+def predict_ad(request: EEGRequest):
+    return predict(ad_model, request.features, AD_CLASSES)
 
 # =========================================================
-# PD ENDPOINTS
+# PD ROUTE
 # =========================================================
 
-@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))
-    }
-
-
-@app.post("/predict/pd/eeg")
-def pd_cnn_predict(req: EEGRequest):
-    p, c, prob = predict(pd_cnn, req.features, "cnn")
-
-    return {
-        "model": "PD_EEG_CNN",
-        "prediction": PD_CLASSES[p],
-        "confidence": c,
-        "probabilities": dict(zip(PD_CLASSES, prob))
-    }
+@app.post("/predict/pd")
+def predict_pd(request: EEGRequest):
+    return predict(pd_model, request.features, PD_CLASSES)