# ============================================================================= # 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