Fix L1 gatekeeper ResNet-50 compatibility, integrate L1-L2-L3 pipeline and return Base64 GradCAMs

app.py

@@ -1,53 +1,30 @@
-import io
 import os
+import io
 import uuid
 import tempfile
 import subprocess
-import torch
 from fastapi import FastAPI, File, UploadFile
 from fastapi.responses import JSONResponse
-from PIL import Image
-from torchvision import transforms
 
-from models.level1_gatekeeper import GatekeeperModel
-from models.level2_router import DiseaseRouterModel
+from scripts.inference_pipeline import OCTInferencePipeline
 
 app = FastAPI(title="OCT Image Classification Pipeline")
 
-# Helper to load state dict properly from checkpoints
-def load_checkpoint(model, path):
-    checkpoint = torch.load(path, map_location='cpu')
-    if 'model_state_dict' in checkpoint:
-        model.load_state_dict(checkpoint['model_state_dict'])
-    else:
-        model.load_state_dict(checkpoint)
-
-# Load models globally
-print("Loading Level 1 Gatekeeper...")
-l1_model = GatekeeperModel(pretrained=False, freeze_backbone=False)
-load_checkpoint(l1_model, 'weights/level1.pth')
-l1_model.eval()
-
-print("Loading Level 2 Router...")
-l2_model = DiseaseRouterModel(pretrained=False, freeze_backbone=False)
-load_checkpoint(l2_model, 'weights/level2.pth')
-l2_model.eval()
-
-# Standard ImageNet transforms for EfficientNet
-transform = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-])
-
-L1_CLASSES = {0: "NORMAL", 1: "ABNORMAL"}
-L2_CLASSES = {
-    0: "Macular_Degeneration",
-    1: "Diabetic_Complications",
-    2: "Vascular_Occlusions",
-    3: "Fluid_Accumulation",
-    4: "Structural_Issues"
+# Initialize pipeline once on startup
+print("Initializing OCT Inference Pipeline...")
+l3_ckpts = {
+    "Macular": "weights/level3_macular.pth",
+    "Diabetic": "weights/level3_diabetic.pth",
+    "Vascular": "weights/level3_vascular.pth",
+    "Fluid": "weights/level3_fluid.pth",
+    "Structural": "weights/level3_structural.pth"
 }
+pipeline = OCTInferencePipeline(
+    l1_ckpt="weights/level1.pth",
+    l2_ckpt="weights/level2.pth",
+    l3_ckpts=l3_ckpts,
+    device="cpu" # Space typically uses CPU unless upgraded
+)
 
 @app.get("/")
 def greet_json():
@@ -57,9 +34,8 @@ def greet_json():
 async def predict_image(file: UploadFile = File(...)):
     try:
         contents = await file.read()
-        image = Image.open(io.BytesIO(contents)).convert("RGB")
         
-        # Save to bucket
+        # Save to bucket for storage purposes (optional analytics)
         ext = os.path.splitext(file.filename)[1]
         if not ext:
             ext = ".png"
@@ -70,38 +46,25 @@ async def predict_image(file: UploadFile = File(...)):
             tmp_path = tmp.name
             
         # Upload to HF Bucket via CLI
-        # The space will use the HF_TOKEN environment variable automatically
         bucket_url = f"hf://buckets/NMundhra/OCT-Image-Classifier-Model-storage/uploads/{unique_filename}"
         subprocess.Popen(["hf", "buckets", "cp", tmp_path, bucket_url])
-        # Popen runs it asynchronously so it doesn't block the API response
         
-    except Exception as e:
-        return JSONResponse(status_code=400, content={"error": f"Invalid image file or upload failed: {str(e)}"})
-
-    input_tensor = transform(image).unsqueeze(0)
-
-    with torch.no_grad():
-        # Level 1 Inference
-        l1_logits = l1_model(input_tensor)
-        l1_probs = torch.softmax(l1_logits, dim=1)[0]
-        l1_pred_idx = torch.argmax(l1_probs).item()
-        l1_confidence = l1_probs[l1_pred_idx].item()
-        l1_label = L1_CLASSES[l1_pred_idx]
+        # Run inference pipeline
+        results = pipeline.predict(tmp_path, gradcam=True)
+        
+        if "error" in results:
+            return JSONResponse(status_code=400, content=results)
 
-        result = {
-            "level1_prediction": l1_label,
-            "level1_confidence": float(l1_confidence)
+        # Re-format to match the expected legacy format but with new L3 + GradCAMs
+        return {
+            "level1_prediction": results.get("Level1", {}).get("prediction"),
+            "level1_confidence": results.get("Level1", {}).get("confidence"),
+            "level2_prediction": results.get("Level2", {}).get("prediction"),
+            "level2_confidence": results.get("Level2", {}).get("confidence"),
+            "level3_prediction": results.get("Level3", {}).get("prediction"),
+            "level3_confidence": results.get("Level3", {}).get("confidence"),
+            "final_diagnosis": results.get("Final_Diagnosis"),
+            "gradcams": results.get("gradcams", {})
         }
-
-        # If abnormal, route to Level 2
-        if l1_label == "ABNORMAL":
-            l2_logits = l2_model(input_tensor)
-            l2_probs = torch.softmax(l2_logits, dim=1)[0]
-            l2_pred_idx = torch.argmax(l2_probs).item()
-            l2_confidence = l2_probs[l2_pred_idx].item()
-            l2_label = L2_CLASSES[l2_pred_idx]
-
-            result["level2_prediction"] = l2_label
-            result["level2_confidence"] = float(l2_confidence)
-
-    return result
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"error": f"Internal server error: {str(e)}"})

