Update app.py

app.py

@@ -11,7 +11,7 @@ import numpy as np
 
 app = FastAPI(
     title="NeuroHealth EEG API",
-    version="9.0"
+    version="10.0"
 )
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -31,7 +31,7 @@ class EEGRequest(BaseModel):
     features: list
 
 # =========================================================
-# AD MODEL (CNN - EXACTLY FROM CHECKPOINT STRUCTURE)
+# AD MODEL (UNCHANGED - 19 x 76 CNN)
 # =========================================================
 
 class EEG_CNN_AD(nn.Module):
@@ -66,36 +66,45 @@ class EEG_CNN_AD(nn.Module):
         return self.fc(x)
 
 # =========================================================
-# PD MODEL (STRICT MATCH TO YOUR CHECKPOINT KEYS)
+# PD MODEL (UPDATED CNN - MATCHES YOUR TRAINED MODEL)
 # =========================================================
 
-class EEG_PD(nn.Module):
+class EEG_CNN_PD(nn.Module):
     def __init__(self):
         super().__init__()
 
-        self.net = nn.Sequential(
-            nn.Linear(76, 256),   # net.0
-            nn.ReLU(),
-
-            nn.BatchNorm1d(256),  # net.2
+        self.conv1 = nn.Conv1d(64, 32, 7, padding=3)
+        self.bn1 = nn.BatchNorm1d(32)
 
-            nn.Linear(256, 128),  # net.4
-            nn.ReLU(),
+        self.conv2 = nn.Conv1d(32, 64, 5, padding=2)
+        self.bn2 = nn.BatchNorm1d(64)
 
-            nn.BatchNorm1d(128),  # net.6
+        self.conv3 = nn.Conv1d(64, 128, 3, padding=1)
+        self.bn3 = nn.BatchNorm1d(128)
 
-            nn.Linear(128, 2)     # net.7
-        )
+        self.pool = nn.AdaptiveAvgPool1d(1)
+        self.fc = nn.Linear(128, 2)
 
     def forward(self, x):
-        return self.net(x)
+        # x: (batch, 64, 256)
+        x = torch.relu(self.conv1(x))
+        x = self.bn1(x)
+
+        x = torch.relu(self.conv2(x))
+        x = self.bn2(x)
+
+        x = torch.relu(self.conv3(x))
+        x = self.bn3(x)
+
+        x = self.pool(x).squeeze(-1)
+        return self.fc(x)
 
 # =========================================================
 # PATHS
 # =========================================================
 
 AD_MODEL_PATH = "AD_eeg_cnn_ad_ftd_cn.pt"
-PD_MODEL_PATH = "PD_eeg_cnn.pth"
+PD_MODEL_PATH = "PD_CNN_FINAL.pth"
 
 # =========================================================
 # LOAD MODELS
@@ -108,7 +117,7 @@ ad_model.eval()
 print("AD model loaded")
 
 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_MODEL_PATH, map_location=DEVICE))
 pd_model.eval()
 print("PD model loaded")
@@ -120,13 +129,14 @@ print("System ready")
 # =========================================================
 
 def prepare(features, model_type):
+
     x = torch.tensor(features, dtype=torch.float32)
 
     if model_type == "ad":
         return x.view(1, 19, 76)
 
     if model_type == "pd":
-        return x.view(1, 76)
+        return x.view(1, 64, 256)
 
     raise ValueError("Invalid model type")
 
@@ -161,8 +171,8 @@ def home():
     return {
         "status": "running",
         "ad_input": "19x76 CNN",
-        "pd_input": "76 features MLP",
-        "version": "9.0"
+        "pd_input": "64x256 CNN",
+        "version": "10.0"
     }
 
 @app.get("/health")