ensemble_models.py

"""
TB-Guard-XAI Ensemble Models
Issue #18: Complete type hints
Issue #22: Comprehensive docstrings
"""

from typing import Tuple, Optional
import torch
import torch.nn as nn
import torchxrayvision as xrv
from torchvision import models
import timm


class DenseNetTB(nn.Module):
    """
    DenseNet121 backbone for TB detection
    
    Pretrained on CheXpert, fine-tuned for binary TB classification.
    Processes 224x224 single-channel (grayscale) X-ray images.
    
    Architecture:
        - Input: (B, 1, 224, 224) grayscale X-ray images
        - Feature extraction: DenseNet121 pretrained weights
        - Output head: Linear layer mapping to binary classification
        - Output: (B, 1) logits for sigmoid/BCE loss
    
    Args:
        pretrained (bool): Load CheXpert pretrained weights (default: True)
    
    Methods:
        forward(x): Process batch of images, return logits
    """
    
    def __init__(self, pretrained: bool = True) -> None:
        super().__init__()
        if pretrained:
            self.model = xrv.models.DenseNet(weights="densenet121-res224-all")
            self.model.op_threshs = None
        else:
            self.model = xrv.models.DenseNet(weights=None)

        self.model.classifier = nn.Linear(self.model.classifier.in_features, 1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, 1, 224, 224) grayscale X-ray batch
        
        Returns:
            (B, 1) logits for TB classification
        """
        return self.model(x)


class EfficientNetTB(nn.Module):
    """
    EfficientNet-B3 backbone for TB detection
    
    Lightweight architecture suitable for edge deployment.
    Processes 224x224 single-channel (grayscale) X-ray images.
    
    Architecture:
        - Input: (B, 1, 224, 224) grayscale X-ray images
        - Backbone: EfficientNet-B3 with grayscale input adaptation
        - Output: (B, 1) logits for binary classification
    
    Args:
        pretrained (bool): Load ImageNet pretrained weights (default: True)
    
    Methods:
        forward(x): Process batch of images, return logits
    """
    
    def __init__(self, pretrained: bool = True) -> None:
        super().__init__()
        self.model = timm.create_model(
            'efficientnet_b3',
            pretrained=pretrained,
            num_classes=1,
            in_chans=1  # Grayscale input
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, 1, 224, 224) grayscale X-ray batch
        
        Returns:
            (B, 1) logits for TB classification
        """
        return self.model(x)