data/transforms.py

@@ -0,0 +1,288 @@
+"""
+data/transforms.py
+
+Augmentation pipelines for the OCT hierarchical classification pipeline.
+
+Resolution Strategy (from architectural directives):
+  Level 1 (Gatekeeper):   224×224 — maximum throughput for binary screening.
+  Level 2 (Router):       224×224 — consistent feature space with L1.
+  Level 3 (Specialists):  384×384 — fine-grained structural detail for
+                           CNV vs DRUSEN, RAO vs RVO, etc.
+
+Pipeline Variants:
+  - Standard Train:  Random crop/flip/rotation + ColorJitter + GaussianBlur +
+                     RandomErasing. Used for L1, L2, L3_Macular, L3_Diabetic.
+  - Heavy Train:     Adds RandomAffine + stronger erasing. Used for
+                     extreme minority L3 specialists (Vascular, Fluid, Structural)
+                     where RAO has only 22 samples and CSR has 102.
+  - Val/Test:        Deterministic resize + CenterCrop + normalize only.
+
+All pipelines use ImageNet mean/std for pretrained backbone compatibility.
+"""
+
+import numpy as np
+from torchvision import transforms
+import cv2
+
+# ── ImageNet statistics ───────────────────────────────────────────────────────
+IMAGENET_MEAN = [0.485, 0.456, 0.406]
+IMAGENET_STD  = [0.229, 0.224, 0.225]
+
+# ── Resolution constants ──────────────────────────────────────────────────────
+RES_L1_L2: int = 224   # Level 1 & 2 input resolution
+RES_L3:    int = 384   # Level 3 specialist input resolution
+
+# Intermediate crop sizes (resize target before random/center crop)
+_CROP_L1_L2: int = 256
+_CROP_L3:    int = 416
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# CLAHE Preprocessing
+# ──────────────────────────────────────────────────────────────────────────────
+
+class CLAHETransform:
+    """
+    Contrast Limited Adaptive Histogram Equalization for OCT images.
+
+    Applied per-image BEFORE resize/crop to normalise brightness and local
+    contrast variation across different OCT scanner manufacturers
+    (Zeiss, Heidelberg, Topcon, etc.).
+
+    Without this step, the model may learn scanner-specific intensity
+    distributions rather than pathology — a form of shortcut learning that
+    degrades performance on unseen devices.
+
+    Applied identically at train, val, and test time — this is NOT an
+    augmentation, it is a deterministic preprocessing step.
+
+    Args:
+        clip_limit:  Contrast clip threshold. 2.0 is standard for OCT.
+                     Higher values = more contrast, more noise amplification.
+        tile_grid:   Size of the adaptive tile grid. (8, 8) is standard.
+    """
+
+    def __init__(
+        self,
+        clip_limit: float = 2.0,
+        tile_grid: tuple = (8, 8),
+    ) -> None:
+        self.clip_limit = clip_limit
+        self.tile_grid = tile_grid
+        self._clahe = None
+
+    def __call__(self, img) -> "PIL.Image.Image":
+        from PIL import Image as PILImage
+        if self._clahe is None:
+            self._clahe = cv2.createCLAHE(
+                clipLimit=self.clip_limit,
+                tileGridSize=self.tile_grid,
+            )
+        # Convert to numpy grayscale — OCT images carry most diagnostic
+        # information in luminance; colour channels are usually redundant
+        img_np = np.array(img.convert("L"), dtype=np.uint8)
+        equalized = self._clahe.apply(img_np)
+        # Stack to 3-channel RGB — required for ImageNet-pretrained backbones
+        rgb = np.stack([equalized, equalized, equalized], axis=-1)
+        return PILImage.fromarray(rgb, mode="RGB")
+
+
+# Shared instance used in all transform pipelines
+_CLAHE = CLAHETransform(clip_limit=2.0, tile_grid=(8, 8))
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Transform factory functions
+# ──────────────────────────────────────────────────────────────────────────────
+
+def get_train_transforms(resolution: int = RES_L1_L2) -> transforms.Compose:
+    """
+    Standard training augmentation pipeline.
+
+    Designed to:
+      - Increase geometric diversity (flip, rotate, crop).
+      - Simulate OCT scan artefacts (GaussianBlur, ColorJitter).
+      - Force the network to ignore local texture via RandomErasing.
+
+    Args:
+        resolution: Target output resolution (224 or 384).
+
+    Returns:
+        Composed torchvision transform.
+    """
+    crop_size = _CROP_L3 if resolution == RES_L3 else _CROP_L1_L2
+    return transforms.Compose([
+        _CLAHE,                          # Scanner normalisation (deterministic)
+        transforms.Resize(
+            crop_size,
+            interpolation=transforms.InterpolationMode.BICUBIC,
+        ),
+        transforms.RandomCrop(resolution),
+        transforms.RandomHorizontalFlip(p=0.5),
+        transforms.RandomVerticalFlip(p=0.2),
+        transforms.RandomRotation(degrees=15),
+        transforms.ColorJitter(
+            brightness=0.3,
+            contrast=0.3,
+            saturation=0.1,
+            hue=0.05,
+        ),
+        transforms.RandomApply(
+            [transforms.GaussianBlur(kernel_size=5, sigma=(0.1, 2.0))],
+            p=0.3,
+        ),
+        transforms.ToTensor(),
+        # RandomErasing after ToTensor (operates on tensor, not PIL image)
+        transforms.RandomErasing(
+            p=0.2,
+            scale=(0.02, 0.10),
+            ratio=(0.3, 3.3),
+            value="random",
+        ),
+        transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
+    ])
+
+
+def get_heavy_train_transforms(resolution: int = RES_L3) -> transforms.Compose:
+    """
+    Heavy augmentation pipeline for extreme minority classes.
+
+    Applied to L3_Vascular (RAO=22, RVO=101, MH=102), L3_Fluid (CSR=102),
+    and L3_Structural (ERM=155, VID=76) to maximise synthetic variation.
+
+    Adds on top of the standard pipeline:
+      - RandomAffine (translate, scale, shear)
+      - Stronger rotation (±30°)
+      - Stronger RandomErasing scale
+
+    Args:
+        resolution: Target output resolution (typically 384 for L3).
+    """
+    crop_size = _CROP_L3 if resolution == RES_L3 else _CROP_L1_L2
+    return transforms.Compose([
+        _CLAHE,                          # Scanner normalisation (deterministic)
+        transforms.Resize(
+            crop_size,
+            interpolation=transforms.InterpolationMode.BICUBIC,
+        ),
+        transforms.RandomCrop(resolution),
+        transforms.RandomHorizontalFlip(p=0.5),
+        transforms.RandomVerticalFlip(p=0.5),
+        transforms.RandomRotation(degrees=30),
+        transforms.RandomAffine(
+            degrees=20,
+            translate=(0.10, 0.10),
+            scale=(0.85, 1.15),
+            shear=10,
+            interpolation=transforms.InterpolationMode.BICUBIC,
+        ),
+        transforms.ColorJitter(
+            brightness=0.4,
+            contrast=0.4,
+            saturation=0.2,
+            hue=0.10,
+        ),
+        transforms.RandomApply(
+            [transforms.GaussianBlur(kernel_size=5, sigma=(0.1, 3.0))],
+            p=0.4,
+        ),
+        transforms.ToTensor(),
+        transforms.RandomErasing(
+            p=0.35,
+            scale=(0.02, 0.15),
+            ratio=(0.3, 3.3),
+            value="random",
+        ),
+        transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
+    ])
+
+
+def get_val_transforms(resolution: int = RES_L1_L2) -> transforms.Compose:
+    """
+    Deterministic validation/test pipeline (no augmentation).
+
+    Args:
+        resolution: Target output resolution (224 or 384).
+
+    Returns:
+        Composed torchvision transform.
+    """
+    crop_size = _CROP_L3 if resolution == RES_L3 else _CROP_L1_L2
+    return transforms.Compose([
+        _CLAHE,                          # Scanner normalisation — must match train pipeline
+        transforms.Resize(
+            crop_size,
+            interpolation=transforms.InterpolationMode.BICUBIC,
+        ),
+        transforms.CenterCrop(resolution),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
+    ])
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Registry — keyed by (mode, split)
+# ──────────────────────────────────────────────────────────────────────────────
+
+#: Complete transform registry. Access via :func:`get_transforms`.
+TRANSFORM_REGISTRY: dict = {
+    # Level 1 — 224px, standard augmentation
+    "level1": {
+        "train": get_train_transforms(RES_L1_L2),
+        "val":   get_val_transforms(RES_L1_L2),
+    },
+    # Level 2 — 224px, HEAVY augmentation (minority class collapse prevention)
+    "level2": {
+        "train": get_heavy_train_transforms(RES_L1_L2),
+        "val":   get_val_transforms(RES_L1_L2),
+    },
+    # Level 3 Macular — 384px, standard (large enough dataset)
+    "level3_macular": {
+        "train": get_train_transforms(RES_L3),
+        "val":   get_val_transforms(RES_L3),
+    },
+    # Level 3 Diabetic — 384px, standard (DME=11,495 samples)
+    "level3_diabetic": {
+        "train": get_train_transforms(RES_L3),
+        "val":   get_val_transforms(RES_L3),
+    },
+    # Level 3 Vascular — 384px, HEAVY (MH=102, RVO=101, RAO=22)
+    "level3_vascular": {
+        "train": get_heavy_train_transforms(RES_L3),
+        "val":   get_val_transforms(RES_L3),
+    },
+    # Level 3 Fluid — 384px, HEAVY (CSR=102 only)
+    "level3_fluid": {
+        "train": get_heavy_train_transforms(RES_L3),
+        "val":   get_val_transforms(RES_L3),
+    },
+    # Level 3 Structural — 384px, HEAVY (ERM=155, VID=76)
+    "level3_structural": {
+        "train": get_heavy_train_transforms(RES_L3),
+        "val":   get_val_transforms(RES_L3),
+    },
+}
+
+
+def get_transforms(mode: str, split: str = "train") -> transforms.Compose:
+    """
+    Convenience accessor for the transform registry.
+
+    Args:
+        mode:  Dataset mode (e.g., ``'level1'``, ``'level3_vascular'``).
+        split: ``'train'`` or ``'val'``.
+
+    Returns:
+        A ``torchvision.transforms.Compose`` instance.
+
+    Raises:
+        ValueError: If mode or split is invalid.
+    """
+    if mode not in TRANSFORM_REGISTRY:
+        raise ValueError(
+            f"Unknown mode: '{mode}'. "
+            f"Choose from: {sorted(TRANSFORM_REGISTRY.keys())}"
+        )
+    if split not in ("train", "val"):
+        raise ValueError(f"Unknown split: '{split}'. Use 'train' or 'val'.")
+    return TRANSFORM_REGISTRY[mode][split]

