Update app.py

app.py

@@ -11,7 +11,7 @@ import numpy as np
 
 app = FastAPI(
     title="NeuroHealth EEG System",
-    version="7.1"
+    version="8.0"
 )
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -24,18 +24,18 @@ AD_CLASSES = ["Alzheimer", "FTD", "Control"]
 PD_CLASSES = ["Parkinson", "Control"]
 
 # =========================================================
-# INPUT SCHEMA
+# INPUT
 # =========================================================
 
 class EEGRequest(BaseModel):
     features: list
 
 # =========================================================
-# AD MODEL (CNN - REAL ARCHITECTURE)
+# AD MODEL (CNN - MUST MATCH TRAINING)
 # =========================================================
 
 class EEG_CNN_AD(nn.Module):
-    def __init__(self, output_dim):
+    def __init__(self):
         super().__init__()
 
         self.conv1 = nn.Conv1d(19, 32, 7, padding=3)
@@ -48,7 +48,7 @@ class EEG_CNN_AD(nn.Module):
         self.bn3 = nn.BatchNorm1d(128)
 
         self.pool = nn.AdaptiveAvgPool1d(1)
-        self.fc = nn.Linear(128, output_dim)
+        self.fc = nn.Linear(128, 3)
 
     def forward(self, x):
         x = x.view(x.size(0), 19, 76)
@@ -61,30 +61,30 @@ class EEG_CNN_AD(nn.Module):
         return self.fc(x)
 
 # =========================================================
-# PD MODEL (FULLY CONNECTED NETWORK - EXACT MATCH)
+# PD MODEL (EXACT MATCH TO CHECKPOINT)
 # =========================================================
 
 class EEG_PD(nn.Module):
-    def __init__(self, output_dim):
+    def __init__(self):
         super().__init__()
 
         self.net = nn.Sequential(
-            nn.Linear(76, 256),     # net.0
+            nn.Linear(76, 256),
             nn.ReLU(),
-            nn.BatchNorm1d(256),    # net.2
+            nn.BatchNorm1d(256),
 
-            nn.Linear(256, 128),    # net.4
+            nn.Linear(256, 128),
             nn.ReLU(),
-            nn.BatchNorm1d(128),    # net.6
+            nn.BatchNorm1d(128),
 
-            nn.Linear(128, output_dim)  # net.7
+            nn.Linear(128, 2)
         )
 
     def forward(self, x):
         return self.net(x)
 
 # =========================================================
-# MODEL PATHS
+# PATHS
 # =========================================================
 
 AD_MODEL_PATH = "AD_eeg_cnn_ad_ftd_cn.pt"
@@ -95,13 +95,13 @@ PD_MODEL_PATH = "PD_eeg_cnn.pth"
 # =========================================================
 
 print("Loading AD model...")
-ad_model = EEG_CNN_AD(3).to(DEVICE)
+ad_model = EEG_CNN_AD().to(DEVICE)
 ad_model.load_state_dict(torch.load(AD_MODEL_PATH, map_location=DEVICE))
 ad_model.eval()
 print("AD model loaded")
 
 print("Loading PD model...")
-pd_model = EEG_PD(2).to(DEVICE)
+pd_model = EEG_PD().to(DEVICE)
 pd_model.load_state_dict(torch.load(PD_MODEL_PATH, map_location=DEVICE))
 pd_model.eval()
 print("PD model loaded")
@@ -109,35 +109,33 @@ print("PD model loaded")
 print("System ready")
 
 # =========================================================
-# INPUT PREPARATION
+# INPUT HANDLER
 # =========================================================
 
-def prepare_input(features, model_type):
+def prepare(features, model_type):
 
     x = torch.tensor(features, dtype=torch.float32)
 
     if model_type == "ad":
-        # CNN expects 19 x 76
         return x.view(1, 19, 76)
 
     elif model_type == "pd":
-        # Dense expects 76 only
         return x.view(1, 76)
 
     else:
-        raise ValueError("Unknown model type")
+        raise ValueError("Invalid model type")
 
 # =========================================================
-# PREDICTION ENGINE
+# PREDICT ENGINE
 # =========================================================
 
 def predict(model, features, classes, model_type):
 
-    x = prepare_input(features, model_type).to(DEVICE)
+    x = prepare(features, model_type).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))
 
@@ -157,8 +155,9 @@ def predict(model, features, classes, model_type):
 def home():
     return {
         "status": "running",
-        "ad_model": "CNN (19x76)",
-        "pd_model": "Dense (76)"
+        "ad_input": "19x76 CNN",
+        "pd_input": "76 Dense",
+        "version": "8.0"
     }
 
 @app.get("/health")