Upload model_manager.py

model_manager.py

@@ -0,0 +1,512 @@
+"""
+Model Manager - Handles multiple CNN models for Alzheimer's (MRI) and
+Parkinson's (DaTscan) classification, with graceful fallback when models
+are not available.
+
+PD imaging = DaTscan ONLY:
+  - densenet121_parkinsonsDATSCAN.keras  (Keras, 2-class)
+  - parkinsons_densenet169DATSCAN.keras  (Keras, 2-class)
+  - parkinsons_densenet201DATSCAN.keras  (Keras, 2-class)
+  - parkinsons_3dcnnDATSCAN.pth         (PyTorch 3D CNN, 2-class)
+
+AD imaging = MRI:
+  - alzheimers_densenet121.pth
+  - alzheimers_densenet169.pth
+  - alzheimers_densenet201.pth
+"""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import Dict, Optional, Tuple, List
+import numpy as np
+import torch
+import torch.nn as nn
+import torchvision.models as tv
+from PIL import Image
+import torchvision.transforms as transforms
+from io import BytesIO
+
+logger = logging.getLogger("app.models.model_manager")
+
+# ── AD MRI Model configurations ───────────────────────────────────────────────
+MODEL_CONFIGS = {
+    # Alzheimer's MRI models (PyTorch .pth)
+    "ad_dn121": {
+        "name": "Alzheimer's DenseNet121 (MRI)",
+        "condition": "alzheimers",
+        "imaging_type": "mri",
+        "architecture": "densenet121",
+        "framework": "pytorch",
+        "num_classes": 4,
+        "filename": "alzheimers_densenet121.pth",
+        "class_names": ["Mild Demented", "Moderate Demented", "Non Demented", "Very Mild Demented"],
+    },
+    "ad_dn169": {
+        "name": "Alzheimer's DenseNet169 (MRI)",
+        "condition": "alzheimers",
+        "imaging_type": "mri",
+        "architecture": "densenet169",
+        "framework": "pytorch",
+        "num_classes": 4,
+        "filename": "alzheimers_densenet169.pth",
+        "class_names": ["Mild Demented", "Moderate Demented", "Non Demented", "Very Mild Demented"],
+    },
+    "ad_dn201": {
+        "name": "Alzheimer's DenseNet201 (MRI)",
+        "condition": "alzheimers",
+        "imaging_type": "mri",
+        "architecture": "densenet201",
+        "framework": "pytorch",
+        "num_classes": 4,
+        "filename": "alzheimers_densenet201.pth",
+        "class_names": ["Mild Demented", "Moderate Demented", "Non Demented", "Very Mild Demented"],
+    },
+
+    # Parkinson's DaTscan models — prefer retrained .pth, fall back to .keras
+    "pd_datscan_dn121": {
+        "name": "Parkinson's DaTscan DenseNet121",
+        "condition": "parkinsons",
+        "imaging_type": "datscan",
+        "architecture": "densenet121",
+        "framework": "pytorch",
+        "num_classes": 2,
+        # Retrained .pth takes priority; .keras kept as fallback filename
+        "filename": "parkinsons_densenet121.pth",
+        "filename_fallback": "densenet121_parkinsonsDATSCAN.keras",
+        "class_names": ["No Parkinson's", "Parkinson's Disease"],
+    },
+    "pd_datscan_dn169": {
+        "name": "Parkinson's DaTscan DenseNet169",
+        "condition": "parkinsons",
+        "imaging_type": "datscan",
+        "architecture": "densenet169",
+        "framework": "pytorch",
+        "num_classes": 2,
+        "filename": "parkinsons_densenet169.pth",
+        "filename_fallback": "parkinsons_densenet169DATSCAN.keras",
+        "class_names": ["No Parkinson's", "Parkinson's Disease"],
+    },
+    "pd_datscan_dn201": {
+        "name": "Parkinson's DaTscan DenseNet201",
+        "condition": "parkinsons",
+        "imaging_type": "datscan",
+        "architecture": "densenet201",
+        "framework": "pytorch",
+        "num_classes": 2,
+        "filename": "parkinsons_densenet201.pth",
+        "filename_fallback": "parkinsons_densenet201DATSCAN.keras",
+        "class_names": ["No Parkinson's", "Parkinson's Disease"],
+    },
+    "pd_datscan_3dcnn": {
+        "name": "Parkinson's DaTscan 3D CNN",
+        "condition": "parkinsons",
+        "imaging_type": "datscan",
+        "architecture": "3dcnn",
+        "framework": "pytorch",
+        "num_classes": 2,
+        "filename": "parkinsons_3dcnnDATSCAN.pth",
+        "class_names": ["No Parkinson's", "Parkinson's Disease"],
+        "input_3d": True,
+    },
+}
+
+# ── Ensemble configurations ────────────────────────────────────────────────────
+ENSEMBLE_CONFIGS = {
+    "ad_homogeneous": {
+        "name": "Alzheimer's MRI Homogeneous Ensemble (DenseNet 121+169+201)",
+        "condition": "alzheimers",
+        "imaging_type": "mri",
+        "models": ["ad_dn121", "ad_dn169", "ad_dn201"],
+        "weights": [0.4, 0.3, 0.3],
+    },
+    "pd_datscan_ensemble": {
+        "name": "Parkinson's DaTscan Ensemble (DenseNet 121+169+201)",
+        "condition": "parkinsons",
+        "imaging_type": "datscan",
+        "models": ["pd_datscan_dn121", "pd_datscan_dn169", "pd_datscan_dn201"],
+        "weights": [0.4, 0.3, 0.3],
+    },
+}
+
+# Accepted DaTscan file extensions
+DATSCAN_EXTENSIONS = {".nii", ".gz", ".dcm", ".png", ".jpg", ".jpeg"}
+
+
+class ModelManager:
+    """Manages MRI (AD) and DaTscan (PD) models with graceful fallback."""
+
+    def __init__(self):
+        self.models_dir = Path(__file__).resolve().parent.parent.parent / "saved_models"
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.loaded_models: Dict[str, Optional[object]] = {}
+        self.model_status: Dict[str, str] = {}
+        self.image_transform = self._get_image_transform()
+        self._initialize_models()
+
+    def _get_image_transform(self):
+        return transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.ToTensor(),
+            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
+        ])
+
+    # ── PyTorch model builder ──────────────────────────────────────────────────
+
+    def _build_pytorch_model(self, config: dict) -> nn.Module:
+        arch = config["architecture"]
+        num_classes = config["num_classes"]
+
+        if arch == "densenet121":
+            model = tv.densenet121(weights=None)
+            in_features = model.classifier.in_features
+            model.classifier = nn.Sequential(
+                nn.Linear(in_features, 256), nn.ReLU(), nn.Dropout(0.3),
+                nn.Linear(256, num_classes),
+            )
+        elif arch == "densenet169":
+            model = tv.densenet169(weights=None)
+            in_features = model.classifier.in_features
+            model.classifier = nn.Sequential(
+                nn.Linear(in_features, 256), nn.ReLU(), nn.Dropout(0.3),
+                nn.Linear(256, num_classes),
+            )
+        elif arch == "densenet201":
+            model = tv.densenet201(weights=None)
+            in_features = model.classifier.in_features
+            model.classifier = nn.Sequential(
+                nn.Linear(in_features, 256), nn.ReLU(), nn.Dropout(0.3),
+                nn.Linear(256, num_classes),
+            )
+        elif arch == "3dcnn":
+            model = self._build_3dcnn(num_classes)
+        else:
+            raise ValueError(f"Unsupported pytorch architecture: {arch}")
+        return model
+
+    def _build_3dcnn(self, num_classes: int = 2) -> nn.Module:
+        """Simple 3D CNN for DaTscan volumetric input."""
+        class Simple3DCNN(nn.Module):
+            def __init__(self, n_classes):
+                super().__init__()
+                self.features = nn.Sequential(
+                    nn.Conv3d(1, 32, 3, padding=1), nn.BatchNorm3d(32), nn.ReLU(),
+                    nn.MaxPool3d(2),
+                    nn.Conv3d(32, 64, 3, padding=1), nn.BatchNorm3d(64), nn.ReLU(),
+                    nn.MaxPool3d(2),
+                    nn.Conv3d(64, 128, 3, padding=1), nn.BatchNorm3d(128), nn.ReLU(),
+                    nn.AdaptiveAvgPool3d((4, 4, 4)),
+                )
+                self.classifier = nn.Sequential(
+                    nn.Flatten(),
+                    nn.Linear(128 * 4 * 4 * 4, 256), nn.ReLU(), nn.Dropout(0.4),
+                    nn.Linear(256, n_classes),
+                )
+            def forward(self, x):
+                return self.classifier(self.features(x))
+        return Simple3DCNN(num_classes)
+
+    # ── Keras model loader ─────────────────────────────────────────────────────
+
+    def _load_keras_model(self, model_key: str) -> Tuple[Optional[object], str]:
+        config = MODEL_CONFIGS[model_key]
+        model_path = self.models_dir / config["filename"]
+        if not model_path.exists():
+            logger.warning("Keras model file not found: %s", model_path)
+            return None, "Model file not found"
+        try:
+            import os
+            os.environ["TF_USE_LEGACY_KERAS"] = "1"
+            import tensorflow as tf
+            model = tf.keras.models.load_model(str(model_path), compile=False)
+            logger.info("Loaded Keras model: %s", config["name"])
+            return model, "Active"
+        except Exception as e:
+            logger.error("Failed to load Keras model %s: %s", config["name"], e)
+            return None, f"Load error: {str(e)}"
+
+    # ── PyTorch model loader ───────────────────────────────────────────────────
+
+    def _load_pytorch_model(self, model_key: str) -> Tuple[Optional[nn.Module], str]:
+        config = MODEL_CONFIGS[model_key]
+        model_path = self.models_dir / config["filename"]
+
+        # If primary .pth not found, try fallback (old .keras → skip, just report missing)
+        if not model_path.exists():
+            fallback = config.get("filename_fallback")
+            if fallback:
+                fallback_path = self.models_dir / fallback
+                if fallback_path.exists() and fallback_path.suffix in (".keras", ".h5"):
+                    # Keras fallback — delegate to keras loader
+                    return self._load_keras_model_from_path(config, fallback_path)
+            logger.warning("Model file not found: %s", model_path)
+            return None, "Model file not found"
+        try:
+            model = self._build_pytorch_model(config)
+            state_dict = torch.load(str(model_path), map_location=self.device)
+            if isinstance(state_dict, dict) and "model_state_dict" in state_dict:
+                state_dict = state_dict["model_state_dict"]
+            # strict=False allows loading models whose classifier head differs slightly
+            model.load_state_dict(state_dict, strict=False)
+            model.to(self.device)
+            model.eval()
+            logger.info("Loaded PyTorch model: %s", config["name"])
+            return model, "Active"
+        except Exception as e:
+            logger.error("Failed to load PyTorch model %s: %s", config["name"], e)
+            return None, f"Load error: {str(e)}"
+
+    def _load_keras_model_from_path(self, config: dict, model_path: Path) -> Tuple[Optional[object], str]:
+        try:
+            import os
+            os.environ["TF_USE_LEGACY_KERAS"] = "1"
+            import tensorflow as tf
+            model = tf.keras.models.load_model(str(model_path), compile=False)
+            logger.info("Loaded Keras fallback model: %s", config["name"])
+            return model, "Active (Keras fallback)"
+        except Exception as e:
+            logger.error("Failed to load Keras fallback %s: %s", config["name"], e)
+            return None, f"Load error: {str(e)}"
+
+    def _initialize_models(self):
+        logger.info("Initializing model manager (AD-MRI + PD-DaTscan)...")
+        for model_key, config in MODEL_CONFIGS.items():
+            # All models now use PyTorch; Keras fallback handled inside _load_pytorch_model
+            model, status = self._load_pytorch_model(model_key)
+            self.loaded_models[model_key] = model
+            self.model_status[model_key] = status
+        logger.info("Model initialization complete")
+
+    # ── Availability ───────────────────────────────────────────────────────────
+
+    def get_available_models(self, condition: str = None) -> List[dict]:
+        available = []
+        for model_key, config in MODEL_CONFIGS.items():
+            if condition and config["condition"] != condition:
+                continue
+            if self.model_status.get(model_key) == "Active":
+                available.append({
+                    "key": model_key,
+                    "name": config["name"],
+                    "condition": config["condition"],
+                    "imaging_type": config.get("imaging_type", "mri"),
+                    "architecture": config["architecture"],
+                    "framework": config.get("framework", "pytorch"),
+                    "status": "Active",
+                })
+        return available
+
+    def get_model_status(self) -> Dict[str, str]:
+        return self.model_status.copy()
+
+    # ── PyTorch image prediction (AD MRI) ─────────────────────────────────────
+
+    def predict_image(self, model_key: str, image_bytes: bytes, filename: str = "") -> dict:
+        """Make prediction using a PyTorch model on standard image bytes."""
+        if model_key not in MODEL_CONFIGS:
+            return {"error": f"Unknown model: {model_key}"}
+
+        config = MODEL_CONFIGS[model_key]
+
+        # Validate DaTscan extensions
+        if config.get("imaging_type") == "datscan" and filename:
+            ext = Path(filename).suffix.lower()
+            # .nii.gz has compound suffix
+            if filename.endswith(".nii.gz"):
+                ext = ".nii.gz"
+            if ext not in DATSCAN_EXTENSIONS and ext != ".nii.gz":
+                return {"error": f"Invalid file type '{ext}' for DaTscan analysis. Accepted: .nii, .nii.gz, .dcm, .png, .jpg"}
+
+        # 3D CNN needs special handling
+        if config.get("input_3d"):
+            return self.predict_3dcnn(model_key, image_bytes, filename)
+
+        model = self.loaded_models.get(model_key)
+        if model is None:
+            return {"error": f"Model {model_key} is not available"}
+
+        try:
+            image = Image.open(BytesIO(image_bytes)).convert("RGB")
+            inputs = self.image_transform(image).unsqueeze(0).to(self.device)
+
+            with torch.no_grad():
+                outputs = model(inputs)
+                probs = torch.softmax(outputs, dim=1)
+                pred_class = torch.argmax(probs, dim=1).item()
+                confidence = float(probs[0][pred_class].item())
+
+            return {
+                "model_key": model_key,
+                "model_name": config["name"],
+                "condition": config["condition"],
+                "imaging_type": config.get("imaging_type", "mri"),
+                "prediction": pred_class,
+                "confidence": confidence,
+                "class_name": config["class_names"][pred_class],
+                "all_probabilities": {
+                    cn: float(p) for cn, p in zip(config["class_names"], probs[0].cpu().numpy())
+                },
+                "status": "success",
+            }
+        except Exception as e:
+            logger.error("PyTorch prediction failed for %s: %s", model_key, e)
+            return {"error": f"Prediction failed: {str(e)}"}
+
+    # ── Keras image prediction (PD DaTscan DenseNet) ──────────────────────────
+
+    def predict_keras_image(self, model_key: str, image_bytes: bytes, filename: str = "") -> dict:
+        """Run a Keras DaTscan model on 2D image/slice bytes."""
+        if model_key not in MODEL_CONFIGS:
+            return {"error": f"Unknown model: {model_key}"}
+
+        config = MODEL_CONFIGS[model_key]
+
+        # Extension check
+        if filename:
+            ext = Path(filename).suffix.lower()
+            fname_lower = filename.lower()
+            if fname_lower.endswith(".nii.gz"):
+                ext = ".nii.gz"
+            if ext not in DATSCAN_EXTENSIONS:
+                return {"error": f"Invalid file type '{ext}' for DaTscan. Accepted: .nii, .nii.gz, .dcm, .png, .jpg"}
+
+        model = self.loaded_models.get(model_key)
+        if model is None:
+            return {"error": f"Keras model {model_key} is not available"}
+
+        try:
+            from app.preprocessing.datscan_preprocessor import DaTscanPreprocessor
+            preprocessor = DaTscanPreprocessor()
+            img_array = preprocessor.preprocess_2d(image_bytes, filename)  # (224, 224, 3) float32
+
+            import numpy as _np
+            batch = _np.expand_dims(img_array, 0)  # (1, 224, 224, 3)
+            preds = model.predict(batch, verbose=0)  # (1, num_classes)
+            probs = preds[0]
+            pred_class = int(_np.argmax(probs))
+            confidence = float(probs[pred_class])
+
+            return {
+                "model_key": model_key,
+                "model_name": config["name"],
+                "condition": config["condition"],
+                "imaging_type": "datscan",
+                "prediction": pred_class,
+                "confidence": confidence,
+                "class_name": config["class_names"][pred_class],
+                "all_probabilities": {
+                    cn: float(p) for cn, p in zip(config["class_names"], probs)
+                },
+                "status": "success",
+            }
+        except Exception as e:
+            logger.error("Keras DaTscan prediction failed for %s: %s", model_key, e)
+            return {"error": f"DaTscan prediction failed: {str(e)}"}
+
+    # ── 3D CNN prediction (PD DaTscan volumetric) ─────────────────────────────
+
+    def predict_3dcnn(self, model_key: str, volume_bytes: bytes, filename: str = "") -> dict:
+        """Run the 3D CNN on a NIfTI volume (.nii or .nii.gz required)."""
+        if model_key not in MODEL_CONFIGS:
+            return {"error": f"Unknown model: {model_key}"}
+
+        config = MODEL_CONFIGS[model_key]
+        fname_lower = (filename or "").lower()
+        if not (fname_lower.endswith(".nii") or fname_lower.endswith(".nii.gz")):
+            return {"error": "3D CNN requires a NIfTI file (.nii or .nii.gz)."}
+
+        model = self.loaded_models.get(model_key)
+        if model is None:
+            return {"error": f"3D CNN model {model_key} is not available"}
+
+        try:
+            from app.preprocessing.datscan_preprocessor import DaTscanPreprocessor
+            preprocessor = DaTscanPreprocessor()
+            volume_tensor = preprocessor.preprocess_3d(volume_bytes, filename)  # (1, 1, D, H, W)
+            volume_tensor = volume_tensor.to(self.device)
+
+            with torch.no_grad():
+                outputs = model(volume_tensor)
+                probs = torch.softmax(outputs, dim=1)
+                pred_class = int(torch.argmax(probs, dim=1).item())
+                confidence = float(probs[0][pred_class].item())
+
+            return {
+                "model_key": model_key,
+                "model_name": config["name"],
+                "condition": config["condition"],
+                "imaging_type": "datscan",
+                "prediction": pred_class,
+                "confidence": confidence,
+                "class_name": config["class_names"][pred_class],
+                "all_probabilities": {
+                    cn: float(p) for cn, p in zip(config["class_names"], probs[0].cpu().numpy())
+                },
+                "status": "success",
+            }
+        except Exception as e:
+            logger.error("3D CNN prediction failed for %s: %s", model_key, e)
+            return {"error": f"3D CNN prediction failed: {str(e)}"}
+
+    # ── Ensemble prediction ────────────────────────────────────────────────────
+
+    def predict_ensemble(self, ensemble_key: str, image_bytes: bytes, filename: str = "") -> dict:
+        if ensemble_key not in ENSEMBLE_CONFIGS:
+            return {"error": f"Unknown ensemble: {ensemble_key}"}
+
+        ensemble_config = ENSEMBLE_CONFIGS[ensemble_key]
+        model_predictions = []
+        weights = ensemble_config["weights"]
+
+        for model_key in ensemble_config["models"]:
+            result = self.predict_image(model_key, image_bytes, filename)
+            if "error" not in result:
+                model_predictions.append(result)
+
+        if not model_predictions:
+            return {"error": "No models available in ensemble"}
+
+        if len(weights) != len(model_predictions):
+            weights = [1.0 / len(model_predictions)] * len(model_predictions)
+
+        combined_probs: Dict[str, float] = {}
+        total_weight = 0.0
+        for pred, weight in zip(model_predictions, weights):
+            total_weight += weight
+            for class_name, prob in pred["all_probabilities"].items():
+                combined_probs[class_name] = combined_probs.get(class_name, 0) + prob * weight
+
+        for cn in combined_probs:
+            combined_probs[cn] /= total_weight
+
+        final_class = max(combined_probs, key=lambda x: combined_probs[x])
+        final_confidence = combined_probs[final_class]
+        first_config = MODEL_CONFIGS[ensemble_config["models"][0]]
+
+        return {
+            "ensemble_key": ensemble_key,
+            "ensemble_name": ensemble_config["name"],
+            "condition": ensemble_config["condition"],
+            "imaging_type": ensemble_config.get("imaging_type", "mri"),
+            "prediction": first_config["class_names"].index(final_class),
+            "confidence": final_confidence,
+            "class_name": final_class,
+            "all_probabilities": combined_probs,
+            "model_contributions": [
+                {"model": p["model_name"], "weight": w, "confidence": p["confidence"]}
+                for p, w in zip(model_predictions, weights)
+            ],
+            "status": "success",
+        }
+
+
+# ── Singleton ──────────────────────────────────────────────────────────────────
+_model_manager: Optional[ModelManager] = None
+
+
+def get_model_manager() -> ModelManager:
+    global _model_manager
+    if _model_manager is None:
+        _model_manager = ModelManager()
+    return _model_manager