Update app.py

app.py

@@ -1,43 +1,31 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
-
 import torch
 import torch.nn as nn
 import numpy as np
 
-# =========================================================
-# APP
-# =========================================================
-
-app = FastAPI(
-    title="NeuroHealth EEG API",
-    version="10.0"
-)
+app = FastAPI(title="NeuroHealth EEG API", version="RESET-1")
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# =========================================================
+# ======================
 # LABELS
-# =========================================================
-
+# ======================
 AD_CLASSES = ["Alzheimer", "FTD", "Control"]
 PD_CLASSES = ["Parkinson", "Control"]
 
-# =========================================================
+# ======================
 # INPUT
-# =========================================================
-
+# ======================
 class EEGRequest(BaseModel):
     features: list
 
-# =========================================================
-# AD MODEL (UNCHANGED - 19 x 76 CNN)
-# =========================================================
-
+# ======================
+# AD MODEL (UNCHANGED CNN)
+# ======================
 class EEG_CNN_AD(nn.Module):
     def __init__(self):
         super().__init__()
-
         self.conv1 = nn.Conv1d(19, 32, 7, padding=3)
         self.bn1 = nn.BatchNorm1d(32)
 
@@ -53,101 +41,71 @@ class EEG_CNN_AD(nn.Module):
     def forward(self, x):
         x = x.view(x.size(0), 19, 76)
 
-        x = torch.relu(self.conv1(x))
-        x = self.bn1(x)
-
-        x = torch.relu(self.conv2(x))
-        x = self.bn2(x)
-
-        x = torch.relu(self.conv3(x))
-        x = self.bn3(x)
+        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 (UPDATED CNN - MATCHES YOUR TRAINED MODEL)
-# =========================================================
-
-class EEG_CNN_PD(nn.Module):
+# ======================
+# PD MODEL (MATCH CHECKPOINT EXACTLY)
+# ======================
+class EEG_PD(nn.Module):
     def __init__(self):
         super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(76, 256),
+            nn.ReLU(),
+            nn.BatchNorm1d(256),
 
-        self.conv1 = nn.Conv1d(64, 32, 7, padding=3)
-        self.bn1 = nn.BatchNorm1d(32)
-
-        self.conv2 = nn.Conv1d(32, 64, 5, padding=2)
-        self.bn2 = nn.BatchNorm1d(64)
+            nn.Linear(256, 128),
+            nn.ReLU(),
+            nn.BatchNorm1d(128),
 
-        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)
+            nn.Linear(128, 2)
+        )
 
     def forward(self, x):
-        # x: (batch, 64, 256)
-        x = torch.relu(self.conv1(x))
-        x = self.bn1(x)
+        return self.net(x)
 
-        x = torch.relu(self.conv2(x))
-        x = self.bn2(x)
-
-        x = torch.relu(self.conv3(x))
-        x = self.bn3(x)
-
-        x = self.pool(x).squeeze(-1)
-        return self.fc(x)
-
-# =========================================================
-# PATHS
-# =========================================================
-
-AD_MODEL_PATH = "AD_eeg_cnn_ad_ftd_cn.pt"
-PD_MODEL_PATH = "PD_CNN_FINAL.pth"
-
-# =========================================================
+# ======================
 # LOAD MODELS
-# =========================================================
+# ======================
+AD_PATH = "AD_eeg_model.pth"
+PD_PATH = "PD_eeg_model.pth"
 
 print("Loading AD model...")
 ad_model = EEG_CNN_AD().to(DEVICE)
-ad_model.load_state_dict(torch.load(AD_MODEL_PATH, map_location=DEVICE))
+ad_model.load_state_dict(torch.load(AD_PATH, map_location=DEVICE))
 ad_model.eval()
-print("AD model loaded")
+print("AD loaded")
 
 print("Loading PD model...")
-pd_model = EEG_CNN_PD().to(DEVICE)
-pd_model.load_state_dict(torch.load(PD_MODEL_PATH, map_location=DEVICE))
+pd_model = EEG_PD().to(DEVICE)
+pd_model.load_state_dict(torch.load(PD_PATH, map_location=DEVICE))
 pd_model.eval()
-print("PD model loaded")
-
-print("System ready")
-
-# =========================================================
-# PREPROCESSING
-# =========================================================
-
-def prepare(features, model_type):
+print("PD loaded")
 
+# ======================
+# PREPROCESS
+# ======================
+def preprocess(features, mode):
     x = torch.tensor(features, dtype=torch.float32)
 
-    if model_type == "ad":
+    if mode == "ad":
         return x.view(1, 19, 76)
 
-    if model_type == "pd":
-        return x.view(1, 64, 256)
-
-    raise ValueError("Invalid model type")
-
-# =========================================================
-# PREDICTION ENGINE
-# =========================================================
-
-def predict(model, features, classes, model_type):
+    if mode == "pd":
+        # IMPORTANT: PD expects FLATTENED 76-D FEATURES
+        return x.view(1, 76)
 
-    x = prepare(features, model_type).to(DEVICE)
+    raise ValueError("unknown mode")
 
+# ======================
+# PREDICT ENGINE
+# ======================
+def predict(model, x, classes):
     with torch.no_grad():
         out = model(x)
         probs = torch.softmax(out, dim=1).cpu().numpy()[0]
@@ -162,27 +120,19 @@ def predict(model, features, classes, model_type):
         }
     }
 
-# =========================================================
+# ======================
 # ROUTES
-# =========================================================
-
+# ======================
 @app.get("/")
 def home():
-    return {
-        "status": "running",
-        "ad_input": "19x76 CNN",
-        "pd_input": "64x256 CNN",
-        "version": "10.0"
-    }
-
-@app.get("/health")
-def health():
-    return {"status": "ok"}
+    return {"status": "ready (RESET architecture)", "year": 2026}
 
 @app.post("/predict/ad")
 def predict_ad(req: EEGRequest):
-    return predict(ad_model, req.features, AD_CLASSES, "ad")
+    x = preprocess(req.features, "ad").to(DEVICE)
+    return predict(ad_model, x, AD_CLASSES)
 
 @app.post("/predict/pd")
 def predict_pd(req: EEGRequest):
-    return predict(pd_model, req.features, PD_CLASSES, "pd")
+    x = preprocess(req.features, "pd").to(DEVICE)
+    return predict(pd_model, x, PD_CLASSES)