models/transmil_q2l.py

"""
Hybrid TransMIL + Query2Label Architecture

Combines:
- TransMIL's instance-level feature aggregation (with Nystrom attention)
- Query2Label's learnable label queries with cross-attention decoder
- End-to-end training with ResNet-50 backbone

Key Innovation: Extract sequence features from TransMIL BEFORE CLS aggregation,
allowing Q2L label queries to cross-attend across all ultrasound images per case.
"""

import os
import sys

# Add models directory to path for local imports
_models_dir = os.path.dirname(os.path.abspath(__file__))
if _models_dir not in sys.path:
    sys.path.insert(0, _models_dir)

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import numpy as np
from torch.utils.checkpoint import checkpoint_sequential

# Import TransMIL components (from nystrom-attention package)
from nystrom_attention import NystromAttention

# Import Q2L Transformer components from local transformer.py
try:
    from models.transformer import TransformerDecoder, TransformerDecoderLayer
except ImportError:
    try:
        from transformer import TransformerDecoder, TransformerDecoderLayer
    except ImportError:
        print("Warning: Could not import Q2L Transformer components.")

# ============================================================================
# TransMIL Components (Modified)
# ============================================================================

class TransLayer(nn.Module):
    """Transformer layer with Nystrom attention (from TransMIL)"""

    def __init__(self, norm_layer=nn.LayerNorm, dim=512):
        super().__init__()
        self.norm = norm_layer(dim)
        self.attn = NystromAttention(
            dim=dim,
            dim_head=dim // 8,
            heads=8,
            num_landmarks=dim // 2,
            pinv_iterations=6,
            residual=True,
            dropout=0.1
        )

    def forward(self, x):
        x = x + self.attn(self.norm(x))
        return x


class TransMILFeatureExtractor(nn.Module):
    """
    Modified TransMIL that outputs sequence features instead of aggregated CLS token.

    Based on TransMIL.py but extracts features BEFORE CLS aggregation (line 83 output).
    Uses learned 1D position encoding instead of PPEG for simplicity.

    Args:
        input_dim: Dimension of input features (2048 for ResNet-50)
        hidden_dim: Dimension of hidden features (512 default)
        use_ppeg: Whether to use PPEG (2D positional encoding) or learned 1D encoding
        max_seq_len: Maximum sequence length for position encoding
    """

    def __init__(self, input_dim=2048, hidden_dim=512, use_ppeg=False, max_seq_len=1024):
        super().__init__()

        # Feature projection (TransMIL line 50)
        self.fc1 = nn.Sequential(nn.Linear(input_dim, hidden_dim), nn.ReLU())

        # Learnable CLS token (TransMIL line 51)
        self.cls_token = nn.Parameter(torch.randn(1, 1, hidden_dim))

        # Transformer layers (TransMIL lines 53-54)
        self.layer1 = TransLayer(dim=hidden_dim)
        self.layer2 = TransLayer(dim=hidden_dim)

        # LayerNorm (TransMIL line 55)
        self.norm = nn.LayerNorm(hidden_dim)

        # Position encoding
        self.use_ppeg = use_ppeg
        if not use_ppeg:
            # Learned 1D position encoding (simpler than PPEG)
            self.pos_embedding = nn.Parameter(torch.randn(1, max_seq_len, hidden_dim))
        else:
            # PPEG: Position-aware Patch Embedding Generator (requires 2D reshaping)
            self.pos_layer = PPEG(dim=hidden_dim)

    def forward(self, features, mask=None):
        """
        Args:
            features: [B, N, input_dim] - Instance features (e.g., from ResNet-50)
            mask: [B, N] - Padding mask (True = valid instance, False = padded)

        Returns:
            seq_features: [B, 1+N, hidden_dim] - Sequence features (CLS + instances)
            attn_mask: [B, 1+N] - Attention mask for decoder
        """
        B, N, _ = features.shape

        # Project features (TransMIL line 63)
        h = self.fc1(features)  # [B, N, hidden_dim]

        # Handle PPEG padding if needed
        if self.use_ppeg:
            # Pad to nearest square for PPEG (TransMIL lines 65-69)
            H = h.shape[1]
            _H, _W = int(np.ceil(np.sqrt(H))), int(np.ceil(np.sqrt(H)))
            add_length = _H * _W - H
            if add_length > 0:
                h = torch.cat([h, h[:, :add_length, :]], dim=1)  # [B, N_padded, hidden_dim]

                # Update mask
                if mask is not None:
                    pad_mask = torch.zeros(B, add_length, dtype=torch.bool, device=mask.device)
                    mask = torch.cat([mask, pad_mask], dim=1)

        # Add CLS token (TransMIL lines 72-74)
        cls_tokens = self.cls_token.expand(B, -1, -1)
        h = torch.cat([cls_tokens, h], dim=1)  # [B, 1+N, hidden_dim]

        # Update mask to include CLS (always valid)
        if mask is not None:
            cls_mask = torch.ones(B, 1, dtype=torch.bool, device=mask.device)
            attn_mask = torch.cat([cls_mask, mask], dim=1)  # [B, 1+N]
        else:
            attn_mask = torch.ones(B, h.shape[1], dtype=torch.bool, device=h.device)

        # TransLayer 1 (TransMIL line 77)
        h = self.layer1(h)  # [B, 1+N, hidden_dim]

        # Position encoding
        if self.use_ppeg:
            # PPEG (TransMIL line 80)
            h = self.pos_layer(h, _H, _W)
        else:
            # Learned 1D position encoding
            seq_len = h.shape[1]
            h = h + self.pos_embedding[:, :seq_len, :]

        # TransLayer 2 (TransMIL line 83)
        h = self.layer2(h)  # [B, 1+N, hidden_dim]

        # LayerNorm (TransMIL line 86, but keep full sequence)
        h = self.norm(h)  # [B, 1+N, hidden_dim]

        # CRITICAL: Return full sequence, not just CLS token
        return h, attn_mask


