Update app.py

app.py

@@ -10,8 +10,8 @@ import numpy as np
 # =========================================================
 
 app = FastAPI(
-    title="NeuroHealth EEG API",
-    version="4.0"
+    title="NeuroHealth EEG CNN API",
+    version="5.0"
 )
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -20,51 +20,50 @@ DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 # LABELS
 # =========================================================
 
-AD_CLASSES = [
-    "Alzheimer",
-    "FTD",
-    "Control"
-]
-
-PD_CLASSES = [
-    "Parkinson",
-    "Control"
-]
+AD_CLASSES = ["Alzheimer", "FTD", "Control"]
+PD_CLASSES = ["Parkinson", "Control"]
 
 # =========================================================
-# REQUEST MODEL
+# REQUEST SCHEMA
 # =========================================================
 
 class EEGRequest(BaseModel):
     features: list
 
 # =========================================================
-# MODEL ARCHITECTURE
+# CNN MODEL (MATCHES YOUR CHECKPOINT)
 # =========================================================
 
-class EEG_MLP(nn.Module):
+class EEG_CNN(nn.Module):
     def __init__(self, input_dim, output_dim):
         super().__init__()
 
-        self.net = nn.Sequential(
-            nn.Linear(input_dim, 256),
-            nn.ReLU(),
+        self.conv1 = nn.Conv1d(1, 16, kernel_size=3, padding=1)
+        self.bn1 = nn.BatchNorm1d(16)
 
-            nn.BatchNorm1d(256),
+        self.conv2 = nn.Conv1d(16, 32, kernel_size=3, padding=1)
+        self.bn2 = nn.BatchNorm1d(32)
 
-            nn.Linear(256, 128),
-            nn.ReLU(),
+        self.conv3 = nn.Conv1d(32, 64, kernel_size=3, padding=1)
+        self.bn3 = nn.BatchNorm1d(64)
 
-            nn.BatchNorm1d(128),
+        self.pool = nn.AdaptiveAvgPool1d(1)
 
-            nn.Linear(128, output_dim)
-        )
+        self.fc = nn.Linear(64, output_dim)
 
     def forward(self, x):
-        return self.net(x)
+        x = x.unsqueeze(1)  # (batch, 1, 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)
 
 # =========================================================
-# SETTINGS
+# CONFIG
 # =========================================================
 
 INPUT_DIM = 76
@@ -76,69 +75,48 @@ PD_MODEL_PATH = "PD_eeg_cnn.pth"
 # LOAD MODELS
 # =========================================================
 
-print("Loading AD model...")
-
-ad_model = EEG_MLP(INPUT_DIM, 3).to(DEVICE)
-
-ad_model.load_state_dict(
-    torch.load(
-        AD_MODEL_PATH,
-        map_location=DEVICE
-    )
-)
+print("Loading AD CNN model...")
 
+ad_model = EEG_CNN(INPUT_DIM, 3).to(DEVICE)
+ad_model.load_state_dict(torch.load(AD_MODEL_PATH, map_location=DEVICE))
 ad_model.eval()
 
-print("AD model loaded successfully")
-
-print("Loading PD model...")
-
-pd_model = EEG_MLP(INPUT_DIM, 2).to(DEVICE)
+print("AD model loaded")
 
-pd_model.load_state_dict(
-    torch.load(
-        PD_MODEL_PATH,
-        map_location=DEVICE
-    )
-)
+print("Loading PD CNN model...")
 
+pd_model = EEG_CNN(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 successfully")
+print("PD model loaded")
 
-print("All models loaded")
+print("All models ready")
 
 # =========================================================
-# PREDICTION FUNCTION
+# PREDICTION ENGINE
 # =========================================================
 
-def predict(model, features, class_names):
+def predict(model, features, classes):
 
-    x = torch.tensor(
-        features,
-        dtype=torch.float32
-    )
+    x = torch.tensor(features, dtype=torch.float32)
 
     if x.numel() != INPUT_DIM:
-        raise ValueError(
-            f"Expected {INPUT_DIM} features but received {x.numel()}"
-        )
+        raise ValueError(f"Expected {INPUT_DIM} features, got {x.numel()}")
 
     x = x.unsqueeze(0).to(DEVICE)
 
     with torch.no_grad():
-        outputs = model(x)
-        probs = torch.softmax(outputs, dim=1)
-        probs = probs.cpu().numpy()[0]
+        logits = model(x)
+        probs = torch.softmax(logits, dim=1).cpu().numpy()[0]
 
-    pred_idx = int(np.argmax(probs))
+    pred = int(np.argmax(probs))
 
     return {
-        "prediction": class_names[pred_idx],
-        "confidence": float(probs[pred_idx]),
+        "prediction": classes[pred],
+        "confidence": float(probs[pred]),
         "probabilities": {
-            class_names[i]: float(probs[i])
-            for i in range(len(class_names))
+            classes[i]: float(probs[i]) for i in range(len(classes))
         }
     }
 
@@ -149,31 +127,18 @@ def predict(model, features, class_names):
 @app.get("/")
 def home():
     return {
-        "status": "running",
-        "device": str(DEVICE),
-        "input_features": INPUT_DIM,
-        "ad_classes": AD_CLASSES,
-        "pd_classes": PD_CLASSES
+        "status": "NeuroHealth EEG CNN API running",
+        "input_dim": INPUT_DIM
     }
 
 @app.get("/health")
 def health():
-    return {
-        "status": "healthy"
-    }
+    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)
 
 @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)