Update app.py

app.py

@@ -4,7 +4,7 @@ import torch
 import torch.nn as nn
 import numpy as np
 
-app = FastAPI(title="NeuroHealth EEG API", version="RESET-2")
+app = FastAPI(title="NeuroHealth EEG API", version="DEMO-AD-ONLY")
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -12,7 +12,6 @@ DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 # LABELS
 # ======================
 AD_CLASSES = ["Alzheimer", "FTD", "Control"]
-PD_CLASSES = ["Parkinson", "Control"]
 
 # ======================
 # INPUT MODEL
@@ -28,35 +27,15 @@ class EEG_CNN_AD(nn.Module):
         super().__init__()
         self.conv1 = nn.Conv1d(19, 32, 7, padding=3)
         self.bn1 = nn.BatchNorm1d(32)
-        self.conv2 = nn.Conv1d(32, 64, 5, padding=2)
-        self.bn2 = nn.BatchNorm1d(64)
-        self.conv3 = nn.Conv1d(64, 128, 3, padding=1)
-        self.bn3 = nn.BatchNorm1d(128)
-        self.pool = nn.AdaptiveAvgPool1d(1)
-        self.fc = nn.Linear(128, 3)
 
-    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 = self.pool(x).squeeze(-1)
-        return self.fc(x)
-
-# ======================
-# PD MODEL - 1D CNN (Fixed to match checkpoint)
-# ======================
-class EEG_CNN_PD(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.conv1 = nn.Conv1d(19, 32, 7, padding=3)
-        self.bn1 = nn.BatchNorm1d(32)
         self.conv2 = nn.Conv1d(32, 64, 5, padding=2)
         self.bn2 = nn.BatchNorm1d(64)
+
         self.conv3 = nn.Conv1d(64, 128, 3, padding=1)
         self.bn3 = nn.BatchNorm1d(128)
+
         self.pool = nn.AdaptiveAvgPool1d(1)
-        self.fc = nn.Linear(128, 2)
+        self.fc = nn.Linear(128, 3)
 
     def forward(self, x):
         x = x.view(x.size(0), 19, 76)
@@ -67,36 +46,32 @@ class EEG_CNN_PD(nn.Module):
         return self.fc(x)
 
 # ======================
-# LOAD MODELS
+# LOAD MODEL
 # ======================
 AD_PATH = "AD_eeg_cnn_ad_ftd_cn.pt"
-PD_PATH = "PD_CNN_FINAL.pth"
 
 print("Loading AD model...")
 ad_model = EEG_CNN_AD().to(DEVICE)
-ad_model.load_state_dict(torch.load(AD_PATH, map_location=DEVICE))
-ad_model.eval()
-print("AD loaded successfully")
 
-print("Loading PD model...")
-pd_model = EEG_CNN_PD().to(DEVICE)
-pd_model.load_state_dict(torch.load(PD_PATH, map_location=DEVICE))
-pd_model.eval()
-print("PD loaded successfully")
+ad_model.load_state_dict(
+    torch.load(AD_PATH, map_location=DEVICE)
+)
+
+ad_model.eval()
+print("AD model loaded successfully")
 
 # ======================
 # PREPROCESS
 # ======================
-def preprocess(features, mode):
+def preprocess(features):
     x = torch.tensor(features, dtype=torch.float32).to(DEVICE)
-    if mode == "ad" or mode == "pd":
-        # Ensure shape is (batch, channels, time) -> (1, 19, 76)
-        if len(x.shape) == 1:
-            x = x.view(1, 19, 76)
-        elif len(x.shape) == 2:
-            x = x.unsqueeze(0)
-        return x
-    raise ValueError("Unknown mode")
+
+    if len(x.shape) == 1:
+        x = x.view(1, 19, 76)
+    elif len(x.shape) == 2:
+        x = x.unsqueeze(0)
+
+    return x
 
 # ======================
 # PREDICT FUNCTION
@@ -106,10 +81,13 @@ def predict(model, x, classes):
         out = model(x)
         probs = torch.softmax(out, dim=1).cpu().numpy()[0]
         pred = int(np.argmax(probs))
+
         return {
             "prediction": classes[pred],
             "confidence": float(probs[pred]),
-            "probabilities": {classes[i]: float(probs[i]) for i in range(len(classes))}
+            "probabilities": {
+                classes[i]: float(probs[i]) for i in range(len(classes))
+            }
         }
 
 # ======================
@@ -119,16 +97,11 @@ def predict(model, x, classes):
 def home():
     return {
         "status": "ready",
-        "message": "NeuroHealth EEG API is running (Fixed CNN architecture)",
+        "message": "NeuroHealth EEG API - Alzheimer Demo Mode",
         "year": 2026
     }
 
 @app.post("/predict/ad")
 def predict_ad(req: EEGRequest):
-    x = preprocess(req.features, "ad")
-    return predict(ad_model, x, AD_CLASSES)
-
-@app.post("/predict/pd")
-def predict_pd(req: EEGRequest):
-    x = preprocess(req.features, "pd")
-    return predict(pd_model, x, PD_CLASSES)
+    x = preprocess(req.features)
+    return predict(ad_model, x, AD_CLASSES)