model_loader.py

# =============================================================================
# model_loader.py — Load ThermoCAN model for deployment
# =============================================================================
# This file handles loading your trained DenseNet121 model checkpoint.
# It is called once when the app starts, not on every prediction.
# =============================================================================

import os
import torch
import torch.nn as nn
import timm
import pickle
import numpy as np

# ── Paste your model architecture here ───────────────────────────────────────
# (Copied from your thermocan.py and cbam.py — needed for loading weights)

class ChannelAttention(nn.Module):
    def __init__(self, in_channels, reduction=16):
        super().__init__()
        mid = max(in_channels // reduction, 1)
        self.mlp = nn.Sequential(
            nn.Linear(in_channels, mid, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(mid, in_channels, bias=False),
        )
    def forward(self, x):
        avg = x.mean(dim=[2, 3])
        mx  = x.amax(dim=[2, 3])
        w   = torch.sigmoid(self.mlp(avg) + self.mlp(mx))
        return x * w.unsqueeze(-1).unsqueeze(-1)


class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super().__init__()
        self.conv = nn.Conv2d(2, 1, kernel_size,
                              padding=kernel_size // 2, bias=False)
    def forward(self, x):
        avg = x.mean(dim=1, keepdim=True)
        mx  = x.amax(dim=1, keepdim=True)
        w   = torch.sigmoid(self.conv(torch.cat([avg, mx], dim=1)))
        return x * w


class CBAM(nn.Module):
    def __init__(self, in_channels, reduction=16):
        super().__init__()
        self.channel = ChannelAttention(in_channels, reduction)
        self.spatial = SpatialAttention()
    def forward(self, x):
        return self.spatial(self.channel(x))


class MetadataMLP(nn.Module):
    def __init__(self, feature_dim, hidden_dims=None, dropout=0.3):
        super().__init__()
        self.feature_dim = feature_dim
        if feature_dim == 0:
            self.net = None
            self.output_dim = 0
            return
        if hidden_dims is None:
            hidden_dims = [64, 32]
        layers = []
        in_dim = feature_dim
        for i, h in enumerate(hidden_dims):
            layers += [nn.Linear(in_dim, h), nn.BatchNorm1d(h), nn.ReLU(inplace=True)]
            if i < len(hidden_dims) - 1:
                layers.append(nn.Dropout(dropout))
            in_dim = h
        self.net = nn.Sequential(*layers)
        self.output_dim = hidden_dims[-1]
    def forward(self, x):
        if self.net is None:
            return torch.zeros(x.shape[0], 0, device=x.device)
        return self.net(x)


class ThermoCAN(nn.Module):
    def __init__(self, backbone_name="densenet121", pretrained=False,
                 num_views=3, cbam_reduction=16, clinical_dim=0,
                 clinical_hidden=None, hidden_dims=None, fusion_dropout=0.5):
        super().__init__()
        self.num_views    = num_views
        self.clinical_dim = clinical_dim
        if clinical_hidden is None: clinical_hidden = [64, 32]
        if hidden_dims     is None: hidden_dims     = [256, 64]

        self.backbone = timm.create_model(
            backbone_name, pretrained=pretrained,
            num_classes=0, global_pool="",
        )
        with torch.no_grad():
            feat = self.backbone(torch.zeros(1, 3, 224, 224))
        backbone_channels = feat.shape[1]

        self.cbam_modules = nn.ModuleList([
            CBAM(backbone_channels, cbam_reduction) for _ in range(num_views)
        ])
        self.gap      = nn.AdaptiveAvgPool2d(1)
        self.meta_mlp = MetadataMLP(clinical_dim, clinical_hidden)
        clinical_out  = self.meta_mlp.output_dim

        fused_dim   = backbone_channels * num_views + clinical_out
        head_layers = []
        in_dim      = fused_dim
        for out_dim in hidden_dims:
            head_layers += [nn.Linear(in_dim, out_dim), nn.BatchNorm1d(out_dim),
                            nn.ReLU(inplace=True), nn.Dropout(fusion_dropout)]
            in_dim = out_dim
        head_layers.append(nn.Linear(in_dim, 1))
        self.classifier = nn.Sequential(*head_layers)

    def forward(self, views, clinical=None):
        vecs = []
        for i, img in enumerate(views):
            f = self.backbone(img)
            f = self.cbam_modules[i](f)
            f = self.gap(f).squeeze(-1).squeeze(-1)
            vecs.append(f)
        fused = torch.cat(vecs, dim=1)
        if clinical is not None and self.clinical_dim > 0:
            fused = torch.cat([fused, self.meta_mlp(clinical)], dim=1)
        return self.classifier(fused)


# ── Model loader function ─────────────────────────────────────────────────────

def load_model(checkpoint_path, config, device):
    """
    Load the trained ThermoCAN model from a checkpoint file.

    Parameters:
        checkpoint_path: path to your best.pth file
        config:          loaded config.yaml dictionary
        device:          torch.device (cpu or cuda)

    Returns:
        model in eval mode, ready for inference
    """
    m         = config["model"]
    d         = config["dataset"]
    use_all   = d.get("use_all_five_views", False)
    primary   = d.get("primary_views", [0, 1, 3])
    num_views = 5 if use_all else len(primary)

    # Load clinical encoder to get feature_dim
    # Ensure checkpoint_path is absolute to find encoder in same dir
    checkpoint_abs = os.path.abspath(checkpoint_path)
    enc_path = os.path.join(
        os.path.dirname(checkpoint_abs), "clinical_encoder.pkl"
    )
    clinical_dim = 0
    encoder      = None
    if os.path.exists(enc_path):
        with open(enc_path, "rb") as f:
            encoder = pickle.load(f)
        clinical_dim = encoder.feature_dim
        print(f"Clinical encoder loaded: {clinical_dim} features")
    else:
        print("No clinical encoder found -- clinical branch disabled")

    # Build model (pretrained=False since we load our own weights)
    model = ThermoCAN(
        backbone_name   = m.get("backbone", "densenet121"),
        pretrained      = False,
        num_views       = num_views,
        cbam_reduction  = m.get("cbam_reduction", 16),
        clinical_dim    = clinical_dim,
        clinical_hidden = m.get("clinical_hidden", [64, 32]),
        hidden_dims     = m.get("hidden_dims", [256, 64]),
        fusion_dropout  = m.get("fusion_dropout", 0.5),
    ).to(device)

    # Load trained weights
    if not os.path.exists(checkpoint_abs):
        raise FileNotFoundError(f"Checkpoint not found at: {checkpoint_abs}")
        
    ckpt = torch.load(checkpoint_abs, map_location=device)
    model.load_state_dict(ckpt["model_state"])
    model.eval()

    trained_epoch = ckpt.get("epoch", "unknown")
    val_metrics   = ckpt.get("metrics", {})
    print(f"Model loaded from epoch {trained_epoch}")
    if val_metrics:
        print(f"   Val AUC: {val_metrics.get('auc', 'N/A')} | "
              f"Val Accuracy: {val_metrics.get('accuracy', 'N/A')}")

    return model, encoder, num_views