Update app.py

app.py

@@ -12,35 +12,36 @@ app = FastAPI(title="Alzheimer Ensemble API")
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # -----------------------------
-# LABELS (YOUR REAL ONES)
+# LABELS
 # -----------------------------
 LABELS = [
     "Mild Demented",
     "Moderate Demented",
     "Non Demented",
-    "Very Mild Demented"
+    "Very Very Mild Demented"
 ]
 
 # -----------------------------
-# IMAGE PREPROCESSING
+# TRANSFORM
 # -----------------------------
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
-    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
+    transforms.Normalize([0.5]*3, [0.5]*3)
 ])
 
 # -----------------------------
-# MODEL ARCHITECTURE (4-LAYER HEAD)
+# MODEL BUILDER (MATCH YOUR TRAINING)
 # -----------------------------
 def build_model():
     model = models.densenet121(weights=None)
 
     model.classifier = nn.Sequential(
-        nn.Linear(1024, 512),
+        nn.Dropout(0.5),
+        nn.Linear(1024, 256),
         nn.ReLU(),
         nn.Dropout(0.3),
-        nn.Linear(512, 4)
+        nn.Linear(256, 4)
     )
 
     return model
@@ -52,68 +53,115 @@ 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)
+    if isinstance(state, dict) and "model_state_dict" in state:
+        state = state["model_state_dict"]
 
+    model.load_state_dict(state, strict=False)
     model.to(DEVICE)
     model.eval()
     return model
 
 # -----------------------------
-# MODEL PATHS
+# LOAD ALL MODELS
 # -----------------------------
+BASE = Path("saved_models")
 
-
-model_121 = "alzheimers_densenet121.pth"
-model_169 = "alzheimers_densenet169.pth"
-model_201 = "alzheimers_densenet201.pth"
+model_121 = load_model(BASE / "alzheimers_densenet121.pth")
+model_169 = load_model(BASE / "alzheimers_densenet169.pth")
+model_201 = load_model(BASE / "alzheimers_densenet201.pth")
 
 # -----------------------------
-# SINGLE MODEL PREDICTION
+# SINGLE PREDICT FUNCTION
 # -----------------------------
-def predict_single(model, image_tensor):
+def predict(model, image_tensor):
     with torch.no_grad():
         out = model(image_tensor)
-        probs = torch.softmax(out, dim=1)
-    return probs[0]
+        probs = torch.softmax(out, dim=1)[0]
+        conf, cls = torch.max(probs, dim=0)
+
+    return {
+        "class_id": cls.item(),
+        "class_name": LABELS[cls.item()],
+        "confidence": float(conf.item()),
+        "probabilities": {
+            LABELS[i]: float(probs[i].item())
+            for i in range(len(LABELS))
+        }
+    }
 
 # -----------------------------
-# ENSEMBLE
+# IMAGE PROCESSING
 # -----------------------------
-def ensemble_predict(image_bytes):
-
+def process_image(image_bytes):
     image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-    image = transform(image).unsqueeze(0).to(DEVICE)
+    return transform(image).unsqueeze(0).to(DEVICE)
+
+# -----------------------------
+# INDIVIDUAL ENDPOINTS
+# -----------------------------
+@app.post("/predict/121")
+async def predict_121(file: UploadFile = File(...)):
+    img = process_image(await file.read())
+    return {"model": "densenet121", **predict(model_121, img)}
 
-    p1 = predict_single(model_121, image)
-    p2 = predict_single(model_169, image)
-    p3 = predict_single(model_201, image)
+@app.post("/predict/169")
+async def predict_169(file: UploadFile = File(...)):
+    img = process_image(await file.read())
+    return {"model": "densenet169", **predict(model_169, img)}
 
-    avg = (p1 + p2 + p3) / 3
+@app.post("/predict/201")
+async def predict_201(file: UploadFile = File(...)):
+    img = process_image(await file.read())
+    return {"model": "densenet201", **predict(model_201, img)}
 
-    confidence, cls = torch.max(avg, dim=0)
+# -----------------------------
+# ENSEMBLE PREDICTION (UPGRADED)
+# -----------------------------
+@app.post("/predict/ensemble")
+async def ensemble(file: UploadFile = File(...)):
+
+    img = process_image(await file.read())
+
+    r1 = predict(model_121, img)
+    r2 = predict(model_169, img)
+    r3 = predict(model_201, img)
+
+    # average probabilities
+    avg_probs = {}
+    for i in range(len(LABELS)):
+        avg_probs[LABELS[i]] = (
+            r1["probabilities"][LABELS[i]] +
+            r2["probabilities"][LABELS[i]] +
+            r3["probabilities"][LABELS[i]]
+        ) / 3
+
+    final_class = max(avg_probs, key=avg_probs.get)
 
     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))
-        }
+        "final_prediction": final_class,
+        "final_confidence": avg_probs[final_class],
+
+        "ensemble_breakdown": {
+            "model_121": r1,
+            "model_169": r2,
+            "model_201": r3
+        },
+
+        "averaged_probabilities": avg_probs
     }
 
 # -----------------------------
-# ROUTES
+# ROOT
 # -----------------------------
 @app.get("/")
 def home():
     return {
-        "message": "Alzheimer Ensemble API Running",
-        "classes": LABELS
-    }
-
-@app.post("/predict")
-async def predict(file: UploadFile = File(...)):
-    image_bytes = await file.read()
-    return ensemble_predict(image_bytes)
+        "status": "running",
+        "models": ["121", "169", "201"],
+        "endpoints": [
+            "/predict/121",
+            "/predict/169",
+            "/predict/201",
+            "/predict/ensemble"
+        ]
+    }