Update app.py

app.py

@@ -28,67 +28,63 @@ PD_CLASSES = [
 ]
 
 # =========================================================
-# MODELS
+# CONFIG
 # =========================================================
 
 INPUT_DIM = 95
 
 # =========================================================
-# MLP MODEL
+# CNN MODEL (MATCHES SAVED WEIGHTS)
 # =========================================================
 
-class MLP(nn.Module):
-
+class EEG_CNN(nn.Module):
     def __init__(self, output_dim):
-
         super().__init__()
 
-        self.network = nn.Sequential(
+        self.conv1 = nn.Conv1d(1, 32, kernel_size=3, padding=1)
+        self.bn1 = nn.BatchNorm1d(32)
 
-            nn.Linear(INPUT_DIM, 512),
-            nn.BatchNorm1d(512),
-            nn.GELU(),
-            nn.Dropout(0.3),
+        self.conv2 = nn.Conv1d(32, 64, kernel_size=3, padding=1)
+        self.bn2 = nn.BatchNorm1d(64)
 
-            nn.Linear(512, 256),
-            nn.BatchNorm1d(256),
-            nn.GELU(),
-            nn.Dropout(0.3),
+        self.conv3 = nn.Conv1d(64, 128, kernel_size=3, padding=1)
+        self.bn3 = nn.BatchNorm1d(128)
 
-            nn.Linear(256, 128),
-            nn.BatchNorm1d(128),
-            nn.GELU(),
+        # Global pooling removes dependence on sequence length
+        self.pool = nn.AdaptiveAvgPool1d(1)
 
-            nn.Linear(128, output_dim)
-        )
+        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)))
 
-        return self.network(x)
+        x = self.pool(x)  # (batch, 128, 1)
+        x = x.squeeze(-1)  # (batch, 128)
+
+        return self.fc(x)
 
 # =========================================================
 # LOAD MODELS
 # =========================================================
 
-ad_mlp = MLP(3).to(DEVICE)
-pd_mlp = MLP(2).to(DEVICE)
+ad_model = EEG_CNN(3).to(DEVICE)
+pd_model = EEG_CNN(2).to(DEVICE)
 
-ad_mlp.load_state_dict(
-    torch.load(
-        "AD_MLP.pt",
-        map_location=DEVICE
-    )
+ad_model.load_state_dict(
+    torch.load("AD_MLP.pt", map_location=DEVICE)
 )
 
-pd_mlp.load_state_dict(
-    torch.load(
-        "PD_MLP.pt",
-        map_location=DEVICE
-    )
+pd_model.load_state_dict(
+    torch.load("PD_MLP.pt", map_location=DEVICE)
 )
 
-ad_mlp.eval()
-pd_mlp.eval()
+ad_model.eval()
+pd_model.eval()
 
 print("Models Loaded Successfully")
 
@@ -97,7 +93,6 @@ print("Models Loaded Successfully")
 # =========================================================
 
 class EEGRequest(BaseModel):
-
     features: list
 
 # =========================================================
@@ -105,23 +100,13 @@ class EEGRequest(BaseModel):
 # =========================================================
 
 def predict_model(model, features):
-
-    x = torch.tensor(
-        features,
-        dtype=torch.float32
-    ).unsqueeze(0).to(DEVICE)
+    x = torch.tensor(features, dtype=torch.float32).unsqueeze(0).to(DEVICE)
 
     with torch.no_grad():
-
         outputs = model(x)
-
-        probs = torch.softmax(
-            outputs,
-            dim=1
-        ).cpu().numpy()[0]
+        probs = torch.softmax(outputs, dim=1).cpu().numpy()[0]
 
     pred = int(np.argmax(probs))
-
     confidence = float(probs[pred])
 
     return pred, confidence, probs.tolist()
@@ -131,9 +116,7 @@ def predict_model(model, features):
 # =========================================================
 
 @app.get("/")
-
 def home():
-
     return {
         "message": "NeuroHealth EEG API Running"
     }
@@ -143,24 +126,15 @@ def home():
 # =========================================================
 
 @app.post("/predict/alzheimer")
-
 def predict_alzheimer(request: EEGRequest):
 
-    pred, confidence, probs = predict_model(
-        ad_mlp,
-        request.features
-    )
+    pred, confidence, probs = predict_model(ad_model, request.features)
 
     return {
-
         "prediction": AD_CLASSES[pred],
-
         "confidence": confidence,
-
         "probabilities": {
-
             AD_CLASSES[i]: float(probs[i])
-
             for i in range(len(AD_CLASSES))
         }
     }
@@ -170,24 +144,15 @@ def predict_alzheimer(request: EEGRequest):
 # =========================================================
 
 @app.post("/predict/parkinson")
-
 def predict_parkinson(request: EEGRequest):
 
-    pred, confidence, probs = predict_model(
-        pd_mlp,
-        request.features
-    )
+    pred, confidence, probs = predict_model(pd_model, request.features)
 
     return {
-
         "prediction": PD_CLASSES[pred],
-
         "confidence": confidence,
-
         "probabilities": {
-
             PD_CLASSES[i]: float(probs[i])
-
             for i in range(len(PD_CLASSES))
         }
     }