class PPEG(nn.Module):
    """
    Position-aware Patch Embedding Generator (from TransMIL)
    Uses 2D depthwise convolutions to inject spatial positional information.
    """

    def __init__(self, dim=512):
        super().__init__()
        self.proj = nn.Conv2d(dim, dim, 7, 1, 7 // 2, groups=dim)
        self.proj1 = nn.Conv2d(dim, dim, 5, 1, 5 // 2, groups=dim)
        self.proj2 = nn.Conv2d(dim, dim, 3, 1, 3 // 2, groups=dim)

    def forward(self, x, H, W):
        """
        Args:
            x: [B, 1+N, C] - Token sequence (CLS + instances)
            H, W: Grid dimensions (H * W >= N)
        """
        B, _, C = x.shape

        # Separate CLS token and feature tokens
        cls_token, feat_token = x[:, 0], x[:, 1:]

        # Reshape to 2D grid
        cnn_feat = feat_token.transpose(1, 2).view(B, C, H, W)

        # Apply 2D convolutions
        x = self.proj(cnn_feat) + cnn_feat + self.proj1(cnn_feat) + self.proj2(cnn_feat)

        # Flatten back to sequence
        x = x.flatten(2).transpose(1, 2)

        # Concatenate CLS token back
        x = torch.cat((cls_token.unsqueeze(1), x), dim=1)

        return x


# ============================================================================
# Query2Label Components (Adapted for Sequences)
# ============================================================================

class GroupWiseLinear(nn.Module):
    """
    Group-wise linear layer for per-class classification (from Q2L).
    Applies a separate linear transformation for each class.
    """

    def __init__(self, num_class, hidden_dim, bias=True):
        super().__init__()
        self.num_class = num_class
        self.hidden_dim = hidden_dim
        self.bias = bias

        self.W = nn.Parameter(torch.Tensor(1, num_class, hidden_dim))
        if bias:
            self.b = nn.Parameter(torch.Tensor(1, num_class))
        self.reset_parameters()

    def reset_parameters(self):
        import math
        stdv = 1. / math.sqrt(self.W.size(2))
        for i in range(self.num_class):
            self.W[0][i].data.uniform_(-stdv, stdv)
        if self.bias:
            for i in range(self.num_class):
                self.b[0][i].data.uniform_(-stdv, stdv)

    def forward(self, x):
        """
        Args:
            x: [B, num_class, hidden_dim]
        Returns:
            logits: [B, num_class]
        """
        # Element-wise multiplication and sum over hidden_dim
        x = (self.W * x).sum(-1)  # [B, num_class]
        if self.bias:
            x = x + self.b
        return x


class HybridQuery2Label(nn.Module):
    """
    Query2Label decoder adapted for sequence inputs (not spatial features).

    Uses learnable label queries to cross-attend to instance sequence from TransMIL.
    Based on query2label.py but modified to accept [B, 1+N, hidden_dim] sequences
    instead of [B, C, H, W] spatial features.

    Args:
        num_class: Number of label classes
        hidden_dim: Dimension of features (512)
        nheads: Number of attention heads
        num_decoder_layers: Number of transformer decoder layers
        dim_feedforward: Dimension of feedforward network
        dropout: Dropout rate
    """

    def __init__(
        self,
        num_class,
        hidden_dim=512,
        nheads=8,
        num_decoder_layers=2,
        dim_feedforward=2048,
        dropout=0.1,
        normalize_before=False
    ):
        super().__init__()
        self.num_class = num_class
        self.hidden_dim = hidden_dim

        # Label query embeddings (Q2L line 68)
        self.query_embed = nn.Embedding(num_class, hidden_dim)

        # Transformer decoder (Q2L uses transformer.py)
        decoder_layer = TransformerDecoderLayer(
            d_model=hidden_dim,
            nhead=nheads,
            dim_feedforward=dim_feedforward,
            dropout=dropout,
            normalize_before=normalize_before
        )
        decoder_norm = nn.LayerNorm(hidden_dim)
        self.decoder = TransformerDecoder(
            decoder_layer,
            num_decoder_layers,
            decoder_norm,
            return_intermediate=False
        )

        # Group-wise linear classifier (Q2L line 69)
        self.fc = GroupWiseLinear(num_class, hidden_dim, bias=True)

    def forward(self, sequence_features, memory_key_padding_mask=None):
        """
        Args:
            sequence_features: [B, 1+N, hidden_dim] - Sequence from TransMIL
            memory_key_padding_mask: [B, 1+N] - Padding mask (True = ignore, False = valid)
                                                 NOTE: PyTorch convention is inverted!

        Returns:
            logits: [B, num_class] - Multi-label classification logits
        """
        B = sequence_features.shape[0]

        # Transpose for decoder: expects [seq_len, B, hidden_dim]
        memory = sequence_features.permute(1, 0, 2)  # [1+N, B, hidden_dim]

        # Label queries (Q2L line 77)
        query_embed = self.query_embed.weight  # [num_class, hidden_dim]
        query_embed = query_embed.unsqueeze(1).repeat(1, B, 1)  # [num_class, B, hidden_dim]

        # Initialize target (zero tensor)
        tgt = torch.zeros_like(query_embed)  # [num_class, B, hidden_dim]

        # Cross-attention decoder (Q2L line 78)
        # Queries attend to instance sequence via cross-attention
        hs = self.decoder(
            tgt=tgt,
            memory=memory,
            memory_key_padding_mask=memory_key_padding_mask,
            pos=None,  # No positional encoding (already in TransMIL)
            query_pos=query_embed
        )  # [1, num_class, B, hidden_dim] if return_intermediate=False

        # Handle output shape
        if hs.dim() == 4:
            hs = hs[-1]  # Take last layer: [num_class, B, hidden_dim]

        # Transpose to [B, num_class, hidden_dim]
        hs = hs.permute(1, 0, 2)  # [B, num_class, hidden_dim]

        # Group-wise linear classification (Q2L line 79)
        logits = self.fc(hs)  # [B, num_class]

        return logits


# ============================================================================
# ResNet-50 Backbone
# ============================================================================

class ResNet50Backbone(nn.Module):
    """
    ResNet-50 feature extractor with Global Average Pooling.

    Extracts 2048-dimensional features from images for TransMIL input.
    Supports gradient checkpointing for memory efficiency.

    Args:
        pretrained: Use ImageNet pre-trained weights
        use_checkpointing: Enable gradient checkpointing (saves memory)
    """

    def __init__(self, pretrained=True, use_checkpointing=False):
        super().__init__()

        # Load ResNet-50
        resnet = torchvision.models.resnet50(pretrained=pretrained)

        # Remove final FC layer and avgpool
        # Output of layer4: [B, 2048, 7, 7] for 224x224 input
        self.features = nn.Sequential(*list(resnet.children())[:-2])

        # Global Average Pooling
        self.gap = nn.AdaptiveAvgPool2d(1)

        self.use_checkpointing = use_checkpointing

    def forward(self, images):
        """
        Args:
            images: [B*N, 3, H, W] - Batch of images (flattened across cases)

        Returns:
            features: [B*N, 2048] - Instance features
        """
        if self.training and self.use_checkpointing:
            # Gradient checkpointing: segment backbone into chunks
            # Trades compute for memory (recomputes activations during backward)
            x = checkpoint_sequential(self.features, segments=4, input=images)
        else:
            x = self.features(images)  # [B*N, 2048, 7, 7]

        x = self.gap(x)  # [B*N, 2048, 1, 1]
        x = x.flatten(1)  # [B*N, 2048]

        return x


# ============================================================================
# Complete End-to-End Model
# ============================================================================

class TransMIL_Query2Label_E2E(nn.Module):
    """
    Complete end-to-end model: Images → ResNet-50 → TransMIL → Q2L → Logits

    Pipeline:
        1. ResNet-50 extracts features from each ultrasound image
        2. TransMIL aggregates variable-length instance sequences with attention
        3. Query2Label decoder performs multi-label classification via cross-attention

    Args:
        num_class: Number of label classes (default 30)
        hidden_dim: Hidden dimension for TransMIL and Q2L (default 512)
        nheads: Number of attention heads in Q2L decoder
        num_decoder_layers: Number of Q2L decoder layers
        pretrained_resnet: Use ImageNet pre-trained ResNet-50
        use_checkpointing: Enable gradient checkpointing for ResNet-50
        use_ppeg: Use PPEG position encoding (vs learned 1D)
    """

    def __init__(
        self,
        num_class=30,
        hidden_dim=512,
        nheads=8,
        num_decoder_layers=2,
        pretrained_resnet=True,
        use_checkpointing=False,
        use_ppeg=False
    ):
        super().__init__()

        # ResNet-50 backbone
        self.backbone = ResNet50Backbone(
            pretrained=pretrained_resnet,
            use_checkpointing=use_checkpointing
        )

        # TransMIL feature extractor (no PPEG by default, learned 1D position encoding)
        self.feature_extractor = TransMILFeatureExtractor(
            input_dim=2048,
            hidden_dim=hidden_dim,
            use_ppeg=use_ppeg
        )

        # Query2Label decoder
        self.q2l_decoder = HybridQuery2Label(
            num_class=num_class,
            hidden_dim=hidden_dim,
            nheads=nheads,
            num_decoder_layers=num_decoder_layers
        )

    def forward(self, images, num_instances_per_case):
        """
        Args:
            images: [B*N_total, 3, H, W] - All images flattened across batch
            num_instances_per_case: [B] or list - Number of images per case

        Returns:
            logits: [B, num_class] - Multi-label classification logits
        """
        # Convert to tensor if list
        if isinstance(num_instances_per_case, list):
            num_instances_per_case = torch.tensor(num_instances_per_case, device=images.device)

        B = len(num_instances_per_case)

        # Step 1: Extract features from all images
        all_features = self.backbone(images)  # [B*N_total, 2048]

        # Step 2: Reshape to [B, max_N, 2048] with padding
        max_N = int(num_instances_per_case.max().item())
        features_padded = torch.zeros(B, max_N, 2048, device=images.device)
        masks = torch.zeros(B, max_N, dtype=torch.bool, device=images.device)

        idx = 0
        for i, n in enumerate(num_instances_per_case):
            n = int(n.item()) if torch.is_tensor(n) else int(n)
            features_padded[i, :n] = all_features[idx:idx+n]
            masks[i, :n] = True  # True = valid instance
            idx += n

        # Step 3: TransMIL sequence features
        seq_features, attn_mask = self.feature_extractor(features_padded, masks)
        # seq_features: [B, 1+max_N, 512]
        # attn_mask: [B, 1+max_N] where True = valid, False = padded

        # Step 4: Q2L decoder
        # IMPORTANT: PyTorch MultiheadAttention uses inverted mask convention!
        # memory_key_padding_mask: True = ignore, False = attend
        # So we need to invert our mask
        decoder_mask = ~attn_mask  # Invert: True = padded (ignore)
        logits = self.q2l_decoder(seq_features, memory_key_padding_mask=decoder_mask)
        # logits: [B, num_class]

        return logits

    def freeze_backbone(self):
        """Freeze ResNet-50 backbone for training only TransMIL+Q2L"""
        for param in self.backbone.parameters():
            param.requires_grad = False

    def unfreeze_backbone(self):
        """Unfreeze ResNet-50 for end-to-end fine-tuning"""
        for param in self.backbone.parameters():
            param.requires_grad = True


# ============================================================================
# Testing
# ============================================================================

if __name__ == "__main__":
    print("Testing TransMIL_Query2Label_E2E model...")

    # Model config
    num_class = 30
    batch_size = 2
    num_instances = [8, 12]  # Variable N per case
    img_size = 224

    # Create model
    model = TransMIL_Query2Label_E2E(
        num_class=num_class,
        hidden_dim=512,
        nheads=8,
        num_decoder_layers=2,
        pretrained_resnet=False,  # Faster for testing
        use_checkpointing=False,
        use_ppeg=False
    )

    # Create dummy data
    total_images = sum(num_instances)
    images = torch.randn(total_images, 3, img_size, img_size)

    print(f"\nInput shapes:")
    print(f"  Images: {images.shape}")
    print(f"  Num instances per case: {num_instances}")

    # Forward pass
    model.eval()
    with torch.no_grad():
        logits = model(images, num_instances)

    print(f"\nOutput shape:")
    print(f"  Logits: {logits.shape}")
    print(f"  Expected: [{batch_size}, {num_class}]")

    assert logits.shape == (batch_size, num_class), "Output shape mismatch!"
    print("\n✓ Model test passed!")

    # Test individual components
    print("\n" + "="*60)
    print("Testing individual components...")
    print("="*60)

    # Test TransMILFeatureExtractor
    print("\n1. TransMILFeatureExtractor")
    feature_extractor = TransMILFeatureExtractor(input_dim=2048, hidden_dim=512)
    features = torch.randn(2, 10, 2048)
    mask = torch.ones(2, 10, dtype=torch.bool)
    seq_features, attn_mask = feature_extractor(features, mask)
    print(f"   Input: {features.shape}, Output: {seq_features.shape}")
    assert seq_features.shape == (2, 11, 512)  # 1 CLS + 10 instances
    print("   ✓ Passed")

    # Test HybridQuery2Label
    print("\n2. HybridQuery2Label")
    decoder = HybridQuery2Label(num_class=30, hidden_dim=512)
    seq_features = torch.randn(2, 11, 512)
    logits = decoder(seq_features)
    print(f"   Input: {seq_features.shape}, Output: {logits.shape}")
    assert logits.shape == (2, 30)
    print("   ✓ Passed")

    # Test ResNet50Backbone
    print("\n3. ResNet50Backbone")
    backbone = ResNet50Backbone(pretrained=False)
    images = torch.randn(20, 3, 224, 224)
    features = backbone(images)
    print(f"   Input: {images.shape}, Output: {features.shape}")
    assert features.shape == (20, 2048)
    print("   ✓ Passed")

    print("\n" + "="*60)
    print("All tests passed! ✓")
    print("="*60)