Update app.py

app.py

@@ -7,9 +7,9 @@ import io
 import torchvision.transforms as transforms
 
 # =========================
-# APP
+# APP INIT
 # =========================
-app = FastAPI(title="Alzheimer DATSCAN Ensemble API")
+app = FastAPI(title="Alzheimer Ensemble API")
 
 # =========================
 # DEVICE
@@ -18,7 +18,7 @@ DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print("Using device:", DEVICE)
 
 # =========================
-# CLASSES (VERY IMPORTANT - YOUR ORIGINAL SETUP)
+# CLASS LABELS (FIXED)
 # =========================
 CLASSES = [
     "Mild Demented",
@@ -37,7 +37,7 @@ transform = transforms.Compose([
 ])
 
 # =========================
-# MODEL BUILDER (MATCH TRAINING)
+# MODEL BUILDER
 # =========================
 def build_model(version="121"):
     if version == "121":
@@ -52,30 +52,28 @@ def build_model(version="121"):
 
     model.classifier = nn.Sequential(
         nn.Dropout(0.4),
-        nn.Linear(in_features, 4)
+        nn.Linear(in_features, len(CLASSES))
     )
 
     return model
 
 # =========================
-# SAFE MODEL LOADER (FIX ALL YOUR ERRORS)
+# SAFE MODEL LOADER
 # =========================
 def load_model(path, version):
     model = build_model(version)
 
     try:
-        state = torch.load(path, map_location=DEVICE)
+        checkpoint = torch.load(path, map_location=DEVICE)
 
-        # handle different checkpoint formats
-        if isinstance(state, dict):
-            if "state_dict" in state:
-                state = state["state_dict"]
-            elif "model_state_dict" in state:
-                state = state["model_state_dict"]
+        if isinstance(checkpoint, dict):
+            if "state_dict" in checkpoint:
+                checkpoint = checkpoint["state_dict"]
+            elif "model_state_dict" in checkpoint:
+                checkpoint = checkpoint["model_state_dict"]
 
-        model.load_state_dict(state, strict=False)
-
-        print(f"Loaded {path}")
+        model.load_state_dict(checkpoint, strict=False)
+        print(f"Loaded: {path}")
 
     except Exception as e:
         print(f"Failed loading {path}: {e}")
@@ -100,17 +98,19 @@ def process_image(image_bytes):
     return img
 
 # =========================
-# PREDICT FUNCTION (FIXED)
+# PREDICTION FUNCTION (FIXED)
 # =========================
 def predict(model, x):
     with torch.no_grad():
-        out = model(x)
-        probs = torch.softmax(out, dim=1)[0]
+        output = model(x)
+        probs = torch.softmax(output, dim=1)[0]
+
         conf, cls = torch.max(probs, 0)
+        cls = int(cls.item())
 
         return {
-            "class_id": int(cls.item()),
-            "class_name": CLASSES[int(cls.item())],
+            "prediction": CLASSES[cls],
+            "class_id": cls,
             "confidence": float(conf.item()),
             "probabilities": {
                 CLASSES[i]: float(probs[i].item())
@@ -119,37 +119,50 @@ def predict(model, x):
         }
 
 # =========================
-# ROOT
+# ROOT ENDPOINT
 # =========================
 @app.get("/")
 def home():
     return {
         "status": "running",
-        "models": ["densenet121", "densenet169", "densenet201"],
+        "models": ["121", "169", "201"],
         "classes": CLASSES,
-        "endpoints": ["/predict/121", "/predict/169", "/predict/201", "/predict/ensemble"]
+        "endpoints": [
+            "/predict/121",
+            "/predict/169",
+            "/predict/201",
+            "/predict/ensemble"
+        ]
     }
 
 # =========================
-# SINGLE MODEL ROUTES
+# SINGLE MODEL PREDICTIONS
 # =========================
 @app.post("/predict/121")
 async def predict_121(file: UploadFile = File(...)):
     img = process_image(await file.read())
-    return {"model": "121", **predict(model_121, img)}
+    result = predict(model_121, img)
+    result["model"] = "121"
+    return result
+
 
 @app.post("/predict/169")
 async def predict_169(file: UploadFile = File(...)):
     img = process_image(await file.read())
-    return {"model": "169", **predict(model_169, img)}
+    result = predict(model_169, img)
+    result["model"] = "169"
+    return result
+
 
 @app.post("/predict/201")
 async def predict_201(file: UploadFile = File(...)):
     img = process_image(await file.read())
-    return {"model": "201", **predict(model_201, img)}
+    result = predict(model_201, img)
+    result["model"] = "201"
+    return result
 
 # =========================
-# ENSEMBLE (FINAL FIXED VERSION)
+# ENSEMBLE PREDICTION (FINAL FIXED)
 # =========================
 @app.post("/predict/ensemble")
 async def ensemble(file: UploadFile = File(...)):
@@ -160,22 +173,24 @@ async def ensemble(file: UploadFile = File(...)):
     r2 = predict(model_169, img)
     r3 = predict(model_201, img)
 
-    avg = {}
+    avg_probs = {}
 
     for c in CLASSES:
-        avg[c] = (r1["probabilities"][c] +
-                  r2["probabilities"][c] +
-                  r3["probabilities"][c]) / 3
+        avg_probs[c] = (
+            r1["probabilities"][c] +
+            r2["probabilities"][c] +
+            r3["probabilities"][c]
+        ) / 3
 
-    final_class = max(avg, key=avg.get)
+    final_class = max(avg_probs, key=avg_probs.get)
 
     return {
-        "final_prediction": final_class,
-        "final_confidence": avg[final_class],
-        "individual_models": {
+        "prediction": final_class,
+        "confidence": avg_probs[final_class],
+        "probabilities": avg_probs,
+        "models": {
             "121": r1,
             "169": r2,
             "201": r3
-        },
-        "ensemble_probabilities": avg
+        }
     }