Update app.py

app.py

@@ -1,208 +1,153 @@
-import os
-import io
-import numpy as np
 import torch
 import torch.nn as nn
 import torchvision.models as models
-import tensorflow as tf
 from fastapi import FastAPI, UploadFile, File
 from PIL import Image
-import cv2
-import nibabel as nib
+import io
+import torchvision.transforms as transforms
+import numpy as np
+import tensorflow as tf
 
-app = FastAPI(title="Parkinson + DaTscan AI API")
+app = FastAPI(title="Alzheimer Ensemble API")
 
-DEVICE = "cpu"
+DEVICE = torch.device("cpu")
 
 # =========================
-# LABELS (BINARY CLASS)
+# LABELS
 # =========================
-LABELS = ["No Parkinson's", "Parkinson's Disease"]
+LABELS = [
+    "Mild Demented",
+    "Moderate Demented",
+    "Non Demented",
+    "Very Mild Demented"
+]
 
 # =========================
-# LOAD KERAS MODELS
+# TRANSFORM
 # =========================
-def load_keras(path):
-    return tf.keras.models.load_model(path, compile=False)
-
-model_121 = load_keras("densenet121_parkinsonsDATSCAN.keras")
-model_169 = load_keras("parkinsons_densenet169DATSCAN.keras")
-model_201 = load_keras("parkinsons_densenet201DATSCAN.keras")
-
-# optional fixed model (if better)
-model_fixed = load_keras("densenet121_parkinsonsDATSCAN_fixed.keras")
-
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.5]*3, [0.5]*3)
+])
 
 # =========================
-# 3D CNN MODEL (PyTorch)
+# SAFE KERAS LOADER
 # =========================
-def build_3dcnn():
-    model = nn.Sequential(
-        nn.Conv3d(1, 32, 3, padding=1),
-        nn.ReLU(),
-        nn.MaxPool3d(2),
-
-        nn.Conv3d(32, 64, 3, padding=1),
-        nn.ReLU(),
-        nn.MaxPool3d(2),
-
-        nn.AdaptiveAvgPool3d((4, 4, 4)),
-        nn.Flatten(),
-        nn.Linear(64 * 4 * 4 * 4, 128),
-        nn.ReLU(),
-        nn.Linear(128, 2)
-    )
-    return model.to(DEVICE).eval()
-
-model_3dcnn = build_3dcnn()
-
+def load_keras_model(path):
+    try:
+        return tf.keras.models.load_model(path, compile=False)
+    except Exception as e:
+        print("Keras load failed:", path, e)
+        return None
 
 # =========================
-# IMAGE PREPROCESSING (2D)
+# LOAD MODELS
 # =========================
-def preprocess_2d(image_bytes):
-    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-    image = np.array(image)
-    image = cv2.resize(image, (128, 128))
-    image = image / 255.0
-    image = np.expand_dims(image, axis=0)
-    return image
+model_121 = load_keras_model("densenet121_parkinsonsDATSCAN.keras")
+model_169 = load_keras_model("parkinsons_densenet169DATSCAN.keras")
+model_201 = load_keras_model("parkinsons_densenet201DATSCAN.keras")
 
+models_list = [model_121, model_169, model_201]
 
 # =========================
-# NIFTI PREPROCESSING (3D)
+# SINGLE PREDICTION (KERAS SAFE)
 # =========================
-def preprocess_3d(file_bytes):
-    temp_path = "temp.nii"
-    with open(temp_path, "wb") as f:
-        f.write(file_bytes)
-
-    volume = nib.load(temp_path).get_fdata()
-    volume = np.squeeze(volume)
-
-    if len(volume.shape) == 2:
-        volume = np.stack([volume] * 32, axis=-1)
+def predict_single(model, image_tensor):
+    if model is None:
+        return np.zeros(len(LABELS))
 
-    depth = volume.shape[2]
+    img = image_tensor.permute(0, 2, 3, 1).numpy()
 
-    slices = []
-    for i in np.linspace(0, depth - 1, 32).astype(int):
-        sl = volume[:, :, i]
-        sl = cv2.resize(sl, (64, 64))
-        slices.append(sl)
+    preds = model.predict(img, verbose=0)[0]
+    return preds
 
-    vol = np.stack(slices, axis=0)
-    vol = np.expand_dims(vol, axis=0)
-    vol = np.expand_dims(vol, axis=0)
-
-    return torch.tensor(vol, dtype=torch.float32).to(DEVICE)
-
-
-# =========================
-# KERAS SINGLE PREDICT
 # =========================
-def predict_keras(model, x):
-    pred = model.predict(x, verbose=0)[0]
-    return pred
-
-
-# =========================
-# ENSEMBLE LOGIC (4 MODELS)
+# ENSEMBLE PREDICTION
 # =========================
 def ensemble_predict(image_bytes):
