Update app.py

app.py

@@ -6,8 +6,8 @@ import torch.nn as nn
 import numpy as np
 
 app = FastAPI(
-    title="NeuroHealth Unified EEG API",
-    version="3.0"
+    title="NeuroHealth EEG API",
+    version="3.1"
 )
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
@@ -20,14 +20,14 @@ AD_CLASSES = ["Alzheimer", "FTD", "Control"]
 PD_CLASSES = ["Parkinson", "Control"]
 
 # =========================================================
-# INPUT SCHEMA
+# REQUEST MODEL
 # =========================================================
 
 class EEGRequest(BaseModel):
     features: list
 
 # =========================================================
-# ✔ MLP MODEL (MATCHES YOUR .pth FILES)
+# ✔ TRUE MODEL (MATCHES YOUR TRAINED STATE_DICT)
 # =========================================================
 
 class EEG_MLP(nn.Module):
@@ -35,17 +35,12 @@ class EEG_MLP(nn.Module):
         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.Linear(input_dim, 256),
             nn.BatchNorm1d(256),
             nn.ReLU(),
-            nn.Dropout(0.3),
 
             nn.Linear(256, 128),
+            nn.BatchNorm1d(128),
             nn.ReLU(),
 
             nn.Linear(128, output_dim)
@@ -55,70 +50,36 @@ class EEG_MLP(nn.Module):
         return self.net(x)
 
 # =========================================================
-# OPTIONAL CNN (ONLY IF YOU TRAIN TRUE EEG CNN LATER)
+# LOAD MODELS (CRITICAL: INPUT DIM = 76)
 # =========================================================
 
-class EEG_CNN(nn.Module):
-    def __init__(self, output_dim):
-        super().__init__()
-
-        self.conv1 = nn.Conv1d(19, 32, kernel_size=7)
-        self.bn1 = nn.BatchNorm1d(32)
-
-        self.conv2 = nn.Conv1d(32, 64, kernel_size=5)
-        self.bn2 = nn.BatchNorm1d(64)
-
-        self.conv3 = nn.Conv1d(64, 128, kernel_size=3)
-        self.bn3 = nn.BatchNorm1d(128)
+INPUT_DIM = 76
 
-        self.pool = nn.AdaptiveAvgPool1d(1)
-        self.fc = nn.Linear(128, output_dim)
-
-    def forward(self, x):
-        x = torch.relu(self.bn1(self.conv1(x)))
-        x = torch.relu(self.bn2(self.conv2(x)))
-        x = torch.relu(self.bn3(self.conv3(x)))
-
-        x = self.pool(x).squeeze(-1)
-        return self.fc(x)
-
-# =========================================================
-# LOAD MODELS (BASED ON YOUR REAL FILES)
-# =========================================================
-
-# ✔ THESE ARE MLP MODELS (CONFIRMED BY "net.*" KEYS)
-
-ad_model = EEG_MLP(95, 3).to(DEVICE)
+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.eval()
 
-pd_model = EEG_MLP(95, 2).to(DEVICE)
+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.eval()
 
-print("All models loaded successfully")
+print("Models loaded successfully with correct architecture")
 
 # =========================================================
-# SAFE PREDICTION ENGINE
+# PREDICTION ENGINE
 # =========================================================
 
-def predict(model, features, classes, model_type="mlp"):
+def predict(model, features, classes):
     x = torch.tensor(features, dtype=torch.float32)
 
-    if x.numel() != 95:
-        raise ValueError("Expected exactly 95 features")
-
-    # ✔ MLP MODE: flat input
-    if model_type == "mlp":
-        x = x.unsqueeze(0).to(DEVICE)
+    if x.numel() != INPUT_DIM:
+        raise ValueError(f"Expected {INPUT_DIM} features, got {x.numel()}")
 
-    # ✔ CNN MODE (future use only)
-    elif model_type == "cnn":
-        x = x.view(19, 5).unsqueeze(0).to(DEVICE)
+    x = x.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))
 
@@ -136,20 +97,15 @@ def predict(model, features, classes, model_type="mlp"):
 
 @app.get("/")
 def home():
-    return {"message": "NeuroHealth Unified API Running"}
-
-# =========================================================
-# AD PREDICTION
-# =========================================================
+    return {
+        "message": "NeuroHealth EEG API Running",
+        "input_dim": INPUT_DIM
+    }
 
 @app.post("/predict/ad")
 def predict_ad(req: EEGRequest):
-    return predict(ad_model, req.features, AD_CLASSES, "mlp")
-
-# =========================================================
-# PD PREDICTION
-# =========================================================
+    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, "mlp")
+    return predict(pd_model, req.features, PD_CLASSES)