Update app.py

app.py

@@ -13,59 +13,45 @@ app = FastAPI(
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 # =========================================================
-# CLASS NAMES
+# CLASS LABELS
 # =========================================================
 
-AD_CLASSES = [
-    "Alzheimer",
-    "FTD",
-    "Control"
-]
-
-PD_CLASSES = [
-    "Parkinson",
-    "Control"
-]
+AD_CLASSES = ["Alzheimer", "FTD", "Control"]
+PD_CLASSES = ["Parkinson", "Control"]
 
 # =========================================================
-# CONFIG
+# EXPECTED INPUT
 # =========================================================
 
-INPUT_DIM = 95
+INPUT_FEATURES = 95  # must reshape to (19, 5)
 
 # =========================================================
-# CNN MODEL (MATCHES SAVED WEIGHTS)
+# CNN MODEL (MATCHES YOUR CHECKPOINT EXACTLY)
 # =========================================================
 
 class EEG_CNN(nn.Module):
     def __init__(self, output_dim):
         super().__init__()
 
-        self.conv1 = nn.Conv1d(1, 32, kernel_size=3, padding=1)
+        self.conv1 = nn.Conv1d(in_channels=19, out_channels=32, kernel_size=7)
         self.bn1 = nn.BatchNorm1d(32)
 
-        self.conv2 = nn.Conv1d(32, 64, kernel_size=3, padding=1)
+        self.conv2 = nn.Conv1d(32, 64, kernel_size=5)
         self.bn2 = nn.BatchNorm1d(64)
 
-        self.conv3 = nn.Conv1d(64, 128, kernel_size=3, padding=1)
+        self.conv3 = nn.Conv1d(64, 128, kernel_size=3)
         self.bn3 = nn.BatchNorm1d(128)
 
-        # Global pooling removes dependence on sequence length
         self.pool = nn.AdaptiveAvgPool1d(1)
-
         self.fc = nn.Linear(128, output_dim)
 
     def forward(self, x):
-        # x shape: (batch, features)
-        x = x.unsqueeze(1)  # (batch, 1, 95)
-
         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)  # (batch, 128, 1)
-        x = x.squeeze(-1)  # (batch, 128)
-
+        x = self.pool(x)
+        x = x.squeeze(-1)
         return self.fc(x)
 
 # =========================================================
@@ -75,13 +61,8 @@ class EEG_CNN(nn.Module):
 ad_model = EEG_CNN(3).to(DEVICE)
 pd_model = EEG_CNN(2).to(DEVICE)
 
-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_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_model.eval()
 pd_model.eval()
@@ -89,18 +70,26 @@ pd_model.eval()
 print("Models Loaded Successfully")
 
 # =========================================================
-# REQUEST MODEL
+# REQUEST SCHEMA
 # =========================================================
 
 class EEGRequest(BaseModel):
     features: list
 
 # =========================================================
-# UTILITY
+# CORE PREDICTION FUNCTION
 # =========================================================
 
 def predict_model(model, features):
-    x = torch.tensor(features, dtype=torch.float32).unsqueeze(0).to(DEVICE)
+    x = torch.tensor(features, dtype=torch.float32)
+
+    # reshape 95 -> (19, 5)
+    if x.numel() != 95:
+        raise ValueError("Expected 95 features")
+
+    x = x.view(19, 5)       # (channels, time)
+    x = x.unsqueeze(0)      # (batch, 19, 5)
+    x = x.to(DEVICE)
 
     with torch.no_grad():
         outputs = model(x)
@@ -112,17 +101,15 @@ def predict_model(model, features):
     return pred, confidence, probs.tolist()
 
 # =========================================================
-# ROOT
+# ROUTES
 # =========================================================
 
 @app.get("/")
 def home():
-    return {
-        "message": "NeuroHealth EEG API Running"
-    }
+    return {"message": "NeuroHealth EEG API Running"}
 
 # =========================================================
-# ALZHEIMER PREDICTION
+# ALZHEIMER
 # =========================================================
 
 @app.post("/predict/alzheimer")
@@ -140,7 +127,7 @@ def predict_alzheimer(request: EEGRequest):
     }
 
 # =========================================================
-# PARKINSON PREDICTION
+# PARKINSON
 # =========================================================
 
 @app.post("/predict/parkinson")