Update app.py

app.py

@@ -31,14 +31,14 @@ class EEGRequest(BaseModel):
     features: list
 
 # =========================================================
-# CNN MODEL (MATCHES YOUR CHECKPOINT)
+# CNN MODEL (MATCH YOUR CHECKPOINT)
 # =========================================================
 
 class EEG_CNN(nn.Module):
-    def __init__(self, input_dim, output_dim):
+    def __init__(self, output_dim):
         super().__init__()
 
-        self.conv1 = nn.Conv1d(1, 32, kernel_size=7, padding=3)
+        self.conv1 = nn.Conv1d(19, 32, kernel_size=7, padding=3)
         self.bn1 = nn.BatchNorm1d(32)
 
         self.conv2 = nn.Conv1d(32, 64, kernel_size=5, padding=2)
@@ -52,7 +52,8 @@ class EEG_CNN(nn.Module):
         self.fc = nn.Linear(128, output_dim)
 
     def forward(self, x):
-        x = x.unsqueeze(1)
+        # EXPECTED INPUT: (batch, 19, 76)
+        x = x.view(x.size(0), 19, -1)
 
         x = torch.relu(self.bn1(self.conv1(x)))
         x = torch.relu(self.bn2(self.conv2(x)))
@@ -66,8 +67,6 @@ class EEG_CNN(nn.Module):
 # CONFIG
 # =========================================================
 
-INPUT_DIM = 76
-
 AD_MODEL_PATH = "AD_eeg_cnn_ad_ftd_cn.pt"
 PD_MODEL_PATH = "PD_eeg_cnn.pth"
 
@@ -77,7 +76,7 @@ PD_MODEL_PATH = "PD_eeg_cnn.pth"
 
 print("Loading AD CNN model...")
 
-ad_model = EEG_CNN(INPUT_DIM, 3).to(DEVICE)
+ad_model = EEG_CNN(3).to(DEVICE)
 ad_model.load_state_dict(torch.load(AD_MODEL_PATH, map_location=DEVICE))
 ad_model.eval()
 
@@ -85,7 +84,7 @@ print("AD model loaded")
 
 print("Loading PD CNN model...")
 
-pd_model = EEG_CNN(INPUT_DIM, 2).to(DEVICE)
+pd_model = EEG_CNN(2).to(DEVICE)
 pd_model.load_state_dict(torch.load(PD_MODEL_PATH, map_location=DEVICE))
 pd_model.eval()
 
@@ -94,17 +93,17 @@ print("PD model loaded")
 print("All models ready")
 
 # =========================================================
-# PREDICTION ENGINE
+# PREDICTION FUNCTION
 # =========================================================
 
 def predict(model, features, classes):
 
     x = torch.tensor(features, dtype=torch.float32)
 
-    if x.numel() != INPUT_DIM:
-        raise ValueError(f"Expected {INPUT_DIM} features, got {x.numel()}")
+    if x.numel() != 19 * 76:
+        raise ValueError("Expected 19x76 = 1444 features")
 
-    x = x.unsqueeze(0).to(DEVICE)
+    x = x.view(1, 19, 76).to(DEVICE)
 
     with torch.no_grad():
         logits = model(x)
@@ -128,7 +127,7 @@ def predict(model, features, classes):
 def home():
     return {
         "status": "NeuroHealth EEG CNN API running",
-        "input_dim": INPUT_DIM
+        "input_shape": "19 x 76"
     }
 
 @app.get("/health")