class ResNetTB(nn.Module):
    """
    ResNet50 backbone for TB detection
    
    Classic architecture with strong performance on medical imaging.
    Adapts RGB-trained weights to grayscale by averaging channels.
    
    Architecture:
        - Input: (B, 1, 224, 224) grayscale X-ray images
        - Conv1: Modified to accept 1 channel (averaged RGB weights)
        - Backbone: ResNet50 residual blocks
        - Output: (B, 1) logits for binary classification
    
    Args:
        pretrained (bool): Load ImageNet pretrained weights (default: True)
                          Channel weights averaged to fit grayscale input
    
    Methods:
        forward(x): Process batch of images, return logits
    """
    
    def __init__(self, pretrained: bool = True) -> None:
        super().__init__()
        self.model = models.resnet50(pretrained=pretrained)
        
        # Adapt Conv1 to accept grayscale input
        old_conv = self.model.conv1
        self.model.conv1 = nn.Conv2d(
            1, 64, kernel_size=7, stride=2, padding=3, bias=False
        )
        
        # Transfer weights by averaging RGB channels
        if pretrained:
            with torch.no_grad():
                self.model.conv1.weight.data = old_conv.weight.data.mean(dim=1, keepdim=True)
        
        # Replace classification head
        self.model.fc = nn.Linear(self.model.fc.in_features, 1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, 1, 224, 224) grayscale X-ray batch
        
        Returns:
            (B, 1) logits for TB classification
        """
        return self.model(x)


class TBEnsemble(nn.Module):
    """
    Three-model weighted ensemble for TB detection
    
    Combines DenseNet, EfficientNet, and ResNet for robust predictions
    with Monte Carlo Dropout-based uncertainty quantification.
    
    Architecture:
        - Three parallel backbones: DenseNet121, EfficientNet-B3, ResNet50
        - Learnable attention weights (soft gating per sample)
        - MC Dropout: 20 forward passes for Bayesian uncertainty
    
    Methods:
        forward(x): Ensemble prediction (sigmoid output 0-1)
        predict_with_uncertainty(x, n_samples): Prediction + std dev
    
    Example:
        >>> model = TBEnsemble()
        >>> x = torch.randn(4, 1, 224, 224)  # Batch of 4 X-rays
        >>> prob = model(x)  # (4, 1), TB probability 0-1
        >>> mean_prob, std_prob = model.predict_with_uncertainty(x)
        >>> print(f"TB prob: {prob[0].item():.2%} ± {std_prob[0].item():.3f}")
    """
    
    def __init__(self, weights: Optional[list] = None) -> None:
        """
        Args:
            weights: Initial ensemble weights (default: equal [1/3, 1/3, 1/3])
                    Can be learnable parameter via optimizer
        """
        super().__init__()
        self.densenet = DenseNetTB(pretrained=True)
        self.efficientnet = EfficientNetTB(pretrained=True)
        self.resnet = ResNetTB(pretrained=True)

        if weights is None:
            self.weights = nn.Parameter(torch.tensor([1/3, 1/3, 1/3]))
        else:
            self.weights = nn.Parameter(torch.tensor(weights, dtype=torch.float32))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Ensemble prediction with weighted logit averaging
        
        Args:
            x: (B, 1, 224, 224) grayscale X-ray batch
        
        Returns:
            (B, 1) sigmoid probabilities 0-1
        """
        logit_densenet = self.densenet(x)
        logit_efficientnet = self.efficientnet(x)
        logit_resnet = self.resnet(x)

        # Normalize weights to sum to 1
        logit_weights = torch.softmax(self.weights, dim=0)
        
        # Weighted ensemble
        ensemble_logit = (
            logit_weights[0] * logit_densenet +
            logit_weights[1] * logit_efficientnet +
            logit_weights[2] * logit_resnet
        )
        
        return torch.sigmoid(ensemble_logit)

    def predict_with_uncertainty(
        self,
        x: torch.Tensor,
        n_samples: int = 20
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Predict with uncertainty using Monte Carlo Dropout
        
        Uses MC Dropout to approximate Bayesian inference.
        Multiple forward passes with dropout enabled estimate
        predictive variance (model uncertainty).
        
        Args:
            x: (B, 1, 224, 224) grayscale X-ray batch
            n_samples: Number of MC Dropout samples (default: 20)
        
        Returns:
            mean_prob: (B, 1) average probability across samples
            std_prob: (B, 1) standard deviation (uncertainty)
        
        Uncertainty Interpretation:
            - std < 0.12: Low uncertainty (model confident)
            - 0.12 <= std < 0.20: Medium uncertainty
            - std >= 0.20: High uncertainty (requires expert review)
        """
        def _enable_dropout(m: nn.Module) -> None:
            """Recursively enable dropout during inference"""
            if isinstance(m, (nn.Dropout, nn.Dropout2d)):
                m.train()

        self.eval()
        self.apply(_enable_dropout)

        predictions = []
        with torch.no_grad():
            for _ in range(n_samples):
                pred = self.forward(x)
                predictions.append(pred)

        # Stack all predictions
        predictions = torch.stack(predictions)  # (n_samples, B, 1)
        
        # Compute mean and std
        mean_prob = predictions.mean(dim=0)  # (B, 1)
        std_prob = predictions.std(dim=0)   # (B, 1)
        
        return mean_prob, std_prob


def load_ensemble(
    checkpoint_path: Optional[str] = None,
    device: str = 'cuda'
) -> TBEnsemble:
    """
    Load or initialize ensemble model
    
    Args:
        checkpoint_path: Path to saved model weights (optional)
        device: 'cuda' or 'cpu' for inference
    
    Returns:
        TBEnsemble model on specified device in eval mode
    
    Example:
        >>> model = load_ensemble('models/ensemble_best.pth', device='cuda')
        >>> model.eval()
    """
    torch.manual_seed(42)
    if device == 'cuda':
        torch.cuda.manual_seed_all(42)

    model = TBEnsemble()
    
    if checkpoint_path:
        # Issue #2: Use weights_only=True for security
        state = torch.load(checkpoint_path, map_location=device, weights_only=True)
        model.load_state_dict(state)

    model = model.to(device)
    model.eval()
    
    return model


if __name__ == "__main__":
    # Test ensemble
    model = TBEnsemble()
    x = torch.randn(2, 1, 224, 224)
    
    # Standard forward pass
    output = model(x)
    print(f"Output shape: {output.shape}")
    print(f"Output (probabilities): {output}")

    # With uncertainty
    mean, std = model.predict_with_uncertainty(x, n_samples=10)
    print(f"\nMean prediction: {mean}")
    print(f"Std prediction: {std}")
    print("\nEnsemble model test passed")