Update app.py

app.py

@@ -1,153 +1,366 @@
+# =========================================
+# IMPORTS
+# =========================================
+import io
+import cv2
+import nibabel as nib
+import numpy as np
 import torch
 import torch.nn as nn
 import torchvision.models as models
+
 from fastapi import FastAPI, UploadFile, File
-from PIL import Image
-import io
-import torchvision.transforms as transforms
-import numpy as np
-import tensorflow as tf
+from fastapi.responses import JSONResponse
 
-app = FastAPI(title="Alzheimer Ensemble API")
+# =========================================
+# CONFIG
+# =========================================
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-DEVICE = torch.device("cpu")
+IMG_SIZE = 128
+NUM_CLASSES = 3
+NUM_SLICES = 32
 
-# =========================
-# LABELS
-# =========================
 LABELS = [
-    "Mild Demented",
-    "Moderate Demented",
-    "Non Demented",
-    "Very Mild Demented"
+    "Control",
+    "Prodromal",
+    "Parkinsons"
 ]
 
-# =========================
-# TRANSFORM
-# =========================
-transform = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize([0.5]*3, [0.5]*3)
-])
-
-# =========================
-# SAFE KERAS LOADER
-# =========================
-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
-
-# =========================
+print("Using device:", DEVICE)
+
+# =========================================
+# FASTAPI
+# =========================================
+app = FastAPI(
+    title="Parkinsons DATSCAN Ensemble API",
+    version="1.0"
+)
+
+# =========================================
+# PREPROCESSING
+# =========================================
+def load_nifti_from_bytes(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)
+
+    return volume
+
+
+def preprocess_2d(volume):
+    depth = volume.shape[2]
+
+    idx1 = np.linspace(0, depth//3 - 1, 10).astype(int)
+    idx2 = np.linspace(depth//3, 2*depth//3 - 1, 10).astype(int)
+    idx3 = np.linspace(2*depth//3, depth - 1, 12).astype(int)
+
+    def make_channel(indices):
+        slices = []
+
+        for i in indices:
+            sl = volume[:, :, i]
+
+            sl = sl - sl.min()
+            sl = sl / (sl.max() + 1e-6)
+
+            sl = cv2.resize(sl, (IMG_SIZE, IMG_SIZE))
+            slices.append(sl)
+
+        return np.mean(slices, axis=0)
+
+    r = make_channel(idx1)
+    g = make_channel(idx2)
+    b = make_channel(idx3)
+
+    img = np.stack([r, g, b], axis=0)
+
+    return torch.tensor(img, dtype=torch.float32).unsqueeze(0)
+
+
+def preprocess_3d(volume):
+    depth = volume.shape[2]
+
+    indices = np.linspace(0, depth - 1, NUM_SLICES).astype(int)
+
+    slices = []
+
+    for i in indices:
+        sl = volume[:, :, i]
+
+        sl = sl - sl.min()
+        sl = sl / (sl.max() + 1e-6)
+
+        sl = cv2.resize(sl, (IMG_SIZE, IMG_SIZE))
+        slices.append(sl)
+
+    vol = np.stack(slices)
+
+    vol = torch.tensor(vol, dtype=torch.float32)
+
+    vol = vol.unsqueeze(0).unsqueeze(0)
+
+    return vol
+
+
+# =========================================
+# DENSENET MODELS
+# =========================================
+class DenseNet121Model(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        self.base = models.densenet121(weights=None)
+
+        self.base.features.conv0 = nn.Conv2d(
+            3,
+            64,
+            kernel_size=7,
+            stride=2,
+            padding=3,
+            bias=False
+        )
+
+        self.base.classifier = nn.Linear(
+            self.base.classifier.in_features,
+            NUM_CLASSES
+        )
+
+    def forward(self, x):
+        return self.base(x)
+
+
+class DenseNet169Model(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        self.base = models.densenet169(weights=None)
+
+        self.base.features.conv0 = nn.Conv2d(
+            3,
+            64,
+            kernel_size=7,
+            stride=2,
+            padding=3,
+            bias=False
+        )
+
+        self.base.classifier = nn.Linear(
+            self.base.classifier.in_features,
+            NUM_CLASSES
+        )
+
+    def forward(self, x):
+        return self.base(x)
+
+
+class DenseNet201Model(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        self.base = models.densenet201(weights=None)
+
+        self.base.features.conv0 = nn.Conv2d(
+            3,
+            64,
+            kernel_size=7,
+            stride=2,
+            padding=3,
+            bias=False
+        )
+
+        self.base.classifier = nn.Linear(
+            self.base.classifier.in_features,
+            NUM_CLASSES
+        )
+
+    def forward(self, x):
+        return self.base(x)
+
+
+# =========================================
+# 3D CNN
+# =========================================
+class CNN3D(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        self.net = nn.Sequential(
+
+            nn.Conv3d(1, 16, 3, padding=1),
+            nn.ReLU(),
+            nn.MaxPool3d(2),
+
+            nn.Conv3d(16, 32, 3, padding=1),
+            nn.ReLU(),
+            nn.MaxPool3d(2),
+
+            nn.Conv3d(32, 64, 3, padding=1),
+            nn.ReLU(),
+            nn.MaxPool3d(2)
+
+        )
+
+        self.fc = nn.Sequential(
+            nn.Linear(64 * 4 * 16 * 16, 256),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(256, NUM_CLASSES)
+        )
+
+    def forward(self, x):
+        x = self.net(x)
+        x = x.view(x.size(0), -1)
+        return self.fc(x)
+
+
+# =========================================
 # LOAD MODELS
-# =========================
-model_121 = load_keras_model("densenet121_parkinsonsDATSCAN.keras")
-model_169 = load_keras_model("parkinsons_densenet169DATSCAN.keras")
-model_201 = load_keras_model("parkinsons_densenet201DATSCAN.keras")
+# =========================================
+def load_model(model, path):
+    model.load_state_dict(
+        torch.load(path, map_location=DEVICE)
+    )
 
-models_list = [model_121, model_169, model_201]
+    model.to(DEVICE)
+    model.eval()
 
-# =========================
-# SINGLE PREDICTION (KERAS SAFE)
-# =========================
-def predict_single(model, image_tensor):
-    if model is None:
-        return np.zeros(len(LABELS))
+    print(f"Loaded: {path}")
 
-    img = image_tensor.permute(0, 2, 3, 1).numpy()
+    return model
 
-    preds = model.predict(img, verbose=0)[0]
-    return preds
 
-# =========================
-# ENSEMBLE PREDICTION
-# =========================
-def ensemble_predict(image_bytes):
-    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-    image = transform(image).unsqueeze(0)
+model121 = load_model(DenseNet121Model(), "densenet121.pth")
+model169 = load_model(DenseNet169Model(), "densenet169.pth")
+model201 = load_model(DenseNet201Model(), "densenet201.pth")
+model3d = load_model(CNN3D(), "cnn3d.pth")
 
-    preds_all = []
-    model_confidence_report = []
 
-    for i, m in enumerate(models_list):
-        preds = predict_single(m, image)
-        preds_all.append(preds)
+# =========================================
+# PREDICTION HELPERS
+# =========================================
+def predict_model(model, tensor):
+    tensor = tensor.to(DEVICE)
 
-        model_confidence_report.append({
-            "model": f"model_{i+1}",
-            "confidence": float(np.max(preds)),
-            "prediction": LABELS[int(np.argmax(preds))]
-        })
+    with torch.no_grad():
+        out = model(tensor)
 
-    avg = np.mean(preds_all, axis=0)
+        probs = torch.softmax(out, dim=1)[0]
 
-    cls = int(np.argmax(avg))
-    conf = float(np.max(avg))
+    pred_idx = torch.argmax(probs).item()
 
     return {
-        "prediction": LABELS[cls],
-        "class_id": cls,
-        "confidence": round(conf * 100, 2),
+        "prediction": LABELS[pred_idx],
+        "class_id": pred_idx,
+        "confidence": round(float(probs[pred_idx]) * 100, 2),
         "probabilities": {
-            LABELS[i]: round(float(avg[i]) * 100, 2)
-            for i in range(len(LABELS))
-        },
-        "model_breakdown": model_confidence_report
+            LABELS[i]: round(float(probs[i]) * 100, 2)
+            for i in range(NUM_CLASSES)
+        }
+    }, probs
+
+
+# =========================================
+# ROOT
+# =========================================
+@app.get("/")
+def home():
+    return {
+        "message": "Parkinson DATSCAN Ensemble API Running",
+        "classes": LABELS
     }
 
-# =========================
+
+# =========================================
 # INDIVIDUAL ENDPOINTS
-# =========================
-@app.post("/predict/121")
+# =========================================
+@app.post("/predict/densenet121")
 async def predict_121(file: UploadFile = File(...)):
-    img = await file.read()
-    return ensemble_predict_single(img, model_121)
 
-@app.post("/predict/169")
+    volume = load_nifti_from_bytes(await file.read())
+    tensor = preprocess_2d(volume)
+
+    result, _ = predict_model(model121, tensor)
+
+    return JSONResponse(result)
+
+
+@app.post("/predict/densenet169")
 async def predict_169(file: UploadFile = File(...)):
-    img = await file.read()
-    return ensemble_predict_single(img, model_169)
 
-@app.post("/predict/201")
+    volume = load_nifti_from_bytes(await file.read())
+    tensor = preprocess_2d(volume)
+
+    result, _ = predict_model(model169, tensor)
+
+    return JSONResponse(result)
+
+
+@app.post("/predict/densenet201")
 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)
+    volume = load_nifti_from_bytes(await file.read())
+    tensor = preprocess_2d(volume)
 
-    preds = predict_single(model, image)
-    cls = int(np.argmax(preds))
+    result, _ = predict_model(model201, tensor)
 
-    return {
-        "prediction": LABELS[cls],
-        "confidence": float(np.max(preds)),
-        "probabilities": {
-            LABELS[i]: float(preds[i])
-            for i in range(len(LABELS))
-        }
-    }
+    return JSONResponse(result)
+
+
+@app.post("/predict/cnn3d")
+async def predict_3d(file: UploadFile = File(...)):
+
+    volume = load_nifti_from_bytes(await file.read())
+    tensor = preprocess_3d(volume)
+
+    result, _ = predict_model(model3d, tensor)
+
+    return JSONResponse(result)
+
+
+# =========================================
+# ENSEMBLE
+# =========================================
+@app.post("/predict/ensemble")
+async def predict_ensemble(file: UploadFile = File(...)):
+
+    volume = load_nifti_from_bytes(await file.read())
+
+    tensor2d = preprocess_2d(volume)
+    tensor3d = preprocess_3d(volume)
+
+    r121, p121 = predict_model(model121, tensor2d)
+    r169, p169 = predict_model(model169, tensor2d)
+    r201, p201 = predict_model(model201, tensor2d)
+    r3d, p3d = predict_model(model3d, tensor3d)
+
+    avg_probs = (p121 + p169 + p201 + p3d) / 4
+
+    pred_idx = torch.argmax(avg_probs).item()
 
-# =========================
-# ENSEMBLE ENDPOINT
-# =========================
-@app.post("/predict")
-async def predict(file: UploadFile = File(...)):
-    image_bytes = await file.read()
-    return ensemble_predict(image_bytes)
-
-# =========================
-# HEALTH CHECK
-# =========================
-@app.get("/")
-def home():
     return {
-        "status": "running",
-        "models": ["121", "169", "201"],
-        "ensemble": True
+
+        "ensemble_prediction": LABELS[pred_idx],
+
+        "ensemble_confidence": round(
+            float(avg_probs[pred_idx]) * 100,
+            2
+        ),
+
+        "ensemble_probabilities": {
+            LABELS[i]: round(float(avg_probs[i]) * 100, 2)
+            for i in range(NUM_CLASSES)
+        },
+
+        "individual_models": {
+
+            "DenseNet121": r121,
+            "DenseNet169": r169,
+            "DenseNet201": r201,
+            "CNN3D": r3d
+        }
     }