Update app.py

app.py

@@ -6,7 +6,10 @@ from PIL import Image
 import io
 import torchvision.transforms as transforms
 
-app = FastAPI(title="Parkinson DATSCAN Ensemble API")
+# =========================
+# APP
+# =========================
+app = FastAPI(title="Alzheimer DATSCAN Ensemble API")
 
 # =========================
 # DEVICE
@@ -15,9 +18,14 @@ DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print("Using device:", DEVICE)
 
 # =========================
-# LABELS
+# CLASSES (VERY IMPORTANT - YOUR ORIGINAL SETUP)
 # =========================
-LABELS = ["Control", "Prodromal", "Parkinsons"]
+CLASSES = [
+    "Mild Demented",
+    "Moderate Demented",
+    "Non Demented",
+    "Very Mild Demented"
+]
 
 # =========================
 # IMAGE TRANSFORM
@@ -25,43 +33,47 @@ LABELS = ["Control", "Prodromal", "Parkinsons"]
 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
+# MODEL BUILDER (MATCH TRAINING)
 # =========================
-def build_densenet(name="121"):
-    if name == "121":
+def build_model(version="121"):
+    if version == "121":
         model = models.densenet121(weights=None)
-        in_f = 1024
-    elif name == "169":
+        in_features = 1024
+    elif version == "169":
         model = models.densenet169(weights=None)
-        in_f = 1664
+        in_features = 1664
     else:
         model = models.densenet201(weights=None)
-        in_f = 1920
+        in_features = 1920
+
+    model.classifier = nn.Sequential(
+        nn.Dropout(0.4),
+        nn.Linear(in_features, 4)
+    )
 
-    model.classifier = nn.Linear(in_f, len(LABELS))
     return model
 
 # =========================
-# SAFE LOADER (IMPORTANT FIX)
+# SAFE MODEL LOADER (FIX ALL YOUR ERRORS)
 # =========================
-def load_model(path, name):
-    model = build_densenet(name)
+def load_model(path, version):
+    model = build_model(version)
 
     try:
-        checkpoint = torch.load(path, map_location=DEVICE)
+        state = torch.load(path, map_location=DEVICE)
 
-        # handle dict checkpoint formats
-        if isinstance(checkpoint, dict):
-            if "state_dict" in checkpoint:
-                checkpoint = checkpoint["state_dict"]
-            elif "model_state_dict" in checkpoint:
-                checkpoint = checkpoint["model_state_dict"]
+        # 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"]
 
-        model.load_state_dict(checkpoint, strict=False)
+        model.load_state_dict(state, strict=False)
 
         print(f"Loaded {path}")
 
@@ -75,35 +87,34 @@ def load_model(path, name):
 # =========================
 # LOAD MODELS
 # =========================
-model_121 = load_model("densenet121.pth", "121")
-model_169 = load_model("densenet169.pth", "169")
-model_201 = load_model("densenet201.pth", "201")
+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):
-    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-    tensor = transform(image).unsqueeze(0).to(DEVICE)
-    return tensor
+    img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+    img = transform(img).unsqueeze(0).to(DEVICE)
+    return img
 
 # =========================
-# PREDICTION FUNCTION (FIXED)
+# PREDICT FUNCTION (FIXED)
 # =========================
-def predict(model, image_tensor):
+def predict(model, x):
     with torch.no_grad():
-        out = model(image_tensor)
+        out = model(x)
         probs = torch.softmax(out, dim=1)[0]
-
-        conf, cls = torch.max(probs, dim=0)
+        conf, cls = torch.max(probs, 0)
 
         return {
             "class_id": int(cls.item()),
-            "class_name": LABELS[int(cls.item())],
+            "class_name": CLASSES[int(cls.item())],
             "confidence": float(conf.item()),
             "probabilities": {
-                LABELS[i]: float(probs[i].item())
-                for i in range(len(LABELS))
+                CLASSES[i]: float(probs[i].item())
+                for i in range(len(CLASSES))
             }
         }
 
@@ -114,30 +125,31 @@ def predict(model, image_tensor):
 def home():
     return {
         "status": "running",
-        "models": ["121", "169", "201"],
+        "models": ["densenet121", "densenet169", "densenet201"],
+        "classes": CLASSES,
         "endpoints": ["/predict/121", "/predict/169", "/predict/201", "/predict/ensemble"]
     }
 
 # =========================
-# ENDPOINTS
+# SINGLE MODEL ROUTES
 # =========================
 @app.post("/predict/121")
 async def predict_121(file: UploadFile = File(...)):
     img = process_image(await file.read())
-    return {"model": "densenet121", **predict(model_121, img)}
+    return {"model": "121", **predict(model_121, img)}
 
 @app.post("/predict/169")
 async def predict_169(file: UploadFile = File(...)):
     img = process_image(await file.read())
-    return {"model": "densenet169", **predict(model_169, img)}
+    return {"model": "169", **predict(model_169, img)}
 
 @app.post("/predict/201")
 async def predict_201(file: UploadFile = File(...)):
     img = process_image(await file.read())
-    return {"model": "densenet201", **predict(model_201, img)}
+    return {"model": "201", **predict(model_201, img)}
 
 # =========================
-# ENSEMBLE (FIXED)
+# ENSEMBLE (FINAL FIXED VERSION)
 # =========================
 @app.post("/predict/ensemble")
 async def ensemble(file: UploadFile = File(...)):
@@ -148,24 +160,22 @@ async def ensemble(file: UploadFile = File(...)):
     r2 = predict(model_169, img)
     r3 = predict(model_201, img)
 
-    avg_probs = {}
+    avg = {}
 
-    for i in range(len(LABELS)):
-        avg_probs[LABELS[i]] = (
-            r1["probabilities"][LABELS[i]] +
-            r2["probabilities"][LABELS[i]] +
-            r3["probabilities"][LABELS[i]]
-        ) / 3
+    for c in CLASSES:
+        avg[c] = (r1["probabilities"][c] +
+                  r2["probabilities"][c] +
+                  r3["probabilities"][c]) / 3
 
-    final_class = max(avg_probs, key=avg_probs.get)
+    final_class = max(avg, key=avg.get)
 
     return {
         "final_prediction": final_class,
-        "final_confidence": avg_probs[final_class],
-        "models": {
+        "final_confidence": avg[final_class],
+        "individual_models": {
             "121": r1,
             "169": r2,
             "201": r3
         },
-        "averaged_probabilities": avg_probs
+        "ensemble_probabilities": avg
     }