models/level1_gatekeeper.py

@@ -59,3 +59,16 @@ class GatekeeperModel(nn.Module):
         x = self.avgpool(x)
         x = self.classifier(x)
         return x
+
+def build_gatekeeper(
+    num_classes: int = 2,
+    dropout_rate: float = 0.3,
+    pretrained: bool = True,
+    freeze_backbone: bool = True,
+) -> GatekeeperModel:
+    return GatekeeperModel(
+        num_classes=num_classes,
+        dropout_rate=dropout_rate,
+        pretrained=pretrained,
+        freeze_backbone=freeze_backbone,
+    )

models/level3_specialist.py

@@ -0,0 +1,236 @@
+"""
+models/level3_specialist.py
+
+Level 3 Specialist Models — EfficientNet-B0 fine-grained classifiers.
+
+Input:  384×384 RGB tensors (2.9× more pixels than L1/L2 for structural detail)
+Backbone: EfficientNet-B0 (smaller/faster than B2 — appropriate for specialists
+          that operate on a much smaller subset of the total dataset)
+
+Specialist Instances:
+  ┌───────────────┬──────────────────────────────────────────────┬────────┐
+  │ Key           │ Classes                                      │ Total  │
+  ├───────────────┼──────────────────────────────────────────────┼────────┤
+  │ Macular       │ CNV (Wet AMD) / DRUSEN (Dry AMD) / Generic_AMD│ 47,107│
+  │ Diabetic      │ DME / DR                                     │ 11,602 │
+  │ Vascular      │ MH / RVO / RAO                               │    225 │
+  │ Fluid         │ CSR (single-class anomaly)                   │    102 │
+  │ Structural    │ ERM / VID                                    │    231 │
+  └───────────────┴──────────────────────────────────────────────┴────────┘
+
+AMD Mapping (as per architectural directive):
+  - CNV:         Wet AMD / Choroidal Neovascularization (class 0)
+  - DRUSEN:      Dry AMD / Drusen deposits (class 1)
+  - Generic_AMD: Unclassified AMD from OCTID source (class 2)
+  These are STRICTLY SEPARATED — never merged.
+
+Vascular Mapping (as per architectural directive):
+  - L2 routes all Vascular into a single Vascular_Occlusions bucket.
+  - L3_Vascular re-separates them: MH (0) / RVO (1) / RAO (2).
+
+384px Justification:
+  At 224px, the subtle textural difference between drusen deposits (dry AMD)
+  and sub-retinal fluid (wet AMD/CNV) can be lost. 384px preserves the
+  fine-grained structural detail needed for specialist discrimination.
+  Batch size is reduced to 16 to fit within 32GB MPS unified memory.
+"""
+
+import logging
+from typing import Dict, List
+
+import torch
+import torch.nn as nn
+from torchvision import models
+from torchvision.models import EfficientNet_B0_Weights
+
+logger = logging.getLogger(__name__)
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Registry of all specialist configurations
+# ──────────────────────────────────────────────────────────────────────────────
+SPECIALIST_CONFIGS: Dict[str, Dict] = {
+    "Macular": {
+        "num_classes":     3,
+        "specialist_name": "L3_Macular",
+        "description":     "CNV (Wet AMD) vs DRUSEN (Dry AMD) vs Generic_AMD — strictly separated",
+        "classes":         {0: "CNV", 1: "DRUSEN", 2: "Generic_AMD"},
+    },
+    "Diabetic": {
+        "num_classes":     2,
+        "specialist_name": "L3_Diabetic",
+        "description":     "Diabetic Macular Edema (DME) vs Diabetic Retinopathy (DR)",
+        "classes":         {0: "DME", 1: "DR"},
+    },
+    "Vascular": {
+        "num_classes":     3,
+        "specialist_name": "L3_Vascular",
+        "description":     "Macular Hole (MH) vs RVO vs RAO (re-separated from L2 aggregate)",
+        "classes":         {0: "MH", 1: "RVO", 2: "RAO"},
+    },
+    "Fluid": {
+        "num_classes":     1,
+        "specialist_name": "L3_Fluid",
+        "description":     "Central Serous Retinopathy — single-class anomaly detection",
+        "classes":         {0: "CSR"},
+    },
+    "Structural": {
+        "num_classes":     2,
+        "specialist_name": "L3_Structural",
+        "description":     "Epiretinal Membrane (ERM) vs Vitreomacular Interface Disease (VID)",
+        "classes":         {0: "ERM", 1: "VID"},
+    },
+}
+
+
+class SpecialistModel(nn.Module):
+    """
+    EfficientNet-B0 fine-grained classifier for Level 3 specialist tasks.
+
+    Operates on 384×384 input for maximum structural resolution.
+
+    Architecture:
+        EfficientNet-B0 features (1280-d after avgpool)
+        → Dropout(dropout_rate)
+        → Linear(1280, 512) + SiLU (Swish — native EfficientNet activation)
+        → Dropout(dropout_rate / 2)
+        → Linear(512, num_classes)
+
+    Args:
+        num_classes:     Number of fine-grained classes for this specialist.
+        specialist_name: Human-readable name for logging.
+        dropout_rate:    Dropout probability in classifier head.
+        pretrained:      Load IMAGENET1K_V1 weights if True.
+        freeze_backbone: Start with backbone frozen.
+    """
+
+    def __init__(
+        self,
+        num_classes: int,
+        specialist_name: str = "Specialist",
+        dropout_rate: float = 0.4,
+        pretrained: bool = True,
+        freeze_backbone: bool = True,
+    ) -> None:
+        super().__init__()
+        self.specialist_name = specialist_name
+
+        weights = EfficientNet_B0_Weights.IMAGENET1K_V1 if pretrained else None
+        backbone = models.efficientnet_b0(weights=weights)
+
+        self.features = backbone.features    # MBConv blocks
+        self.avgpool  = backbone.avgpool     # AdaptiveAvgPool2d(1, 1)
+
+        in_features = backbone.classifier[-1].in_features  # 1280 for B0
+
+        # Richer head than the router — specialists need more discriminative power
+        # for subtle inter-class differences (e.g., CNV vs DRUSEN fluid patterns)
+        self.classifier = nn.Sequential(
+            nn.Dropout(p=dropout_rate),
+            nn.Linear(in_features, 512),
+            nn.SiLU(inplace=True),    # Swish — consistent with EfficientNet internals
+            nn.Dropout(p=dropout_rate / 2),
+            nn.Linear(512, num_classes),
+        )
+
+        if freeze_backbone:
+            self.freeze_backbone()
+
+        logger.info(
+            "%s ready | backbone=EfficientNet-B0 | in_features=%d | "
+            "num_classes=%d | input=384×384 | frozen=%s",
+            specialist_name, in_features, num_classes, freeze_backbone,
+        )
+
+    # ──────────────────────────────────────────────────────────────────────────
+    # Freeze / Unfreeze API
+    # ──────────────────────────────────────────────────────────────────────────
+
+    def freeze_backbone(self) -> None:
+        for param in self.features.parameters():
+            param.requires_grad = False
+
+    def unfreeze_backbone(self) -> None:
+        for param in self.features.parameters():
+            param.requires_grad = True
+        logger.info("%s: backbone UNFROZEN.", self.specialist_name)
+
+    def get_param_groups(
+        self,
+        backbone_lr: float = 5e-5,
+        head_lr: float = 5e-4,
+    ) -> List[Dict]:
+        """Differential LR groups for Phase 2 fine-tuning."""
+        return [
+            {"params": self.features.parameters(),   "lr": backbone_lr},
+            {"params": self.classifier.parameters(), "lr": head_lr},
+        ]
+
+    # ──────────────────────────────────────────────────────────────────────────
+    # Forward
+    # ──────────────────────────────────────────────────────────────────────────
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: Float tensor, shape ``(B, 3, 384, 384)``.
+
+        Returns:
+            Logits tensor, shape ``(B, num_classes)``.
+        """
+        x = self.features(x)       # (B, 1280, H', W')
+        x = self.avgpool(x)        # (B, 1280, 1, 1)
+        x = torch.flatten(x, 1)    # (B, 1280)
+        x = self.classifier(x)     # (B, num_classes)
+        return x
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Factory
+# ──────────────────────────────────────────────────────────────────────────────
+
+def build_specialist(
+    specialist_key: str,
+    dropout_rate: float = 0.4,
+    pretrained: bool = True,
+    freeze_backbone: bool = True,
+) -> SpecialistModel:
+    """
+    Factory function for Level 3 specialist models.
+
+    Args:
+        specialist_key:  One of 'Macular', 'Diabetic', 'Vascular',
+                         'Fluid', 'Structural'.
+        dropout_rate:    Head dropout probability.
+        pretrained:      Use ImageNet pretrained weights.
+        freeze_backbone: Start with frozen backbone (Phase 1 warm-up).
+
+    Returns:
+        Configured :class:`SpecialistModel` instance.
+
+    Raises:
+        ValueError: If specialist_key is not in SPECIALIST_CONFIGS.
+
+    Example::
+
+        model = build_specialist('Macular')
+        # → L3_Macular: 3 classes (CNV / DRUSEN / Generic_AMD), 384×384 input
+    """
+    if specialist_key not in SPECIALIST_CONFIGS:
+        raise ValueError(
+            f"Unknown specialist: '{specialist_key}'. "
+            f"Valid keys: {list(SPECIALIST_CONFIGS.keys())}"
+        )
+
+    cfg = SPECIALIST_CONFIGS[specialist_key]
+    logger.info(
+        "Building specialist [%s]: %s",
+        specialist_key, cfg["description"],
+    )
+
+    return SpecialistModel(
+        num_classes=cfg["num_classes"],
+        specialist_name=cfg["specialist_name"],
+        dropout_rate=dropout_rate,
+        pretrained=pretrained,
+        freeze_backbone=freeze_backbone,
+    )

requirements.txt

@@ -5,3 +5,6 @@ uvicorn[standard]
 python-multipart
 Pillow
 huggingface_hub[cli]
+opencv-python-headless
+matplotlib
+numpy

scripts/inference_pipeline.py

@@ -0,0 +1,258 @@
+"""
+scripts/inference_pipeline.py
+
+End-to-End Inference Pipeline for Hierarchical OCT Classification.
+Connects L1 -> L2 -> L3 into a single callable function.
+Takes a raw OCT scan and returns a final diagnosis with confidence scores.
+"""
+
+import sys
+import os
+import json
+import logging
+from pathlib import Path
+from typing import Dict, Any, Optional
+
+import torch
+import torch.nn.functional as F
+from PIL import Image
+
+# Add parent directory to path so we can import from models, data, and utils
+sys.path.append(str(Path(__file__).resolve().parent.parent))
+
+from models.level1_gatekeeper import build_gatekeeper
+from models.level2_router import build_router
+from models.level3_specialist import build_specialist, SPECIALIST_CONFIGS
+from data.transforms import get_transforms
+from utils.gradcam import GradCAM
+
+logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
+logger = logging.getLogger(__name__)
+
+class OCTInferencePipeline:
+    def __init__(
+        self,
+        l1_ckpt: Optional[str] = None,
+        l2_ckpt: Optional[str] = None,
+        l3_ckpts: Optional[Dict[str, str]] = None,
+        device: str = "auto",
+    ):
+        """
+        Initializes the entire L1 -> L2 -> L3 inference pipeline.
+        
+        Args:
+            l1_ckpt: Path to Level 1 (Gatekeeper) checkpoint.
+            l2_ckpt: Path to Level 2 (Router) checkpoint.
+            l3_ckpts: Dict mapping specialist names ('Macular', etc.) to checkpoint paths.
+            device: 'cuda', 'mps', 'cpu', or 'auto'.
+        """
+        if device == "auto":
+            if torch.backends.mps.is_available():
+                self.device = torch.device("mps")
+            elif torch.cuda.is_available():
+                self.device = torch.device("cuda")
+            else:
+                self.device = torch.device("cpu")
+        else:
+            self.device = torch.device(device)
+            
+        logger.info(f"Initialising OCT Inference Pipeline on device: {self.device}")
+        
+        # 1. Load Transforms (we use 'val' split for deterministic preprocessing)
+        self.transform_l1_l2 = get_transforms("level1", "val")
+        # For L3, we can just use level3_macular's val transform as all L3 val transforms are identical (384px)
+        self.transform_l3 = get_transforms("level3_macular", "val")
+        
+        # 2. Build Models
+        logger.info("Building Level 1 Gatekeeper...")
+        self.l1_model = build_gatekeeper(pretrained=True).to(self.device)
+        self._load_ckpt(self.l1_model, l1_ckpt)
+        self.l1_model.eval()
+        
+        logger.info("Building Level 2 Router...")
+        self.l2_model = build_router(pretrained=True).to(self.device)
+        self._load_ckpt(self.l2_model, l2_ckpt)
+        self.l2_model.eval()
+        
+        self.l3_models = {}
+        l3_ckpts = l3_ckpts or {}
+        for spec_name in SPECIALIST_CONFIGS.keys():
+            logger.info(f"Building Level 3 Specialist: {spec_name}...")
+            model = build_specialist(spec_name, pretrained=True).to(self.device)
+            self._load_ckpt(model, l3_ckpts.get(spec_name))
+            model.eval()
+            self.l3_models[spec_name] = model
+            
+        # 3. Label Mappings
+        self.l1_mapping = {0: "NORMAL", 1: "ABNORMAL"}
+        self.l2_mapping = {
+            0: "Macular",
+            1: "Diabetic",
+            2: "Vascular",
+            3: "Fluid",
+            4: "Structural"
+        }
+
+    def _load_ckpt(self, model: torch.nn.Module, ckpt_path: Optional[str]):
+        """Helper to load state dict if path is provided."""
+        if ckpt_path and os.path.exists(ckpt_path):
+            state = torch.load(ckpt_path, map_location=self.device)
+            if "model_state_dict" in state:
+                model.load_state_dict(state["model_state_dict"])
+            else:
+                model.load_state_dict(state)
+            logger.info(f"  -> Loaded weights from {ckpt_path}")
+        else:
+            logger.warning(f"  -> No checkpoint provided for {model.__class__.__name__}. Using random initialization!")
+
+    def _get_heatmap_base64(self, img_pil, cam_array):
+        """Helper to generate a base64 encoded overlay image."""
+        import base64
+        import io
+        overlay = GradCAM.overlay_cam(img_pil, cam_array, alpha=0.5)
+        buffered = io.BytesIO()
+        overlay.save(buffered, format="JPEG")
+        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+        return f"data:image/jpeg;base64,{img_str}"
+
+    def predict(
+        self, 
+        image_path: str, 
+        gradcam: bool = False, 
+        output_dir: str = "output/explanations"
+    ) -> Dict[str, Any]:
+        """
+        Runs the end-to-end inference pipeline on a single image.
+        """
+        logger.info(f"Processing image: {image_path}")
+        
+        # Load Image
+        try:
+            img = Image.open(image_path).convert("RGB")
+        except Exception as e:
+            return {"error": f"Failed to load image: {e}"}
+            
+        # Prepare Tensors
+        tensor_224 = self.transform_l1_l2(img).unsqueeze(0).to(self.device)
+        tensor_384 = self.transform_l3(img).unsqueeze(0).to(self.device)
+        
+        results = {
+            "Level1": {},
+            "Level2": {},
+            "Level3": {},
+            "Final_Diagnosis": None,
+            "Path": [],
+            "gradcams": {}
+        }
+        
+        if gradcam:
+            # Need gradients enabled for Grad-CAM
+            grad_context = torch.enable_grad()
+            # Also require gradients for input tensors
+            tensor_224.requires_grad = True
+            tensor_384.requires_grad = True
+        else:
+            grad_context = torch.no_grad()
+            
+        with grad_context:
+            # --- LEVEL 1: Gatekeeper ---
+            if gradcam:
+                l1_cam_gen = GradCAM(self.l1_model, self.l1_model.features[-1])
+                
+            logits_l1 = self.l1_model(tensor_224)
+            probs_l1 = F.softmax(logits_l1, dim=1).squeeze(0)
+            pred_l1_idx = torch.argmax(probs_l1).item()
+            pred_l1_label = self.l1_mapping[pred_l1_idx]
+            conf_l1 = probs_l1[pred_l1_idx].item()
+            
+            results["Level1"] = {
+                "prediction": pred_l1_label,
+                "confidence": conf_l1,
+                "probs": {self.l1_mapping[i]: probs_l1[i].item() for i in range(2)}
+            }
+            results["Path"].append(f"L1: {pred_l1_label}")
+            
+            if gradcam:
+                heatmap = l1_cam_gen.generate_cam(tensor_224, pred_l1_idx)
+                results["gradcams"]["L1"] = self._get_heatmap_base64(img, heatmap)
+                # Cleanup gradcam to free hooks
+                l1_cam_gen.target_layer._forward_hooks.clear()
+                l1_cam_gen.target_layer._backward_hooks.clear()
+
+            if pred_l1_label == "NORMAL":
+                results["Final_Diagnosis"] = "NORMAL"
+                logger.info("Pipeline terminated at Level 1 (NORMAL)")
+                # return results # Disabled for Grad-CAM testing
+                
+            # --- LEVEL 2: Disease Router ---
+            if gradcam:
+                l2_cam_gen = GradCAM(self.l2_model, self.l2_model.features[-1])
+                
+            logits_l2 = self.l2_model(tensor_224)
+            probs_l2 = F.softmax(logits_l2, dim=1).squeeze(0)
+            pred_l2_idx = torch.argmax(probs_l2).item()
+            pred_l2_label = self.l2_mapping[pred_l2_idx]
+            conf_l2 = probs_l2[pred_l2_idx].item()
+            
+            results["Level2"] = {
+                "prediction": pred_l2_label,
+                "confidence": conf_l2,
+                "probs": {self.l2_mapping[i]: probs_l2[i].item() for i in range(5)}
+            }
+            results["Path"].append(f"L2: {pred_l2_label}")
+            
+            if gradcam:
+                heatmap = l2_cam_gen.generate_cam(tensor_224, pred_l2_idx)
+                results["gradcams"]["L2"] = self._get_heatmap_base64(img, heatmap)
+                l2_cam_gen.target_layer._forward_hooks.clear()
+                l2_cam_gen.target_layer._backward_hooks.clear()
+            
+            # --- LEVEL 3: Specialist ---
+            specialist_model = self.l3_models[pred_l2_label]
+            spec_config = SPECIALIST_CONFIGS[pred_l2_label]
+            l3_classes_map = spec_config["classes"]
+            
+            if gradcam:
+                l3_cam_gen = GradCAM(specialist_model, specialist_model.features[-1])
+                
+            logits_l3 = specialist_model(tensor_384)
+            probs_l3 = F.softmax(logits_l3, dim=1).squeeze(0)
+            pred_l3_idx = torch.argmax(probs_l3).item()
+            pred_l3_label = l3_classes_map[pred_l3_idx]
+            conf_l3 = probs_l3[pred_l3_idx].item()
+            
+            results["Level3"] = {
+                "specialist_used": spec_config["specialist_name"],
+                "prediction": pred_l3_label,
+                "confidence": conf_l3,
+                "probs": {l3_classes_map[i]: probs_l3[i].item() for i in range(len(l3_classes_map))}
+            }
+            results["Path"].append(f"L3: {pred_l3_label}")
+            results["Final_Diagnosis"] = pred_l3_label
+            
+            if gradcam:
+                heatmap = l3_cam_gen.generate_cam(tensor_384, pred_l3_idx)
+                results["gradcams"]["L3"] = self._get_heatmap_base64(img, heatmap)
+                l3_cam_gen.target_layer._forward_hooks.clear()
+                l3_cam_gen.target_layer._backward_hooks.clear()
+                
+        return results
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser(description="Run OCT Hierarchical Inference")
+    parser.add_argument("--image", type=str, required=True, help="Path to raw OCT image")
+    parser.add_argument("--l1_ckpt", type=str, default=None, help="L1 model checkpoint path")
+    parser.add_argument("--l2_ckpt", type=str, default=None, help="L2 model checkpoint path")
+    parser.add_argument("--gradcam", action="store_true", help="Generate Grad-CAM heatmaps")
+    parser.add_argument("--output-dir", type=str, default="output/explanations", help="Output directory for heatmaps")
+    args = parser.parse_args()
+    
+    pipeline = OCTInferencePipeline(
+        l1_ckpt=args.l1_ckpt,
+        l2_ckpt=args.l2_ckpt,
+    )
+    
+    res = pipeline.predict(args.image, gradcam=args.gradcam, output_dir=args.output_dir)
+    print("\n--- INFERENCE RESULTS ---")
+    print(json.dumps(res, indent=4))

utils/gradcam.py

@@ -0,0 +1,128 @@
+import cv2
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+
+class GradCAM:
+    """
+    Grad-CAM implementation for PyTorch models.
+    Supports visualizing the regions of the input image that are most important
+    for the model's prediction.
+    """
+    def __init__(self, model, target_layer):
+        self.model = model
+        self.target_layer = target_layer
+        self.gradients = None
+        self.activations = None
+        
+        # Register hooks
+        self.target_layer.register_forward_hook(self.save_activation)
+        self.target_layer.register_full_backward_hook(self.save_gradient)
+
+    def save_activation(self, module, input, output):
+        if torch.backends.mps.is_available():
+            torch.mps.synchronize()
+        self.activations = output.clone().detach()
+        import logging
+        logger = logging.getLogger("GradCAM")
+        logger.info(f"[Hook] Captured Activations - Max: {self.activations.max().item():.4f}, Min: {self.activations.min().item():.4f}")
+
+    def save_gradient(self, module, grad_input, grad_output):
+        # grad_output is a tuple; we want the first element
+        self.gradients = grad_output[0]
+
+    def generate_cam(self, input_tensor, target_class=None):
+        """
+        Generates the Class Activation Map (CAM).
+        
+        Args:
+            input_tensor (torch.Tensor): Preprocessed image tensor (1, C, H, W)
+            target_class (int, optional): The class to generate CAM for. 
+                                          If None, uses the class with highest score.
+        
+        Returns:
+            np.ndarray: The normalized CAM (H, W) in range [0, 1].
+        """
+        # Ensure we have gradients enabled for this forward/backward pass
+        self.model.zero_grad()
+        
+        # Forward pass
+        output = self.model(input_tensor)
+        
+        if target_class is None:
+            target_class = output.argmax(dim=1).item()
+            
+        # Extract the score for the target class
+        score = output[0, target_class]
+        
+        # Backward pass
+        score.backward()
+        
+        # Get activations and gradients from the hooks
+        gradients = self.gradients.detach().cpu().numpy()[0]   # (C, H, W)
+        activations = self.activations.detach().cpu().numpy()[0] # (C, H, W)
+        
+        # Compute the channel weights (global average pooling of gradients)
+        weights = np.mean(gradients, axis=(1, 2))  # (C,)
+        
+        import logging
+        logger = logging.getLogger("GradCAM")
+        logger.info(f"Gradients - Max: {np.max(gradients):.4e}, Min: {np.min(gradients):.4e}, Sum: {np.sum(gradients):.4e}")
+        logger.info(f"Activations - Max: {np.max(activations):.4f}, Min: {np.min(activations):.4f}, Sum: {np.sum(activations):.4f}")
+        logger.info(f"Weights - Max: {np.max(weights):.4e}, Min: {np.min(weights):.4e}, Sum: {np.sum(weights):.4e}")
+        
+        # Compute the weighted sum of activations
+        cam = np.zeros(activations.shape[1:], dtype=np.float32)
+        for i, w in enumerate(weights):
+            cam += w * activations[i, :, :]
+            
+        import logging
+        logger = logging.getLogger("GradCAM")
+        logger.info(f"Raw CAM - Max: {np.max(cam):.4f}, Min: {np.min(cam):.4f}")
+            
+        # Apply ReLU to keep only features that have a positive influence on the target class
+        cam = np.maximum(cam, 0)
+        
+        # Normalize the CAM to [0, 1]
+        cam = cv2.resize(cam, (input_tensor.shape[3], input_tensor.shape[2]))
+        
+        cam_max = np.max(cam)
+        cam_min = np.min(cam)
+        logger.info(f"Post-ReLU Resized CAM - Max: {cam_max:.4f}, Min: {cam_min:.4f}")
+        
+        cam = cam - cam_min
+        cam = cam / (cam_max + 1e-8)
+        
+        return cam
+
+    @staticmethod
+    def overlay_cam(img_pil: Image.Image, cam: np.ndarray, alpha: float = 0.5) -> Image.Image:
+        """
+        Overlays the CAM heatmap onto the original image.
+        
+        Args:
+            img_pil (PIL.Image.Image): Original image (RGB).
+            cam (np.ndarray): Normalized CAM (H, W) in [0, 1].
+            alpha (float): Blending factor.
+            
+        Returns:
+            PIL.Image.Image: Superimposed image.
+        """
+        # Ensure cam is same size as image
+        if cam.shape != img_pil.size[::-1]:
+            cam = cv2.resize(cam, img_pil.size)
+            
+        # Convert PIL to cv2 (RGB to BGR for colormap)
+        img_cv2 = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
+        
+        # Convert CAM to 8-bit heatmap
+        heatmap = np.uint8(255 * cam)
+        heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
+        
+        # Superimpose
+        superimposed = cv2.addWeighted(img_cv2, 1 - alpha, heatmap, alpha, 0)
+        
+        # Convert back to PIL
+        superimposed = cv2.cvtColor(superimposed, cv2.COLOR_BGR2RGB)
+        return Image.fromarray(superimposed)

weights/level3_diabetic.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:487a5efd9faebde7db003ef05380471980ac261d60b87246e1366e3c1cef6e15
+size 18955227

weights/level3_fluid.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e5c681c681b8422af3e03378054a3dff0d45ee7be2df82d7e66a9678c7062a33
+size 18953179

weights/level3_macular.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a426c914029800d55145f3c159519c60d0c01355809c48c35469aa9e0967d86b
+size 18957275

weights/level3_structural.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b9831f160fb3e99904d4669fa62b2219f9c601a95855b75cab6c07b9e71fc280
+size 18955227

weights/level3_vascular.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f7af4c43b091c277cab9d55efe9f9991cc7b28d2a741199aba71cd5f03e3bba5
+size 18957275