Update app.py

app.py

@@ -11,7 +11,7 @@ import numpy as np
 
 app = FastAPI(
     title="NeuroHealth EEG System",
-    version="7.0"
+    version="7.1"
 )
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -31,7 +31,7 @@ class EEGRequest(BaseModel):
     features: list
 
 # =========================================================
-# AD MODEL (CNN - 19 channels)
+# AD MODEL (CNN - REAL ARCHITECTURE)
 # =========================================================
 
 class EEG_CNN_AD(nn.Module):
@@ -61,7 +61,7 @@ class EEG_CNN_AD(nn.Module):
         return self.fc(x)
 
 # =========================================================
-# PD MODEL (Dense Sequential)
+# PD MODEL (FULLY CONNECTED NETWORK - EXACT MATCH)
 # =========================================================
 
 class EEG_PD(nn.Module):
@@ -69,22 +69,22 @@ 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.BatchNorm1d(256),    # net.2
 
-            nn.Linear(256, 128),
+            nn.Linear(256, 128),    # net.4
             nn.ReLU(),
-            nn.BatchNorm1d(128),
+            nn.BatchNorm1d(128),    # net.6
 
-            nn.Linear(128, output_dim)
+            nn.Linear(128, output_dim)  # net.7
         )
 
     def forward(self, x):
         return self.net(x)
 
 # =========================================================
-# PATHS
+# MODEL PATHS
 # =========================================================
 
 AD_MODEL_PATH = "AD_eeg_cnn_ad_ftd_cn.pt"
@@ -109,18 +109,31 @@ print("PD model loaded")
 print("System ready")
 
 # =========================================================
-# PREDICTION ENGINE
+# INPUT PREPARATION
 # =========================================================
 
-def predict(model, features, classes):
+def prepare_input(features, model_type):
 
     x = torch.tensor(features, dtype=torch.float32)
 
-    # AD expects 1444 (19x76), PD expects 76
-    if len(features) == 1444:
-        x = x.view(1, 19, 76)
+    if model_type == "ad":
+        # CNN expects 19 x 76
+        return x.view(1, 19, 76)
+
+    elif model_type == "pd":
+        # Dense expects 76 only
+        return x.view(1, 76)
+
     else:
-        x = x.view(1, 76)
+        raise ValueError("Unknown model type")
+
+# =========================================================
+# PREDICTION ENGINE
+# =========================================================
+
+def predict(model, features, classes, model_type):
+
+    x = prepare_input(features, model_type).to(DEVICE)
 
     with torch.no_grad():
         logits = model(x)
@@ -144,18 +157,18 @@ def predict(model, features, classes):
 def home():
     return {
         "status": "running",
-        "ad_input": "19x76 CNN",
-        "pd_input": "76 Dense"
+        "ad_model": "CNN (19x76)",
+        "pd_model": "Dense (76)"
     }
 
+@app.get("/health")
+def health():
+    return {"status": "ok"}
+
 @app.post("/predict/ad")
 def predict_ad(req: EEGRequest):
-    return predict(ad_model, req.features, AD_CLASSES)
+    return predict(ad_model, req.features, AD_CLASSES, "ad")
 
 @app.post("/predict/pd")
 def predict_pd(req: EEGRequest):
-    return predict(pd_model, req.features, PD_CLASSES)
-
-@app.get("/health")
-def health():
-    return {"status": "ok"}
+    return predict(pd_model, req.features, PD_CLASSES, "pd")