Update app.py

app.py

@@ -10,8 +10,8 @@ import numpy as np
 # =========================================================
 
 app = FastAPI(
-    title="NeuroHealth EEG API",
-    version="6.0"
+    title="NeuroHealth EEG System",
+    version="7.0"
 )
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -24,14 +24,14 @@ AD_CLASSES = ["Alzheimer", "FTD", "Control"]
 PD_CLASSES = ["Parkinson", "Control"]
 
 # =========================================================
-# INPUT
+# INPUT SCHEMA
 # =========================================================
 
 class EEGRequest(BaseModel):
     features: list
 
 # =========================================================
-# AD MODEL (conv style)
+# AD MODEL (CNN - 19 channels)
 # =========================================================
 
 class EEG_CNN_AD(nn.Module):
@@ -51,7 +51,7 @@ class EEG_CNN_AD(nn.Module):
         self.fc = nn.Linear(128, output_dim)
 
     def forward(self, x):
-        x = x.view(x.size(0), 19, -1)
+        x = x.view(x.size(0), 19, 76)
 
         x = torch.relu(self.bn1(self.conv1(x)))
         x = torch.relu(self.bn2(self.conv2(x)))
@@ -61,23 +61,23 @@ class EEG_CNN_AD(nn.Module):
         return self.fc(x)
 
 # =========================================================
-# PD MODEL (SEQUENTIAL net style - YOUR CHECKPOINT)
+# PD MODEL (Dense Sequential)
 # =========================================================
 
-class EEG_CNN_PD(nn.Module):
-    def __init__(self, input_dim, output_dim):
+class EEG_PD(nn.Module):
+    def __init__(self, output_dim):
         super().__init__()
 
         self.net = nn.Sequential(
-            nn.Linear(input_dim, 256),   # net.0
+            nn.Linear(76, 256),
             nn.ReLU(),
-            nn.BatchNorm1d(256),         # net.2
+            nn.BatchNorm1d(256),
 
-            nn.Linear(256, 128),         # net.4
+            nn.Linear(256, 128),
             nn.ReLU(),
-            nn.BatchNorm1d(128),         # net.6
+            nn.BatchNorm1d(128),
 
-            nn.Linear(128, output_dim)   # net.7
+            nn.Linear(128, output_dim)
         )
 
     def forward(self, x):
@@ -90,8 +90,6 @@ class EEG_CNN_PD(nn.Module):
 AD_MODEL_PATH = "AD_eeg_cnn_ad_ftd_cn.pt"
 PD_MODEL_PATH = "PD_eeg_cnn.pth"
 
-INPUT_DIM = 19 * 76
-
 # =========================================================
 # LOAD MODELS
 # =========================================================
@@ -100,27 +98,33 @@ print("Loading AD model...")
 ad_model = EEG_CNN_AD(3).to(DEVICE)
 ad_model.load_state_dict(torch.load(AD_MODEL_PATH, map_location=DEVICE))
 ad_model.eval()
-print("AD loaded")
+print("AD model loaded")
 
 print("Loading PD model...")
-pd_model = EEG_CNN_PD(INPUT_DIM, 2).to(DEVICE)
+pd_model = EEG_PD(2).to(DEVICE)
 pd_model.load_state_dict(torch.load(PD_MODEL_PATH, map_location=DEVICE))
 pd_model.eval()
-print("PD loaded")
+print("PD model loaded")
+
+print("System ready")
 
 # =========================================================
-# PREDICT
+# PREDICTION ENGINE
 # =========================================================
 
 def predict(model, features, classes):
 
     x = torch.tensor(features, dtype=torch.float32)
 
-    x = x.view(1, -1).to(DEVICE)
+    # AD expects 1444 (19x76), PD expects 76
+    if len(features) == 1444:
+        x = x.view(1, 19, 76)
+    else:
+        x = x.view(1, 76)
 
     with torch.no_grad():
-        out = model(x)
-        probs = torch.softmax(out, dim=1).cpu().numpy()[0]
+        logits = model(x)
+        probs = torch.softmax(logits, dim=1).cpu().numpy()[0]
 
     pred = int(np.argmax(probs))
 
@@ -138,7 +142,11 @@ def predict(model, features, classes):
 
 @app.get("/")
 def home():
-    return {"status": "running"}
+    return {
+        "status": "running",
+        "ad_input": "19x76 CNN",
+        "pd_input": "76 Dense"
+    }
 
 @app.post("/predict/ad")
 def predict_ad(req: EEGRequest):
@@ -146,4 +154,8 @@ def predict_ad(req: EEGRequest):
 
 @app.post("/predict/pd")
 def predict_pd(req: EEGRequest):
-    return predict(pd_model, req.features, PD_CLASSES)
+    return predict(pd_model, req.features, PD_CLASSES)
+
+@app.get("/health")
+def health():
+    return {"status": "ok"}