Update app.py

app.py

@@ -6,22 +6,18 @@ from PIL import Image
 import io
 import torchvision.transforms as transforms
 
+app = FastAPI(title="Parkinson DATSCAN Ensemble API")
+
 # =========================
-# APP
+# DEVICE
 # =========================
-app = FastAPI(title="Alzheimer Ensemble API")
-
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print("Using device:", DEVICE)
 
 # =========================
 # LABELS
 # =========================
-LABELS = [
-    "Mild Demented",
-    "Moderate Demented",
-    "Non Demented",
-    "Very Mild Demented"
-]
+LABELS = ["Control", "Prodromal", "Parkinsons"]
 
 # =========================
 # IMAGE TRANSFORM
@@ -29,76 +25,97 @@ LABELS = [
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
-    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
+    transforms.Normalize([0.5]*3, [0.5]*3)
 ])
 
 # =========================
-# MODEL ARCHITECTURE
+# MODEL BUILDER
 # =========================
-def build_model():
-    model = models.densenet121(weights=None)
-
-    model.classifier = nn.Sequential(
-        nn.Dropout(0.5),
-        nn.Linear(1024, 256),
-        nn.ReLU(),
-        nn.Dropout(0.3),
-        nn.Linear(256, 4)
-    )
+def build_densenet(name="121"):
+    if name == "121":
+        model = models.densenet121(weights=None)
+        in_f = 1024
+    elif name == "169":
+        model = models.densenet169(weights=None)
+        in_f = 1664
+    else:
+        model = models.densenet201(weights=None)
+        in_f = 1920
 
+    model.classifier = nn.Linear(in_f, len(LABELS))
     return model
 
 # =========================
-# LOAD MODEL (FIXED)
+# SAFE LOADER (IMPORTANT FIX)
 # =========================
-def load_model(path):
-    model = build_model()
+def load_model(path, name):
+    model = build_densenet(name)
+
+    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 checkpoint formats safely
-    if isinstance(state, dict) and "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}")
+
+    except Exception as e:
+        print(f"Failed loading {path}: {e}")
 
     model.to(DEVICE)
     model.eval()
-
     return model
 
 # =========================
-# LOAD ALL MODELS (REAL FIX HERE)
+# LOAD MODELS
 # =========================
-model_121 = load_model("alzheimers_densenet121.pth")
-model_169 = load_model("alzheimers_densenet169.pth")
-model_201 = load_model("alzheimers_densenet201.pth")
+model_121 = load_model("densenet121.pth", "121")
+model_169 = load_model("densenet169.pth", "169")
+model_201 = load_model("densenet201.pth", "201")
 
 # =========================
 # IMAGE PROCESSING
 # =========================
 def process_image(image_bytes):
     image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-    image = transform(image).unsqueeze(0)
-    return image.to(DEVICE)
+    tensor = transform(image).unsqueeze(0).to(DEVICE)
+    return tensor
 
 # =========================
-# SINGLE MODEL PREDICT
+# PREDICTION FUNCTION (FIXED)
 # =========================
 def predict(model, image_tensor):
     with torch.no_grad():
         out = model(image_tensor)
         probs = torch.softmax(out, dim=1)[0]
+
         conf, cls = torch.max(probs, dim=0)
 
-    return {
-        "class_id": int(cls.item()),
-        "class_name": LABELS[cls.item()],
-        "confidence": round(float(conf.item()), 4),
-        "probabilities": {
-            LABELS[i]: round(float(probs[i].item()), 4)
-            for i in range(len(LABELS))
+        return {
+            "class_id": int(cls.item()),
+            "class_name": LABELS[int(cls.item())],
+            "confidence": float(conf.item()),
+            "probabilities": {
+                LABELS[i]: float(probs[i].item())
+                for i in range(len(LABELS))
+            }
         }
+
+# =========================
+# ROOT
+# =========================
+@app.get("/")
+def home():
+    return {
+        "status": "running",
+        "models": ["121", "169", "201"],
+        "endpoints": ["/predict/121", "/predict/169", "/predict/201", "/predict/ensemble"]
     }
 
 # =========================
@@ -120,7 +137,7 @@ async def predict_201(file: UploadFile = File(...)):
     return {"model": "densenet201", **predict(model_201, img)}
 
 # =========================
-# ENSEMBLE (FIXED + CLEAN)
+# ENSEMBLE (FIXED)
 # =========================
 @app.post("/predict/ensemble")
 async def ensemble(file: UploadFile = File(...)):
@@ -144,30 +161,11 @@ async def ensemble(file: UploadFile = File(...)):
 
     return {
         "final_prediction": final_class,
-        "final_confidence": round(avg_probs[final_class], 4),
-
-        "averaged_probabilities": avg_probs,
-
-        "model_outputs": {
-            "densenet121": r1,
-            "densenet169": r2,
-            "densenet201": r3
-        }
-    }
-
-# =========================
-# ROOT
-# =========================
-@app.get("/")
-def home():
-    return {
-        "status": "running",
-        "device": str(DEVICE),
-        "models_loaded": True,
-        "endpoints": [
-            "/predict/121",
-            "/predict/169",
-            "/predict/201",
-            "/predict/ensemble"
-        ]
+        "final_confidence": avg_probs[final_class],
+        "models": {
+            "121": r1,
+            "169": r2,
+            "201": r3
+        },
+        "averaged_probabilities": avg_probs
     }