Update app.py

app.py

@@ -1,196 +1,193 @@
+from fastapi import FastAPI
+from pydantic import BaseModel
+
 import torch
 import torch.nn as nn
-import torchvision.models as models
-from fastapi import FastAPI, UploadFile, File
-from PIL import Image
-import io
-import torchvision.transforms as transforms
-
-# =========================
-# APP INIT
-# =========================
-app = FastAPI(title="Alzheimer Ensemble API")
-
-# =========================
-# DEVICE
-# =========================
-DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print("Using device:", DEVICE)
-
-# =========================
-# CLASS LABELS (FIXED)
-# =========================
-CLASSES = [
-    "Mild Demented",
-    "Moderate Demented",
-    "Non Demented",
-    "Very Mild Demented"
+import numpy as np
+
+app = FastAPI(
+    title="NeuroHealth EEG API",
+    version="1.0"
+)
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# =========================================================
+# CLASS NAMES
+# =========================================================
+
+AD_CLASSES = [
+    "Alzheimer",
+    "FTD",
+    "Control"
 ]
 
-# =========================
-# IMAGE TRANSFORM
-# =========================
-transform = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
-])
-
-# =========================
-# MODEL BUILDER
-# =========================
-def build_model(version="121"):
-    if version == "121":
-        model = models.densenet121(weights=None)
-        in_features = 1024
-    elif version == "169":
-        model = models.densenet169(weights=None)
-        in_features = 1664
-    else:
-        model = models.densenet201(weights=None)
-        in_features = 1920
-
-    model.classifier = nn.Sequential(
-        nn.Dropout(0.4),
-        nn.Linear(in_features, len(CLASSES))
-    )
+PD_CLASSES = [
+    "Parkinson",
+    "Control"
+]
+
+# =========================================================
+# MODELS
+# =========================================================
+
+INPUT_DIM = 95
+
+# =========================================================
+# MLP MODEL
+# =========================================================
+
+class MLP(nn.Module):
+
+    def __init__(self, output_dim):
+
+        super().__init__()
 
-    return model
+        self.network = nn.Sequential(
 
-# =========================
-# SAFE MODEL LOADER
-# =========================
-def load_model(path, version):
-    model = build_model(version)
+            nn.Linear(INPUT_DIM, 512),
+            nn.BatchNorm1d(512),
+            nn.GELU(),
+            nn.Dropout(0.3),
 
-    try:
-        checkpoint = torch.load(path, map_location=DEVICE)
+            nn.Linear(512, 256),
+            nn.BatchNorm1d(256),
+            nn.GELU(),
+            nn.Dropout(0.3),
 
-        if isinstance(checkpoint, dict):
-            if "state_dict" in checkpoint:
-                checkpoint = checkpoint["state_dict"]
-            elif "model_state_dict" in checkpoint:
-                checkpoint = checkpoint["model_state_dict"]
+            nn.Linear(256, 128),
+            nn.BatchNorm1d(128),
+            nn.GELU(),
 
-        model.load_state_dict(checkpoint, strict=False)
-        print(f"Loaded: {path}")
+            nn.Linear(128, output_dim)
+        )
 
-    except Exception as e:
-        print(f"Failed loading {path}: {e}")
+    def forward(self, x):
 
-    model.to(DEVICE)
-    model.eval()
-    return model
+        return self.network(x)
 
-# =========================
+# =========================================================
 # LOAD MODELS
-# =========================
-model_121 = load_model("alzheimers_densenet121.pth", "121")
-model_169 = load_model("alzheimers_densenet169.pth", "169")
-model_201 = load_model("alzheimers_densenet201.pth", "201")
-
-# =========================
-# IMAGE PROCESSING
-# =========================
-def process_image(image_bytes):
-    img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-    img = transform(img).unsqueeze(0).to(DEVICE)
-    return img
-
-# =========================
-# PREDICTION FUNCTION (FIXED)
-# =========================
-def predict(model, x):
+# =========================================================
+
+ad_mlp = MLP(3).to(DEVICE)
+pd_mlp = MLP(2).to(DEVICE)
+
+ad_mlp.load_state_dict(
+    torch.load(
+        "AD_MLP.pt",
+        map_location=DEVICE
+    )
+)
+
+pd_mlp.load_state_dict(
+    torch.load(
+        "PD_MLP.pt",
+        map_location=DEVICE
+    )
+)
+
+ad_mlp.eval()
+pd_mlp.eval()
+
+print("Models Loaded Successfully")
+
+# =========================================================
+# REQUEST MODEL
+# =========================================================
+
+class EEGRequest(BaseModel):
+
+    features: list
+
+# =========================================================
+# UTILITY
+# =========================================================
+
+def predict_model(model, features):
+
+    x = torch.tensor(
+        features,
+        dtype=torch.float32
+    ).unsqueeze(0).to(DEVICE)
+
     with torch.no_grad():
