Update app.py

app.py

@@ -4,7 +4,7 @@ import torch
 import torch.nn as nn
 import numpy as np
 
-app = FastAPI(title="NeuroHealth EEG API", version="RESET-1")
+app = FastAPI(title="NeuroHealth EEG API", version="RESET-2")
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -15,59 +15,56 @@ AD_CLASSES = ["Alzheimer", "FTD", "Control"]
 PD_CLASSES = ["Parkinson", "Control"]
 
 # ======================
-# INPUT
+# INPUT MODEL
 # ======================
 class EEGRequest(BaseModel):
     features: list
 
 # ======================
-# AD MODEL (UNCHANGED CNN)
+# AD MODEL - 1D CNN
 # ======================
 class EEG_CNN_AD(nn.Module):
     def __init__(self):
         super().__init__()
         self.conv1 = nn.Conv1d(19, 32, 7, padding=3)
         self.bn1 = nn.BatchNorm1d(32)
-
         self.conv2 = nn.Conv1d(32, 64, 5, padding=2)
         self.bn2 = nn.BatchNorm1d(64)
-
         self.conv3 = nn.Conv1d(64, 128, 3, padding=1)
         self.bn3 = nn.BatchNorm1d(128)
-
         self.pool = nn.AdaptiveAvgPool1d(1)
         self.fc = nn.Linear(128, 3)
 
     def forward(self, x):
         x = x.view(x.size(0), 19, 76)
-
         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)
 
 # ======================
-# PD MODEL (MATCH CHECKPOINT EXACTLY)
+# PD MODEL - 1D CNN (Fixed to match checkpoint)
 # ======================
-class EEG_PD(nn.Module):
+class EEG_CNN_PD(nn.Module):
     def __init__(self):
         super().__init__()
-        self.net = nn.Sequential(
-            nn.Linear(76, 256),
-            nn.ReLU(),
-            nn.BatchNorm1d(256),
-
-            nn.Linear(256, 128),
-            nn.ReLU(),
-            nn.BatchNorm1d(128),
-
-            nn.Linear(128, 2)
-        )
+        self.conv1 = nn.Conv1d(19, 32, 7, padding=3)
+        self.bn1 = nn.BatchNorm1d(32)
+        self.conv2 = nn.Conv1d(32, 64, 5, padding=2)
+        self.bn2 = nn.BatchNorm1d(64)
+        self.conv3 = nn.Conv1d(64, 128, 3, padding=1)
+        self.bn3 = nn.BatchNorm1d(128)
+        self.pool = nn.AdaptiveAvgPool1d(1)
+        self.fc = nn.Linear(128, 2)
 
     def forward(self, x):
-        return self.net(x)
+        x = x.view(x.size(0), 19, 76)
+        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
@@ -79,60 +76,59 @@ print("Loading AD model...")
 ad_model = EEG_CNN_AD().to(DEVICE)
 ad_model.load_state_dict(torch.load(AD_PATH, map_location=DEVICE))
 ad_model.eval()
-print("AD loaded")
+print("AD loaded successfully")
 
 print("Loading PD model...")
-pd_model = EEG_PD().to(DEVICE)
+pd_model = EEG_CNN_PD().to(DEVICE)
 pd_model.load_state_dict(torch.load(PD_PATH, map_location=DEVICE))
 pd_model.eval()
-print("PD loaded")
+print("PD loaded successfully")
 
 # ======================
 # PREPROCESS
 # ======================
 def preprocess(features, mode):
-    x = torch.tensor(features, dtype=torch.float32)
-
-    if mode == "ad":
-        return x.view(1, 19, 76)
-
-    if mode == "pd":
-        # IMPORTANT: PD expects FLATTENED 76-D FEATURES
-        return x.view(1, 76)
-
-    raise ValueError("unknown mode")
+    x = torch.tensor(features, dtype=torch.float32).to(DEVICE)
+    if mode == "ad" or mode == "pd":
+        # Ensure shape is (batch, channels, time) -> (1, 19, 76)
+        if len(x.shape) == 1:
+            x = x.view(1, 19, 76)
+        elif len(x.shape) == 2:
+            x = x.unsqueeze(0)
+        return x
+    raise ValueError("Unknown mode")
 
 # ======================
-# PREDICT ENGINE
+# PREDICT FUNCTION
 # ======================
 def predict(model, x, classes):
     with torch.no_grad():
         out = model(x)
         probs = torch.softmax(out, dim=1).cpu().numpy()[0]
-
-    pred = int(np.argmax(probs))
-
-    return {
-        "prediction": classes[pred],
-        "confidence": float(probs[pred]),
-        "probabilities": {
-            classes[i]: float(probs[i]) for i in range(len(classes))
+        pred = int(np.argmax(probs))
+        return {
+            "prediction": classes[pred],
+            "confidence": float(probs[pred]),
+            "probabilities": {classes[i]: float(probs[i]) for i in range(len(classes))}
         }
-    }
 
 # ======================
 # ROUTES
 # ======================
 @app.get("/")
 def home():
-    return {"status": "ready (RESET architecture)", "year": 2026}
+    return {
+        "status": "ready",
+        "message": "NeuroHealth EEG API is running (Fixed CNN architecture)",
+        "year": 2026
+    }
 
 @app.post("/predict/ad")
 def predict_ad(req: EEGRequest):
-    x = preprocess(req.features, "ad").to(DEVICE)
+    x = preprocess(req.features, "ad")
     return predict(ad_model, x, AD_CLASSES)
 
 @app.post("/predict/pd")
 def predict_pd(req: EEGRequest):
-    x = preprocess(req.features, "pd").to(DEVICE)
+    x = preprocess(req.features, "pd")
     return predict(pd_model, x, PD_CLASSES)