Update app.py

app.py

@@ -5,34 +5,36 @@ from fastapi import FastAPI, UploadFile, File
 from PIL import Image
 import io
 import torchvision.transforms as transforms
-from pathlib import Path
 
+# =========================
+# APP
+# =========================
 app = FastAPI(title="Alzheimer Ensemble API")
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# -----------------------------
+# =========================
 # LABELS
-# -----------------------------
+# =========================
 LABELS = [
     "Mild Demented",
     "Moderate Demented",
     "Non Demented",
-    "Very Very Mild Demented"
+    "Very Mild Demented"
 ]
 
-# -----------------------------
-# TRANSFORM
-# -----------------------------
+# =========================
+# IMAGE TRANSFORM
+# =========================
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
-    transforms.Normalize([0.5]*3, [0.5]*3)
+    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
 ])
 
-# -----------------------------
-# MODEL BUILDER (MATCH YOUR TRAINING)
-# -----------------------------
+# =========================
+# MODEL ARCHITECTURE
+# =========================
 def build_model():
     model = models.densenet121(weights=None)
 
@@ -46,31 +48,43 @@ def build_model():
 
     return model
 
-# -----------------------------
-# LOAD MODEL
-# -----------------------------
+# =========================
+# LOAD MODEL (FIXED)
+# =========================
 def load_model(path):
     model = build_model()
+
     state = torch.load(path, map_location=DEVICE)
 
+    # handle checkpoint formats safely
     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
 
-# -----------------------------
-# LOAD ALL MODELS
-# -----------------------------
-model_121 = "alzheimers_densenet121.pth"
-model_169 = "alzheimers_densenet169.pth"
-model_201 = "alzheimers_densenet201.pth"
+# =========================
+# LOAD ALL MODELS (REAL FIX HERE)
+# =========================
+model_121 = load_model("alzheimers_densenet121.pth")
+model_169 = load_model("alzheimers_densenet169.pth")
+model_201 = load_model("alzheimers_densenet201.pth")
+
+# =========================
+# IMAGE PROCESSING
+# =========================
+def process_image(image_bytes):
+    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+    image = transform(image).unsqueeze(0)
+    return image.to(DEVICE)
 
-# -----------------------------
-# SINGLE PREDICT FUNCTION
-# -----------------------------
+# =========================
+# SINGLE MODEL PREDICT
+# =========================
 def predict(model, image_tensor):
     with torch.no_grad():
         out = model(image_tensor)
@@ -78,25 +92,18 @@ def predict(model, image_tensor):
         conf, cls = torch.max(probs, dim=0)
 
     return {
-        "class_id": cls.item(),
+        "class_id": int(cls.item()),
         "class_name": LABELS[cls.item()],
-        "confidence": float(conf.item()),
+        "confidence": round(float(conf.item()), 4),
         "probabilities": {
-            LABELS[i]: float(probs[i].item())
+            LABELS[i]: round(float(probs[i].item()), 4)
             for i in range(len(LABELS))
         }
     }
 
-# -----------------------------
-# IMAGE PROCESSING
-# -----------------------------
-def process_image(image_bytes):
-    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-    return transform(image).unsqueeze(0).to(DEVICE)
-
-# -----------------------------
-# INDIVIDUAL ENDPOINTS
-# -----------------------------
+# =========================
+# ENDPOINTS
+# =========================
 @app.post("/predict/121")
 async def predict_121(file: UploadFile = File(...)):
     img = process_image(await file.read())
@@ -112,9 +119,9 @@ async def predict_201(file: UploadFile = File(...)):
     img = process_image(await file.read())
     return {"model": "densenet201", **predict(model_201, img)}
 
-# -----------------------------
-# ENSEMBLE PREDICTION (UPGRADED)
-# -----------------------------
+# =========================
+# ENSEMBLE (FIXED + CLEAN)
+# =========================
 @app.post("/predict/ensemble")
 async def ensemble(file: UploadFile = File(...)):
 
@@ -124,8 +131,8 @@ async def ensemble(file: UploadFile = File(...)):
     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]] +
@@ -137,25 +144,26 @@ async def ensemble(file: UploadFile = File(...)):
 
     return {
         "final_prediction": final_class,
-        "final_confidence": avg_probs[final_class],
+        "final_confidence": round(avg_probs[final_class], 4),
 
-        "ensemble_breakdown": {
-            "model_121": r1,
-            "model_169": r2,
-            "model_201": r3
-        },
+        "averaged_probabilities": avg_probs,
 
-        "averaged_probabilities": avg_probs
+        "model_outputs": {
+            "densenet121": r1,
+            "densenet169": r2,
+            "densenet201": r3
+        }
     }
 
-# -----------------------------
+# =========================
 # ROOT
-# -----------------------------
+# =========================
 @app.get("/")
 def home():
     return {
         "status": "running",
-        "models": ["121", "169", "201"],
+        "device": str(DEVICE),
+        "models_loaded": True,
         "endpoints": [
             "/predict/121",
             "/predict/169",