-        output = model(x)
-        probs = torch.softmax(output, dim=1)[0]
-
-        conf, cls = torch.max(probs, 0)
-        cls = int(cls.item())
-
-        return {
-            "prediction": CLASSES[cls],
-            "class_id": cls,
-            "confidence": float(conf.item()),
-            "probabilities": {
-                CLASSES[i]: float(probs[i].item())
-                for i in range(len(CLASSES))
-            }
-        }
 
-# =========================
-# ROOT ENDPOINT
-# =========================
+        outputs = model(x)
+
+        probs = torch.softmax(
+            outputs,
+            dim=1
+        ).cpu().numpy()[0]
+
+    pred = int(np.argmax(probs))
+
+    confidence = float(probs[pred])
+
+    return pred, confidence, probs.tolist()
+
+# =========================================================
+# ROOT
+# =========================================================
+
 @app.get("/")
+
 def home():
+
     return {
-        "status": "running",
-        "models": ["121", "169", "201"],
-        "classes": CLASSES,
-        "endpoints": [
-            "/predict/121",
-            "/predict/169",
-            "/predict/201",
-            "/predict/ensemble"
-        ]
+        "message": "NeuroHealth EEG API Running"
     }
 
-# =========================
-# SINGLE MODEL PREDICTIONS
-# =========================
-@app.post("/predict/121")
-async def predict_121(file: UploadFile = File(...)):
-    img = process_image(await file.read())
-    result = predict(model_121, img)
-    result["model"] = "121"
-    return result
+# =========================================================
+# ALZHEIMER PREDICTION
+# =========================================================
+
+@app.post("/predict/alzheimer")
 
+def predict_alzheimer(request: EEGRequest):
 
-@app.post("/predict/169")
-async def predict_169(file: UploadFile = File(...)):
-    img = process_image(await file.read())
-    result = predict(model_169, img)
-    result["model"] = "169"
-    return result
+    pred, confidence, probs = predict_model(
+        ad_mlp,
+        request.features
+    )
+
+    return {
 
+        "prediction": AD_CLASSES[pred],
 
-@app.post("/predict/201")
-async def predict_201(file: UploadFile = File(...)):
-    img = process_image(await file.read())
-    result = predict(model_201, img)
-    result["model"] = "201"
-    return result
+        "confidence": confidence,
 
-# =========================
-# ENSEMBLE PREDICTION (FINAL FIXED)
-# =========================
-@app.post("/predict/ensemble")
-async def ensemble(file: UploadFile = File(...)):
+        "probabilities": {
 
-    img = process_image(await file.read())
+            AD_CLASSES[i]: float(probs[i])
+
+            for i in range(len(AD_CLASSES))
+        }
+    }
 
-    r1 = predict(model_121, img)
-    r2 = predict(model_169, img)
-    r3 = predict(model_201, img)
+# =========================================================
+# PARKINSON PREDICTION
+# =========================================================
 
-    avg_probs = {}
+@app.post("/predict/parkinson")
 
-    for c in CLASSES:
-        avg_probs[c] = (
-            r1["probabilities"][c] +
-            r2["probabilities"][c] +
-            r3["probabilities"][c]
-        ) / 3
+def predict_parkinson(request: EEGRequest):
 
-    final_class = max(avg_probs, key=avg_probs.get)
+    pred, confidence, probs = predict_model(
+        pd_mlp,
+        request.features
+    )
 
     return {
-        "prediction": final_class,
-        "confidence": avg_probs[final_class],
-        "probabilities": avg_probs,
-        "models": {
-            "121": r1,
-            "169": r2,
-            "201": r3
+
+        "prediction": PD_CLASSES[pred],
+
+        "confidence": confidence,
+
+        "probabilities": {
+
+            PD_CLASSES[i]: float(probs[i])
+
+            for i in range(len(PD_CLASSES))
         }
     }