Update app.py

app.py

@@ -2,6 +2,8 @@ import torch
 import torch.nn as nn
 import torchvision.models as models
 from fastapi import FastAPI, UploadFile, File
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
 from PIL import Image
 import io
 import torchvision.transforms as transforms
@@ -9,7 +11,18 @@ import torchvision.transforms as transforms
 # =========================
 # APP INIT
 # =========================
-app = FastAPI(title="Alzheimer Ensemble API")
+app = FastAPI(title="Alzheimer Ensemble API", version="1.0")
+
+# =========================
+# CORS
+# =========================
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
 # =========================
 # DEVICE
@@ -98,7 +111,8 @@ def process_image(image_bytes):
     return img
 
 # =========================
-# PREDICTION FUNCTION (FIXED)
+# PREDICTION FUNCTION
+# All confidence values are returned as floats in range [0.0, 1.0]
 # =========================
 def predict(model, x):
     with torch.no_grad():
@@ -111,6 +125,7 @@ def predict(model, x):
         return {
             "prediction": CLASSES[cls],
             "class_id": cls,
+            # Confidence as 0.0–1.0 decimal
             "confidence": float(conf.item()),
             "probabilities": {
                 CLASSES[i]: float(probs[i].item())
@@ -118,6 +133,23 @@ def predict(model, x):
             }
         }
 
+# =========================
+# HEALTH ENDPOINT
+# =========================
+@app.get("/health")
+def health():
+    return {
+        "status": "running",
+        "service": "Alzheimer MRI Ensemble API",
+        "models_loaded": {
+            "densenet121": True,
+            "densenet169": True,
+            "densenet201": True,
+        },
+        "device": str(DEVICE),
+        "classes": CLASSES,
+    }
+
 # =========================
 # ROOT ENDPOINT
 # =========================
@@ -128,6 +160,7 @@ def home():
         "models": ["121", "169", "201"],
         "classes": CLASSES,
         "endpoints": [
+            "/health",
             "/predict/121",
             "/predict/169",
             "/predict/201",
@@ -142,27 +175,29 @@ def home():
 async def predict_121(file: UploadFile = File(...)):
     img = process_image(await file.read())
     result = predict(model_121, img)
-    result["model"] = "121"
-    return result
+    result["model"] = "densenet121"
+    return JSONResponse(result)
 
 
 @app.post("/predict/169")
 async def predict_169(file: UploadFile = File(...)):
     img = process_image(await file.read())
     result = predict(model_169, img)
-    result["model"] = "169"
-    return result
+    result["model"] = "densenet169"
+    return JSONResponse(result)
 
 
 @app.post("/predict/201")
 async def predict_201(file: UploadFile = File(...)):
     img = process_image(await file.read())
     result = predict(model_201, img)
-    result["model"] = "201"
-    return result
+    result["model"] = "densenet201"
+    return JSONResponse(result)
 
 # =========================
-# ENSEMBLE PREDICTION (FINAL FIXED)
+# ENSEMBLE PREDICTION
+# Returns ensemble result with individual model results
+# All confidence values are 0.0-1.0
 # =========================
 @app.post("/predict/ensemble")
 async def ensemble(file: UploadFile = File(...)):
@@ -173,8 +208,8 @@ async def ensemble(file: UploadFile = File(...)):
     r2 = predict(model_169, img)
     r3 = predict(model_201, img)
 
+    # Average probabilities (all already 0.0-1.0)
     avg_probs = {}
-
     for c in CLASSES:
         avg_probs[c] = (
             r1["probabilities"][c] +
@@ -183,14 +218,17 @@ async def ensemble(file: UploadFile = File(...)):
         ) / 3
 
     final_class = max(avg_probs, key=avg_probs.get)
+    # Ensemble confidence = max averaged probability (0.0-1.0)
+    ensemble_confidence = avg_probs[final_class]
 
-    return {
+    return JSONResponse({
         "prediction": final_class,
-        "confidence": avg_probs[final_class],
+        # confidence as 0.0-1.0
+        "confidence": ensemble_confidence,
         "probabilities": avg_probs,
-        "models": {
-            "121": r1,
-            "169": r2,
-            "201": r3
+        "individual_models": {
+            "densenet121": r1,
+            "densenet169": r2,
+            "densenet201": r3,
         }
-    }
+    })