Update app.py

app.py

@@ -10,8 +10,8 @@ import numpy as np
 # =========================================================
 
 app = FastAPI(
-    title="NeuroHealth EEG System",
-    version="8.0"
+    title="NeuroHealth EEG API",
+    version="9.0"
 )
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -31,7 +31,7 @@ class EEGRequest(BaseModel):
     features: list
 
 # =========================================================
-# AD MODEL (CNN - MUST MATCH TRAINING)
+# AD MODEL (CNN - EXACTLY FROM CHECKPOINT STRUCTURE)
 # =========================================================
 
 class EEG_CNN_AD(nn.Module):
@@ -53,15 +53,20 @@ class EEG_CNN_AD(nn.Module):
     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 = 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)
 
 # =========================================================
-# PD MODEL (EXACT MATCH TO CHECKPOINT)
+# PD MODEL (STRICT MATCH TO YOUR CHECKPOINT KEYS)
 # =========================================================
 
 class EEG_PD(nn.Module):
@@ -69,15 +74,17 @@ class EEG_PD(nn.Module):
         super().__init__()
 
         self.net = nn.Sequential(
-            nn.Linear(76, 256),
+            nn.Linear(76, 256),   # net.0
             nn.ReLU(),
-            nn.BatchNorm1d(256),
 
-            nn.Linear(256, 128),
+            nn.BatchNorm1d(256),  # net.2
+
+            nn.Linear(256, 128),  # net.4
             nn.ReLU(),
-            nn.BatchNorm1d(128),
 
-            nn.Linear(128, 2)
+            nn.BatchNorm1d(128),  # net.6
+
+            nn.Linear(128, 2)     # net.7
         )
 
     def forward(self, x):
@@ -109,24 +116,22 @@ print("PD model loaded")
 print("System ready")
 
 # =========================================================
-# INPUT HANDLER
+# PREPROCESSING
 # =========================================================
 
 def prepare(features, model_type):
-
     x = torch.tensor(features, dtype=torch.float32)
 
     if model_type == "ad":
         return x.view(1, 19, 76)
 
-    elif model_type == "pd":
+    if model_type == "pd":
         return x.view(1, 76)
 
-    else:
-        raise ValueError("Invalid model type")
+    raise ValueError("Invalid model type")
 
 # =========================================================
-# PREDICT ENGINE
+# PREDICTION ENGINE
 # =========================================================
 
 def predict(model, features, classes, model_type):
@@ -156,8 +161,8 @@ def home():
     return {
         "status": "running",
         "ad_input": "19x76 CNN",
-        "pd_input": "76 Dense",
-        "version": "8.0"
+        "pd_input": "76 features MLP",
+        "version": "9.0"
     }
 
 @app.get("/health")