""" 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)