PyTorch Implementation: TPMRI Core Components

The implementation below captures the essential mechanisms of TPMRI: temporal modality reweighting with historical signal propagation, cross-modality attention for semantic alignment, and adaptive global fusion for consistency. The code demonstrates how these components integrate into a unified three-stage progressive interaction framework.

# ─────────────────────────────────────────────────────────────────────────────
# TPMRI: Three-stage Progressive Modality Reweighting and Interaction
# Li, Liang, Kong & Zhang · Knowledge-Based Systems 2026
# Core implementation: TMR, CMI, and MAF modules for RGB-T tracking
# ─────────────────────────────────────────────────────────────────────────────

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple, Optional
import math


# ─── Section 1: Temporal Modality Reweighting (TMR) ──────────────────────────

class EMA(nn.Module):
    """
    Efficient Multi-scale Attention module for feature enhancement.
    Used within TMR to emphasize salient regions before response computation.
    """
    
    def __init__(self, channels: int, factor: int = 8):
        super().__init__()
        self.avg_pool_x = nn.AdaptiveAvgPool2d((None, 1))
        self.avg_pool_y = nn.AdaptiveAvgPool2d((1, None))
        
        self.conv1x1 = nn.Conv2d(channels, channels // factor, kernel_size=1, bias=False)
        self.group_norm = nn.GroupNorm(32, channels // factor)
        self.sigmoid = nn.Sigmoid()
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b, c, h, w = x.size()
        
        # X and Y average pooling
        x_avg = self.avg_pool_x(x).expand(-1, -1, h, w)
        y_avg = self.avg_pool_y(x).expand(-1, -1, h, w)
        
        # Convolution path
        conv_out = self.conv1x1(x)
        conv_out = self.sigmoid(conv_out)
        
        # Attention computation
        attn = self.sigmoid(self.group_norm(x_avg + y_avg))
        out = x * attn * conv_out
        
        return self.sigmoid(out)


class TMRModule(nn.Module):
    """
    Temporal Modality Reweighting module.
    
    Dynamically models modality reliability over time by incorporating
    historical modality preference and response discrepancy.
    """

    def __init__(self, channels: int = 768, hidden_dim: int = 64):
        super().__init__()
        self.ema_rgb = EMA(channels)
        self.ema_tir = EMA(channels)
        
        # MLP for modality preference estimation: R^4 -> R^1
        self.mlp = nn.Sequential(
            nn.Linear(4, hidden_dim),
            nn.ReLU(inplace=True),
            nn.Linear(hidden_dim, 1),
            nn.Sigmoid()
        )
        
        # Initialize historical states
        self.alpha_prev = None  # Previous modality preference
        self.delta_r_prev = None  # Previous response discrepancy

    def compute_response(self, feat: torch.Tensor) -> torch.Tensor:
        """Compute modality response strength via global average pooling."""
        return feat.mean(dim=[1, 2, 3], keepdim=True)  # [B, 1, 1, 1]

    def forward(
        self,
        F_rgb: torch.Tensor,  # [B, C, H, W]
        F_tir: torch.Tensor,  # [B, C, H, W]
        reset_history: bool = False
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Temporal modality reweighting with historical signal propagation.
        
        Args:
            F_rgb: RGB features at current timestep
            F_tir: Thermal features at current timestep
            reset_history: Whether to reset temporal state (new sequence)
            
        Returns:
            F_rgb_weighted: Weighted RGB features
            F_tir_weighted: Weighted thermal features
            alpha: Current modality preference [0, 1]
        """
        b = F_rgb.size(0)
        
        # Reset history for new sequences
        if reset_history or self.alpha_prev is None:
            self.alpha_prev = torch.ones(b, 1, device=F_rgb.device) * 0.5
            self.delta_r_prev = torch.zeros(b, 1, device=F_rgb.device)
        
        # Step 1: Feature enhancement via EMA
        F_rgb_enh = self.ema_rgb(F_rgb)
        F_tir_enh = self.ema_tir(F_tir)
        
        # Step 2: Compute current response strengths
        r_rgb = self.compute_response(F_rgb_enh).squeeze()
        r_tir = self.compute_response(F_tir_enh).squeeze()
        
        # Step 3: Compute response discrepancy
        delta_r = r_rgb - r_tir
        
        # Step 4: Construct temporal context vector z^t
        # z^t = [r_rgb^t, r_tir^t, delta_r^{t-1}, alpha^{t-1}]
        z = torch.stack([
            r_rgb,
            r_tir,
            self.delta_r_prev.squeeze(),
            self.alpha_prev.squeeze()
        ], dim=1)  # [B, 4]
        
        # Step 5: Estimate modality preference via MLP
        alpha = self.mlp(z)  # [B, 1], in (0, 1)
        
        # Step 6: Apply modality reweighting
        # F_hat_rgb = alpha * F_rgb, F_hat_tir = (1-alpha) * F_tir
        alpha_expanded = alpha.view(b, 1, 1, 1)
        F_rgb_weighted = alpha_expanded * F_rgb
        F_tir_weighted = (1 - alpha_expanded) * F_tir
        
        # Step 7: Update historical states for next timestep
        self.alpha_prev = alpha.detach()
        self.delta_r_prev = delta_r.detach().unsqueeze(1)
        
        return F_rgb_weighted, F_tir_weighted, alpha


# ─── Section 2: Cross-Modality Interaction (CMI) ─────────────────────────────

class CrossModalityAttention(nn.Module):
    """
    Cross-modality attention mechanism for CMI module.
    RGB features as Query, TIR features as Key and Value.
    """
    
    def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = self.head_dim ** -0.5
        
        # Separate projections for asymmetric interaction
        self.q_proj = nn.Conv2d(dim, dim // 2, kernel_size=1, bias=qkv_bias)  # RGB -> Q
        self.k_proj = nn.Conv2d(dim, dim // 2, kernel_size=1, bias=qkv_bias)  # TIR -> K
        self.v_proj = nn.Conv2d(dim, dim, kernel_size=1, bias=qkv_bias)      # TIR -> V
        
        self.proj = nn.Conv2d(dim, dim, kernel_size=1)
        
    def forward(self, rgb_feat: torch.Tensor, tir_feat: torch.Tensor) -> torch.Tensor:
        """
        Cross-modality attention: RGB queries attend to TIR keys/values.
        
        Args:
            rgb_feat: RGB features [B, C, H, W]
            tir_feat: Thermal features [B, C, H, W]
            
        Returns:
            enhanced_rgb: RGB features enhanced with TIR information
        """
        b, c, h, w = rgb_feat.shape
        
        # Generate Q from RGB, K and V from TIR
        Q = self.q_proj(rgb_feat)  # [B, C/2, H, W]
        K = self.k_proj(tir_feat)  # [B, C/2, H, W]
        V = self.v_proj(tir_feat)  # [B, C, H, W]
        
        # Reshape to 2D format: [B, H*W, C]
        Q = Q.flatten(2).transpose(1, 2)  # [B, HW, C/2]
        K = K.flatten(2).transpose(1, 2)  # [B, HW, C/2]
        V = V.flatten(2).transpose(1, 2)  # [B, HW, C]
        
        # Compute attention: A = Softmax(Q'K')
        attn = torch.bmm(Q, K.transpose(1, 2)) * self.scale  # [B, HW, HW]
        attn = F.softmax(attn, dim=-1)
        
        # Apply attention to values
        enhanced = torch.bmm(attn, V)  # [B, HW, C]
        
        # Reshape back to spatial format
        enhanced = enhanced.transpose(1, 2).view(b, c, h, w)
        enhanced = self.proj(enhanced)
        
        # Residual connection
        output = enhanced + rgb_feat
        
        return output


class CMIModule(nn.Module):
    """
    Cross-Modality Interaction module.
    
    Establishes fine-grained semantic alignment between RGB and TIR
    through asymmetric cross-attention (TIR guides RGB).
    """
    
    def __init__(self, dim: int = 768):
        super().__init__()
        self.cross_attn = CrossModalityAttention(dim)
        
    def forward(
        self,
        F_rgb: torch.Tensor,  # [B, C, H, W] from TMR
        F_tir: torch.Tensor   # [B, C, H, W] from TMR
    ) -> torch.Tensor:
        """
        Perform cross-modality interaction.
        
        Returns:
            F_cmi: Enhanced RGB features with TIR guidance [B, C, H, W]
        """
        F_cmi = self.cross_attn(F_rgb, F_tir)
        return F_cmi


# ─── Section 3: Modality-Aware Fusion (MAF) ──────────────────────────────────

class MAFModule(nn.Module):
    """
    Modality-Aware Fusion module.
    
    Performs adaptive global aggregation to smooth temporal fluctuations
    and prevent bias accumulation.
    """
    
    def __init__(self, dim: int = 768, reduction: int = 4):
        super().__init__()
        self.gap = nn.AdaptiveAvgPool2d(1)
        
        # Two-layer perceptron for fusion weight generation
        self.fc = nn.Sequential(
            nn.Linear(dim * 2, dim // reduction),
            nn.ReLU(inplace=True),
            nn.Linear(dim // reduction, 2),  # Output 2 weights
            nn.Softmax(dim=1)  # Ensure alpha + beta = 1
        )
        
    def forward(
        self,
        F_rgb: torch.Tensor,   # [B, C, H, W] from TMR
        F_cmi: torch.Tensor    # [B, C, H, W] from CMI
    ) -> torch.Tensor:
        """
        Adaptive global fusion of TMR output and CMI-enhanced features.
        
        Args:
            F_rgb: TMR-weighted RGB features
            F_cmi: CMI-enhanced features (RGB with TIR guidance)
            
        Returns:
            F_fused: Globally fused features [B, C, H, W]
        """
        # Step 1: Global information aggregation
        f_rgb = self.gap(F_rgb).flatten(1)  # [B, C]
        f_cmi = self.gap(F_cmi).flatten(1)  # [B, C]
        
        # Step 2: Joint representation
        f_cat = torch.cat([f_rgb, f_cmi], dim=1)  # [B, 2C]
        
        # Step 3: Adaptive weight generation
        weights = self.fc(f_cat)  # [B, 2], weights sum to 1
        alpha = weights[:, 0].view(-1, 1, 1, 1)  # Weight for F_rgb
        beta = weights[:, 1].view(-1, 1, 1,1)         # Step 4: Weighted fusion
        F_fused = alpha * F_rgb + beta * F_cmi
        
        return F_fused


# ─── Section 4: Three-stage Progressive Interaction (TPI) Module ─────────────

class TPIModule(nn.Module):
    """
    Three-stage Progressive Interaction (TPI) module.
    
    Unifies TMR, CMI, and MAF into a single progressive fusion pipeline
    that evolves from temporal modeling to semantic interaction to
    adaptive fusion decision-making.
    """
    
    def __init__(self, dim: int = 768):
        super().__init__()
        self.tmr = TMRModule(dim)
        self.cmi = CMIModule(dim)
        self.maf = MAFModule(dim)
        
    def forward(
        self,
        F_rgb: torch.Tensor,
        F_tir: torch.Tensor,
        reset_history: bool = False
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Three-stage progressive fusion.
        
        Stage 1 (TMR): Temporal modality reweighting - perceives reliability
        Stage 2 (CMI): Cross-modality interaction - establishes semantic alignment
        Stage 3 (MAF): Modality-aware fusion - performs adaptive global aggregation
        
        Returns:
            F_fused: Final fused features
            alpha: Modality preference from TMR (for monitoring)
        """
        # Stage 1: TMR - Temporal Modality Reweighting
        F_rgb_tmr, F_tir_tmr, alpha = self.tmr(F_rgb, F_tir, reset_history)
        
        # Stage 2: CMI - Cross-Modality Interaction (TIR guides RGB)
        F_cmi = self.cmi(F_rgb_tmr, F_tir_tmr)
        
        # Stage 3: MAF - Modality-Aware Fusion (global aggregation)
        F_fused = self.maf(F_rgb_tmr, F_cmi)
        
        return F_fused, alpha


# ─── Section 5: Complete TPMRI Tracker Integration ───────────────────────────

class TPMRI(nn.Module):
    """
    Complete TPMRI tracker with Transformer backbone and TPI modules.
    
    Integrates TPI modules after each Transformer block for progressive
    multi-modal fusion throughout the feature hierarchy.
    """
    
    def __init__(
        self,
        backbone_dim: int = 768,
        num_tpi_layers: int = 12,
        pretrained_backbone: Optional[str] = None
    ):
        super().__init__()
        
        # Transformer backbone (e.g., ViT-Base)
        self.backbone = self._build_backbone(backbone_dim, pretrained_backbone)
        
        # TPI modules integrated after each Transformer block
        self.tpi_modules = nn.ModuleList([
            TPIModule(backbone_dim) for _ in range(num_tpi_layers)
        ])
        
        # Prediction head (center-based from OSTrack)
        self.head = self._build_head(backbone_dim)
        
    def _build_backbone(self, dim: int, pretrained: Optional[str]) -> nn.Module:
        """Build Transformer backbone for feature extraction."""
        # Implementation would load ViT-Base or similar
        # Returns backbone that processes both RGB and TIR patches
        pass
    
    def _build_head(self, dim: int) -> nn.Module:
        """
        Build center-based prediction head.
        Three branches: classification, center offset, and size regression.
        """
        return nn.ModuleDict({
            'cls': nn.Sequential(
                nn.Conv2d(dim, dim, 3, padding=1),
                nn.BatchNorm2d(dim),
                nn.ReLU(inplace=True),
                nn.Conv2d(dim, 1, 1)  # Classification score
            ),
            'offset': nn.Sequential(
                nn.Conv2d(dim, dim, 3, padding=1),
                nn.BatchNorm2d(dim),
                nn.ReLU(inplace=True),
                nn.Conv2d(dim, 2, 1)  # Center offset (x, y)
            ),
            'size': nn.Sequential(
                nn.Conv2d(dim, dim, 3, padding=1),
                nn.BatchNorm2d(dim),
                nn.ReLU(inplace=True),
                nn.Conv2d(dim, 2, 1)  # Size (w, h)
            )
        })
    
    def forward(
        self,
        rgb_template: torch.Tensor,
        rgb_search: torch.Tensor,
        tir_template: torch.Tensor,
        tir_search: torch.Tensor,
        reset_history: bool = False
    ) -> Dict[str, torch.Tensor]:
        """
        Forward pass of TPMRI tracker.
        
        Args:
            rgb_template: Template RGB patches [B, 3, H, W]
            rgb_search: Search RGB patches [B, 3, H, W]
            tir_template: Template thermal patches [B, 3, H, W]
            tir_search: Search thermal patches [B, 3, H, W]
            reset_history: Reset TMR temporal state (new sequence)
            
        Returns:
            outputs: Dict with 'cls' (classification), 'offset', 'size'
        """
        # Extract features from both modalities via backbone
        # Returns list of features at each layer for both template and search
        rgb_features, tir_features = self.backbone(
            rgb_template, rgb_search, tir_template, tir_search
        )
        
        # Progressive fusion through TPI modules at each layer
        fused_features = []
        for i, (rgb_feat, tir_feat, tpi) in enumerate(
            zip(rgb_features, tir_features, self.tpi_modules)
        ):
            # Apply TPI: TMR -> CMI -> MAF
            fused, alpha = tpi(rgb_feat, tir_feat, reset_history=(reset_history and i == 0))
            fused_features.append(fused)
        
        # Final fused features from last layer
        final_features = fused_features[-1]
        
        # Prediction head
        outputs = {
            'cls': self.head['cls'](final_features),
            'offset': self.head['offset'](final_features),
            'size': self.head['size'](final_features)
        }
        
        return outputs


# ─── Section 6: Loss Functions ───────────────────────────────────────────────

class TPMRILoss(nn.Module):
    """
    Composite loss for TPMRI training.
    
    Combines focal loss for classification with GIoU and L1 losses
    for bounding box regression.
    """
    
    def __init__(self, lambda_giou: float = 2.0, lambda_l1: float = 5.0):
        super().__init__()
        self.lambda_giou = lambda_giou
        self.lambda_l1 = lambda_l1
        self.focal_loss = self._focal_loss
        
    def _focal_loss(self, pred: torch.Tensor, target: torch.Tensor, alpha: float = 0.25, gamma: float = 2.0) -> torch.Tensor:
        """Focal loss for classification."""
        bce = F.binary_cross_entropy_with_logits(pred, target, reduction='none')
        pt = torch.exp(-bce)
        focal = alpha * (1 - pt) ** gamma * bce
        return focal.mean()
    
    def giou_loss(self, pred_boxes: torch.Tensor, target_boxes: torch.Tensor) -> torch.Tensor:
        """Generalized IoU loss for bounding box regression."""
        # Implementation of GIoU calculation
        # ...
        pass
    
    def forward(
        self,
        outputs: Dict[str, torch.Tensor],
        targets: Dict[str, torch.Tensor]
    ) -> torch.Tensor:
        """
        Compute total loss.
        
        L_total = L_cls + lambda_giou * L_GIoU + lambda_l1 * L_L1
        """
        loss_cls = self.focal_loss(outputs['cls'], targets['cls'])
        loss_giou = self.giou_loss(outputs['boxes'], targets['boxes'])
        loss_l1 = F.l1_loss(outputs['boxes'], targets['boxes'])
        
        total_loss = loss_cls + self.lambda_giou * loss_giou + self.lambda_l1 * loss_l1
        
        return total_loss


# ─── Section 7: Usage Example ────────────────────────────────────────────────

if __name__ == "__main__":
    """
    Example: Initialize TPMRI tracker and perform tracking inference.
    """
    
    # Initialize TPMRI with ViT-Base backbone
    model = TPMRI(
        backbone_dim=768,
        num_tpi_layers=12,
        pretrained_backbone='vit_base_patch16_224'
    ).cuda()
    
    # Example input tensors (batch_size=1)
    rgb_template = torch.randn(1, 3, 128, 128).cuda()
    rgb_search = torch.randn(1, 3, 256, 256).cuda()
    tir_template = torch.randn(1, 3, 128, 128).cuda()
    tir_search = torch.randn(1, 3, 256, 256).cuda()
    
    # First frame: reset temporal history
    outputs = model(
        rgb_template, rgb_search,
        tir_template, tir_search,
        reset_history=True
    )
    
    print(f"Classification map shape: {outputs['cls'].shape}")
    print(f"Offset map shape: {outputs['offset'].shape}")
    print(f"Size map shape: {outputs['size'].shape}")
    print("✓ TPMRI forward pass complete — temporal fusion enabled")
    
    # Subsequent frames: temporal state propagates automatically
    # outputs = model(rgb_template, rgb_search, tir_template, tir_search)