Update app.py

app.py

@@ -10,8 +10,8 @@ import numpy as np
 # =========================================================
 
 app = FastAPI(
-    title="NeuroHealth EEG CNN API",
-    version="5.0"
+    title="NeuroHealth EEG API",
+    version="6.0"
 )
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -24,35 +24,33 @@ AD_CLASSES = ["Alzheimer", "FTD", "Control"]
 PD_CLASSES = ["Parkinson", "Control"]
 
 # =========================================================
-# REQUEST SCHEMA
+# INPUT
 # =========================================================
 
 class EEGRequest(BaseModel):
     features: list
 
 # =========================================================
-# CNN MODEL (MATCH YOUR CHECKPOINT)
+# AD MODEL (conv style)
 # =========================================================
 
-class EEG_CNN(nn.Module):
+class EEG_CNN_AD(nn.Module):
     def __init__(self, output_dim):
         super().__init__()
 
-        self.conv1 = nn.Conv1d(19, 32, kernel_size=7, padding=3)
+        self.conv1 = nn.Conv1d(19, 32, 7, padding=3)
         self.bn1 = nn.BatchNorm1d(32)
 
-        self.conv2 = nn.Conv1d(32, 64, kernel_size=5, padding=2)
+        self.conv2 = nn.Conv1d(32, 64, 5, padding=2)
         self.bn2 = nn.BatchNorm1d(64)
 
-        self.conv3 = nn.Conv1d(64, 128, kernel_size=3, padding=1)
+        self.conv3 = nn.Conv1d(64, 128, 3, padding=1)
         self.bn3 = nn.BatchNorm1d(128)
 
         self.pool = nn.AdaptiveAvgPool1d(1)
-
         self.fc = nn.Linear(128, output_dim)
 
     def forward(self, x):
-        # EXPECTED INPUT: (batch, 19, 76)
         x = x.view(x.size(0), 19, -1)
 
         x = torch.relu(self.bn1(self.conv1(x)))
@@ -60,54 +58,69 @@ class EEG_CNN(nn.Module):
         x = torch.relu(self.bn3(self.conv3(x)))
 
         x = self.pool(x).squeeze(-1)
-
         return self.fc(x)
 
 # =========================================================
-# CONFIG
+# PD MODEL (SEQUENTIAL net style - YOUR CHECKPOINT)
+# =========================================================
+
+class EEG_CNN_PD(nn.Module):
+    def __init__(self, input_dim, output_dim):
+        super().__init__()
+
+        self.net = nn.Sequential(
+            nn.Linear(input_dim, 256),   # net.0
+            nn.ReLU(),
+            nn.BatchNorm1d(256),         # net.2
+
+            nn.Linear(256, 128),         # net.4
+            nn.ReLU(),
+            nn.BatchNorm1d(128),         # net.6
+
+            nn.Linear(128, output_dim)   # net.7
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+# =========================================================
+# PATHS
 # =========================================================
 
 AD_MODEL_PATH = "AD_eeg_cnn_ad_ftd_cn.pt"
 PD_MODEL_PATH = "PD_eeg_cnn.pth"
 
+INPUT_DIM = 19 * 76
+
 # =========================================================
 # LOAD MODELS
 # =========================================================
 
-print("Loading AD CNN model...")
-
-ad_model = EEG_CNN(3).to(DEVICE)
+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 CNN model...")
-
-pd_model = EEG_CNN(2).to(DEVICE)
+print("Loading PD model...")
+pd_model = EEG_CNN_PD(INPUT_DIM, 2).to(DEVICE)
 pd_model.load_state_dict(torch.load(PD_MODEL_PATH, map_location=DEVICE))
 pd_model.eval()
-
-print("PD model loaded")
-
-print("All models ready")
+print("PD loaded")
 
 # =========================================================
-# PREDICTION FUNCTION
+# PREDICT
 # =========================================================
 
 def predict(model, features, classes):
 
     x = torch.tensor(features, dtype=torch.float32)
 
-    if x.numel() != 19 * 76:
-        raise ValueError("Expected 19x76 = 1444 features")
-
-    x = x.view(1, 19, 76).to(DEVICE)
+    x = x.view(1, -1).to(DEVICE)
 
     with torch.no_grad():
-        logits = model(x)
-        probs = torch.softmax(logits, dim=1).cpu().numpy()[0]
+        out = model(x)
+        probs = torch.softmax(out, dim=1).cpu().numpy()[0]
 
     pred = int(np.argmax(probs))
 
@@ -125,14 +138,7 @@ def predict(model, features, classes):
 
 @app.get("/")
 def home():
-    return {
-        "status": "NeuroHealth EEG CNN API running",
-        "input_shape": "19 x 76"
-    }
-
-@app.get("/health")
-def health():
-    return {"status": "ok"}
+    return {"status": "running"}
 
 @app.post("/predict/ad")
 def predict_ad(req: EEGRequest):