Update app.py

app.py

@@ -1,90 +1,80 @@
-import sys
-from pathlib import Path
-
-
-def _find_backend_path() -> Path:
-    current = Path(__file__).resolve().parent
-    for parent in [current] + list(current.parents):
-        candidate = parent / "backend"
-        if candidate.is_dir():
-            return candidate
-    raise RuntimeError("Could not find backend directory in parent paths")
-
-BACKEND_PATH = str(_find_backend_path())
-if BACKEND_PATH not in sys.path:
-    sys.path.insert(0, BACKEND_PATH)
-
-from fastapi import FastAPI, UploadFile, File, HTTPException, Query
-from fastapi.middleware.cors import CORSMiddleware
-from app.models.model_manager import get_model_manager, MODEL_CONFIGS, ENSEMBLE_CONFIGS
-
-app = FastAPI(
-    title="NeuroHealth Alzheimer MRI Space",
-    version="1.0.0",
-    description="FastAPI wrapper for Alzheimer's MRI models stored in backend/saved_models.",
-)
-
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-model_manager = get_model_manager()
-AD_KEYS = {"ad_dn121", "ad_dn169", "ad_dn201", "ad_homogeneous"}
-
-
+import torch
+import torchvision.models as models
+import torch.nn as nn
+from fastapi import FastAPI, UploadFile, File, HTTPException
+from PIL import Image
+import io
+import torchvision.transforms as transforms
+
+app = FastAPI(title="Alzheimer MRI API")
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+# -----------------------------
+# MODEL LOADER
+# -----------------------------
+def load_model(path, num_classes=3):
+    model = models.densenet121(weights=None)
+    model.classifier = nn.Linear(model.classifier.in_features, num_classes)
+    model.load_state_dict(torch.load(path, map_location=DEVICE))
+    model.to(DEVICE)
+    model.eval()
+    return model
+
+
+# -----------------------------
+# LOAD ALL MODELS
+# -----------------------------
+MODELS = {
+    "ad_dn121": load_model("alzheimers_densenet121.pth"),
+    "ad_dn169": load_model("alzheimers_densenet169.pth"),
+    "ad_dn201": load_model("alzheimers_densenet201.pth"),
+}
+
+
+# -----------------------------
+# IMAGE TRANSFORM
+# -----------------------------
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+])
+
+
+# -----------------------------
+# PREDICTION FUNCTION
+# -----------------------------
+def predict(model, image_bytes):
+    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+    image = transform(image).unsqueeze(0).to(DEVICE)
+
+    with torch.no_grad():
+        outputs = model(image)
+        probs = torch.softmax(outputs, dim=1)
+
+    return probs[0].tolist()
+
+
+# -----------------------------
+# ROUTES
+# -----------------------------
 @app.get("/")
-async def root():
-    return {
-        "message": "NeuroHealth Alzheimer MRI Space",
-        "supported_endpoints": ["/health", "/models", "/predict/image"],
-    }
+def home():
+    return {"message": "Alzheimer MRI API is running"}
 
 
-@app.get("/health")
-async def health():
-    statuses = model_manager.get_model_status()
-    return {
-        "status": "ok",
-        "model_status": {k: v for k, v in statuses.items() if k in AD_KEYS},
-        "available_models": [m["key"] for m in model_manager.get_available_models("alzheimers")],
-    }
-
-
-@app.get("/models")
-async def get_models():
-    models = model_manager.get_available_models("alzheimers")
-    models.append({
-        "key": "ad_homogeneous",
-        "name": "Alzheimer's MRI Homogeneous Ensemble (DenseNet 121+169+201)",
-        "condition": "alzheimers",
-        "imaging_type": "mri",
-        "type": "ensemble",
-    })
-    return {"models": models}
-
-
-@app.post("/predict/image")
-async def predict_image(
-    model_key: str = Query(..., description="Model key such as ad_dn121 or ad_homogeneous"),
-    file: UploadFile = File(...),
-):
-    if model_key not in AD_KEYS:
-        raise HTTPException(status_code=400, detail=f"Unknown Alzheimer MRI model key: {model_key}")
+@app.post("/predict")
+async def predict_image(model_key: str, file: UploadFile = File(...)):
 
-    image_bytes = await file.read()
-    filename = file.filename or ""
-    config = MODEL_CONFIGS.get(model_key, {})
+    if model_key not in MODELS:
+        raise HTTPException(status_code=400, detail="Invalid model key")
 
-    if model_key == "ad_homogeneous":
-        result = model_manager.predict_ensemble(model_key, image_bytes, filename)
-    else:
-        result = model_manager.predict_image(model_key, image_bytes, filename)
+    image_bytes = await file.read()
 
-    if "error" in result:
-        raise HTTPException(status_code=400, detail=result["error"])
+    result = predict(MODELS[model_key], image_bytes)
 
-    return result
+    return {
+        "model": model_key,
+        "prediction": result
+    }