Update app.py

app.py

@@ -1,80 +1,119 @@
 import torch
-import torchvision.models as models
 import torch.nn as nn
-from fastapi import FastAPI, UploadFile, File, HTTPException
+import torchvision.models as models
+from fastapi import FastAPI, UploadFile, File
 from PIL import Image
 import io
 import torchvision.transforms as transforms
+from pathlib import Path
 
-app = FastAPI(title="Alzheimer MRI API")
+app = FastAPI(title="Alzheimer Ensemble API")
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+# -----------------------------
+# LABELS (YOUR REAL ONES)
+# -----------------------------
+LABELS = [
+    "Mild Demented",
+    "Moderate Demented",
+    "Non Demented",
+    "Very Mild Demented"
+]
 
 # -----------------------------
-# MODEL LOADER
+# IMAGE PREPROCESSING
+# -----------------------------
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
+])
+
 # -----------------------------
-def load_model(path, num_classes=3):
+# MODEL ARCHITECTURE (4-LAYER HEAD)
+# -----------------------------
+def build_model():
     model = models.densenet121(weights=None)
-    model.classifier = nn.Linear(model.classifier.in_features, num_classes)
-    model.load_state_dict(torch.load(path, map_location=DEVICE))
+
+    model.classifier = nn.Sequential(
+        nn.Linear(1024, 512),
+        nn.ReLU(),
+        nn.Dropout(0.3),
+        nn.Linear(512, 4)
+    )
+
+    return model
+
+# -----------------------------
+# LOAD MODEL
+# -----------------------------
+def load_model(path):
+    model = build_model()
+    state = torch.load(path, map_location=DEVICE)
+
+    # safe load for HF spaces
+    model.load_state_dict(state, strict=False)
+
     model.to(DEVICE)
     model.eval()
     return model
 
-
 # -----------------------------
-# LOAD ALL MODELS
+# MODEL PATHS
 # -----------------------------
-MODELS = {
-    "ad_dn121": load_model("alzheimers_densenet121.pth"),
-    "ad_dn169": load_model("alzheimers_densenet169.pth"),
-    "ad_dn201": load_model("alzheimers_densenet201.pth"),
-}
+BASE = Path("saved_models")
 
+model_121 = load_model(BASE / "alzheimers_densenet121.pth")
+model_169 = load_model(BASE / "alzheimers_densenet169.pth")
+model_201 = load_model(BASE / "alzheimers_densenet201.pth")
 
 # -----------------------------
-# IMAGE TRANSFORM
+# SINGLE MODEL PREDICTION
 # -----------------------------
-transform = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-])
-
+def predict_single(model, image_tensor):
+    with torch.no_grad():
+        out = model(image_tensor)
+        probs = torch.softmax(out, dim=1)
+    return probs[0]
 
 # -----------------------------
-# PREDICTION FUNCTION
+# ENSEMBLE
 # -----------------------------
-def predict(model, image_bytes):
+def ensemble_predict(image_bytes):
+
     image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
     image = transform(image).unsqueeze(0).to(DEVICE)
 
-    with torch.no_grad():
-        outputs = model(image)
-        probs = torch.softmax(outputs, dim=1)
+    p1 = predict_single(model_121, image)
+    p2 = predict_single(model_169, image)
+    p3 = predict_single(model_201, image)
 
-    return probs[0].tolist()
+    avg = (p1 + p2 + p3) / 3
 
+    confidence, cls = torch.max(avg, dim=0)
+
+    return {
+        "prediction": LABELS[cls.item()],
+        "class_id": cls.item(),
+        "confidence": round(confidence.item() * 100, 2),
+        "probabilities": {
+            LABELS[i]: round(avg[i].item() * 100, 2)
+            for i in range(len(LABELS))
+        }
+    }
 
 # -----------------------------
 # ROUTES
 # -----------------------------
 @app.get("/")
 def home():
-    return {"message": "Alzheimer MRI API is running"}
-
+    return {
+        "message": "Alzheimer Ensemble API Running",
+        "classes": LABELS
+    }
 
 @app.post("/predict")
-async def predict_image(model_key: str, file: UploadFile = File(...)):
-
-    if model_key not in MODELS:
-        raise HTTPException(status_code=400, detail="Invalid model key")
-
+async def predict(file: UploadFile = File(...)):
     image_bytes = await file.read()
-
-    result = predict(MODELS[model_key], image_bytes)
-
-    return {
-        "model": model_key,
-        "prediction": result
-    }
+    return ensemble_predict(image_bytes)