+    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+    image = transform(image).unsqueeze(0)
 
-    x = preprocess_2d(image_bytes)
+    preds_all = []
+    model_confidence_report = []
 
-    p1 = predict_keras(model_121, x)
-    p2 = predict_keras(model_169, x)
-    p3 = predict_keras(model_201, x)
-    p4 = predict_keras(model_fixed, x)
+    for i, m in enumerate(models_list):
+        preds = predict_single(m, image)
+        preds_all.append(preds)
 
-    preds = np.array([p1, p2, p3, p4])
+        model_confidence_report.append({
+            "model": f"model_{i+1}",
+            "confidence": float(np.max(preds)),
+            "prediction": LABELS[int(np.argmax(preds))]
+        })
 
-    avg = np.mean(preds, axis=0)
+    avg = np.mean(preds_all, axis=0)
 
     cls = int(np.argmax(avg))
-    confidence = float(avg[cls] * 100)
+    conf = float(np.max(avg))
 
     return {
         "prediction": LABELS[cls],
         "class_id": cls,
-        "confidence": round(confidence, 2),
-
-        "model_confidences": {
-            "DenseNet121": round(float(np.max(p1)) * 100, 2),
-            "DenseNet169": round(float(np.max(p2)) * 100, 2),
-            "DenseNet201": round(float(np.max(p3)) * 100, 2),
-            "Fixed121": round(float(np.max(p4)) * 100, 2),
-        },
-
+        "confidence": round(conf * 100, 2),
         "probabilities": {
             LABELS[i]: round(float(avg[i]) * 100, 2)
-            for i in range(2)
-        }
+            for i in range(len(LABELS))
+        },
+        "model_breakdown": model_confidence_report
     }
 
-
 # =========================
-# 3D CNN PREDICTION
+# INDIVIDUAL ENDPOINTS
 # =========================
-def predict_3d(file_bytes):
-    x = preprocess_3d(file_bytes)
+@app.post("/predict/121")
+async def predict_121(file: UploadFile = File(...)):
+    img = await file.read()
+    return ensemble_predict_single(img, model_121)
 
-    with torch.no_grad():
-        out = model_3dcnn(x)
-        probs = torch.softmax(out, dim=1)[0]
+@app.post("/predict/169")
+async def predict_169(file: UploadFile = File(...)):
+    img = await file.read()
+    return ensemble_predict_single(img, model_169)
+
+@app.post("/predict/201")
+async def predict_201(file: UploadFile = File(...)):
+    img = await file.read()
+    return ensemble_predict_single(img, model_201)
+
+# helper
+def ensemble_predict_single(image_bytes, model):
+    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+    image = transform(image).unsqueeze(0)
 
-    cls = int(torch.argmax(probs))
-    confidence = float(probs[cls] * 100)
+    preds = predict_single(model, image)
+    cls = int(np.argmax(preds))
 
     return {
         "prediction": LABELS[cls],
-        "class_id": cls,
-        "confidence": round(confidence, 2),
+        "confidence": float(np.max(preds)),
         "probabilities": {
-            LABELS[i]: round(float(probs[i]) * 100, 2)
-            for i in range(2)
+            LABELS[i]: float(preds[i])
+            for i in range(len(LABELS))
         }
     }
 
-
 # =========================
-# ROUTES
+# ENSEMBLE ENDPOINT
 # =========================
-@app.get("/")
-def home():
-    return {
-        "status": "running",
-        "models": ["121", "169", "201", "fixed", "3dcnn"]
-    }
-
-
 @app.post("/predict")
 async def predict(file: UploadFile = File(...)):
     image_bytes = await file.read()
     return ensemble_predict(image_bytes)
 
-
-@app.post("/predict/121")
-async def p121(file: UploadFile = File(...)):
-    x = preprocess_2d(await file.read())
-    return {"model": "121", "prob": predict_keras(model_121, x).tolist()}
-
-
-@app.post("/predict/169")
-async def p169(file: UploadFile = File(...)):
-    x = preprocess_2d(await file.read())
-    return {"model": "169", "prob": predict_keras(model_169, x).tolist()}
-
-
-@app.post("/predict/201")
-async def p201(file: UploadFile = File(...)):
-    x = preprocess_2d(await file.read())
-    return {"model": "201", "prob": predict_keras(model_201, x).tolist()}
-
-
-@app.post("/predict/3d")
-async def p3d(file: UploadFile = File(...)):
-    return predict_3d(await file.read())
+# =========================
+# HEALTH CHECK
+# =========================
+@app.get("/")
+def home():
+    return {
+        "status": "running",
+        "models": ["121", "169", "201"],
+        "ensemble": True
+    }