model.py

# model.py
# Vision Transformer (ViT) implementation for masked distance map modeling

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from einops import rearrange, repeat
from einops.layers.torch import Rearrange

class PatchEmbedding(nn.Module):
    """
    Convert input distance map into patches and then embed them.
    """
    def __init__(self, img_size=64, patch_size=8, in_channels=1, embed_dim=768):
        super().__init__()
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = (img_size // patch_size) ** 2
        
        # Linear projection
        self.proj = nn.Sequential(
            Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1=patch_size, p2=patch_size),
            nn.Linear(patch_size * patch_size * in_channels, embed_dim)
        )

    def forward(self, x):
        return self.proj(x)

class PositionalEncoding(nn.Module):
    """
    Add positional encoding to the token embeddings.
    """
    def __init__(self, embed_dim, max_len=5000):
        super().__init__()
        
        # Use fixed positional embeddings for now
        self.pos_embed = nn.Parameter(torch.zeros(1, max_len, embed_dim))
        
        # Initialize with sinusoidal positional encoding
        pos_embed = self.get_sinusoidal_encoding(max_len, embed_dim)
        self.pos_embed.data.copy_(pos_embed)
        
    def get_sinusoidal_encoding(self, max_len, embed_dim):
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, embed_dim, 2).float() * (-np.log(10000.0) / embed_dim))
        
        pos_embed = torch.zeros(1, max_len, embed_dim)
        pos_embed[0, :, 0::2] = torch.sin(position * div_term)
        pos_embed[0, :, 1::2] = torch.cos(position * div_term)
        
        return pos_embed
        
    def forward(self, token_embedding):
        # Add positional embedding
        return token_embedding + self.pos_embed[:, :token_embedding.size(1), :]

class MultiHeadAttention(nn.Module):
    """
    Multi-head attention mechanism.
    """
    def __init__(self, embed_dim, num_heads, dropout=0.1):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        assert self.head_dim * num_heads == embed_dim, "embed_dim must be divisible by num_heads"
        
        # Query, Key, Value projections
        self.qkv_proj = nn.Linear(embed_dim, 3 * embed_dim)
        self.out_proj = nn.Linear(embed_dim, embed_dim)
        self.dropout = nn.Dropout(dropout)
        
        self.scale = self.head_dim ** -0.5
        
    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.shape
        
        # Project input to query, key, value
        qkv = self.qkv_proj(x)
        qkv = qkv.reshape(batch_size, seq_len, 3, self.num_heads, self.head_dim)
        qkv = qkv.permute(2, 0, 3, 1, 4)  # (3, batch_size, num_heads, seq_len, head_dim)
        
        q, k, v = qkv[0], qkv[1], qkv[2]
        
        # Compute attention scores
        attn_scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale  # (batch_size, num_heads, seq_len, seq_len)
        
        # Apply mask if provided
        if mask is not None:
            attn_scores = attn_scores.masked_fill(mask == 0, -1e9)
            
        # Apply attention weights
        attn_weights = F.softmax(attn_scores, dim=-1)
        attn_weights = self.dropout(attn_weights)
        
        # Compute output
        out = torch.matmul(attn_weights, v)  # (batch_size, num_heads, seq_len, head_dim)
        out = out.permute(0, 2, 1, 3).contiguous()  # (batch_size, seq_len, num_heads, head_dim)
        out = out.view(batch_size, seq_len, -1)  # (batch_size, seq_len, embed_dim)
        
        # Output projection
        out = self.out_proj(out)
        
        return out

class FeedForward(nn.Module):
    """
    MLP block in Transformer.
    """
    def __init__(self, embed_dim, hidden_dim, dropout=0.1):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(embed_dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, embed_dim),
            nn.Dropout(dropout)
        )
    
    def forward(self, x):
        return self.net(x)

class TransformerBlock(nn.Module):
    """
    Transformer encoder block with multi-head attention and feed-forward network.
    """
    def __init__(self, embed_dim, num_heads, mlp_ratio=4.0, dropout=0.1):
        super().__init__()
        self.attn_norm = nn.LayerNorm(embed_dim)
        self.attn = MultiHeadAttention(embed_dim, num_heads, dropout)
        self.ff_norm = nn.LayerNorm(embed_dim)
        self.ff = FeedForward(embed_dim, int(embed_dim * mlp_ratio), dropout)
        
    def forward(self, x, mask=None):
        # Attention block with residual connection
        residual = x
        x = self.attn_norm(x)
        x = self.attn(x, mask)
        x = x + residual
        
        # Feed-forward block with residual connection
        residual = x
        x = self.ff_norm(x)
        x = self.ff(x)
        x = x + residual
        
        return x

class MaskedAutoencoderViT(nn.Module):
    """
    Vision Transformer for masked distance map modeling.
    """
    def __init__(
        self,
        img_size=64,
        patch_size=8,
        in_channels=1,
        embed_dim=768,
        depth=12,
        num_heads=12,
        mlp_ratio=4.0,
        dropout=0.1,
        mask_ratio=0.75,
        decoder_embed_dim=512,
        decoder_depth=8,
        decoder_num_heads=16,
    ):
        super().__init__()
        
        # Parameters
        self.img_size = img_size
        self.patch_size = patch_size
        self.in_channels = in_channels
        self.embed_dim = embed_dim
        self.num_patches = (img_size // patch_size) ** 2
        self.mask_ratio = mask_ratio
        self.decoder_embed_dim = decoder_embed_dim
        
        # Patch embedding
        self.patch_embed = PatchEmbedding(
            img_size=img_size,
            patch_size=patch_size,
            in_channels=in_channels,
            embed_dim=embed_dim
        )
        
        # Additional learnable [CLS] token
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        
        # Positional embedding
        self.pos_encoding = PositionalEncoding(embed_dim, self.num_patches + 1)
        
        # Transformer encoder blocks
        self.blocks = nn.ModuleList([
            TransformerBlock(embed_dim, num_heads, mlp_ratio, dropout)
            for _ in range(depth)
        ])
        
        # Encoder normalization
        self.norm = nn.LayerNorm(embed_dim)
        
        # Decoder
        # Project encoder to decoder dimensions
        self.decoder_embed = nn.Linear(embed_dim, decoder_embed_dim, bias=True)
        
        # Decoder positional embeddings
        self.decoder_pos_encoding = PositionalEncoding(decoder_embed_dim, self.num_patches + 1)
        
        # Mask token (shared across all masked patches)
        self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim))
        
        # Decoder blocks
        self.decoder_blocks = nn.ModuleList([
            TransformerBlock(decoder_embed_dim, decoder_num_heads, mlp_ratio, dropout)
            for _ in range(decoder_depth)
        ])
        
        # Decoder normalization
        self.decoder_norm = nn.LayerNorm(decoder_embed_dim)
        
        # Decoder prediction head
        self.decoder_pred = nn.Linear(decoder_embed_dim, patch_size * patch_size * in_channels, bias=True)
        
        # Initialize weights
        self.initialize_weights()
        
    def initialize_weights(self):
        # Initialize patch embedding
        nn.init.xavier_uniform_(self.patch_embed.proj[1].weight)
        nn.init.zeros_(self.patch_embed.proj[1].bias)
        
        # Initialize cls token
        nn.init.normal_(self.cls_token, std=0.02)
        
        # Initialize mask token
        nn.init.normal_(self.mask_token, std=0.02)
        
        # Initialize decoder prediction
        nn.init.xavier_uniform_(self.decoder_pred.weight)
        nn.init.zeros_(self.decoder_pred.bias)
        
    def random_masking(self, x, mask_ratio):
        """
        Apply random masking to input tokens.
        
        Args:
            x (torch.Tensor): Input tokens, shape (B, N, D)
            mask_ratio (float): Ratio of tokens to mask
            
        Returns:
            x_masked (torch.Tensor): Masked input tokens
            mask (torch.Tensor): Binary mask (1 = keep, 0 = mask)
            ids_restore (torch.Tensor): Indices to restore original tokens
        """
        batch_size, seq_len, dim = x.shape
        
        # Calculate number of tokens to mask
        len_keep = int(seq_len * (1 - mask_ratio))
        
        # Generate random noise for masking
        noise = torch.rand(batch_size, seq_len, device=x.device)
        
        # Ensure CLS token is always visible (never masked)
        noise[:, 0] = -1
        
        # Sort noise to get indices
        ids_shuffle = torch.argsort(noise, dim=1)  # Ascending order
        ids_restore = torch.argsort(ids_shuffle, dim=1)  # Indices to restore
        
        # Keep first len_keep indices (including CLS token)
        ids_keep = ids_shuffle[:, :len_keep]
        
        # Create binary mask (1 = keep, 0 = mask)
        mask = torch.zeros(batch_size, seq_len, device=x.device)
        mask.scatter_(1, ids_keep, 1.0)
        
        # Generate masked tokens using gather
        x_masked = torch.gather(
            x, dim=1, 
            index=ids_keep.unsqueeze(-1).repeat(1, 1, dim)
        )
        
        return x_masked, mask, ids_restore
    
    def forward_encoder(self, x, mask_ratio):
        """
        Forward pass through the encoder.
        
        Args:
            x (torch.Tensor): Input distance maps, shape (B, C, H, W)
            mask_ratio (float): Ratio of patches to mask
            
        Returns:
            x (torch.Tensor): Encoded features of visible patches
            mask (torch.Tensor): Binary mask (1 = keep, 0 = mask)
            ids_restore (torch.Tensor): Indices to restore original order
        """
        # Convert image to patch embeddings
        x = self.patch_embed(x)  # (B, N, D)
        
        # Add cls token
        cls_token = self.cls_token.expand(x.shape[0], -1, -1)  # (B, 1, D)
        x = torch.cat([cls_token, x], dim=1)  # (B, N+1, D)
        
        # Add positional encoding
        x = self.pos_encoding(x)  # (B, N+1, D)
        
        # Apply random masking
        x, mask, ids_restore = self.random_masking(x, mask_ratio)
        
        # Forward through encoder blocks
        for block in self.blocks:
            x = block(x)
        
        # Apply final normalization
        x = self.norm(x)
        
        return x, mask, ids_restore
    
    def forward_decoder(self, x, ids_restore):
        """
        Forward pass through the decoder.
        
        Args:
            x (torch.Tensor): Encoded features of visible patches, shape (B, N_vis, D)
            ids_restore (torch.Tensor): Indices to restore original order
            
        Returns:
            x (torch.Tensor): Reconstructed patch features
        """
        batch_size = x.shape[0]
        
        # Project encoder features to decoder dimensions
        x = self.decoder_embed(x)  # (B, N_vis, D_d)
        
        # Calculate number of tokens in original sequence (before masking)
        # This is needed to handle edge cases with the restore indices
        n_orig = ids_restore.shape[1]
        
        # Prepare mask tokens
        mask_tokens = self.mask_token.repeat(batch_size, n_orig - x.shape[1] + 1, 1)  # +1 for cls token
        
        # First token is CLS token
        cls_token = x[:, :1, :]
        
        # The rest are patch tokens (after masking)
        patch_tokens = x[:, 1:, :]
        
        # Concatenate visible tokens with mask tokens
        x_ = torch.cat([patch_tokens, mask_tokens], dim=1)  # Skip CLS token, (B, N, D_d)
        
        # Make sure x_ has the same size as ids_restore (minus 1 for cls token)
        x_ = x_[:, :(n_orig - 1), :]
        
        # Restore original order using gather - limiting indices to valid range
        ids_to_gather = ids_restore[:, 1:].clamp(0, x_.size(1) - 1)
        
        # Restore original order
        x_ = torch.gather(
            x_, dim=1,
            index=ids_to_gather.unsqueeze(-1).repeat(1, 1, x.shape[2])
        )  # (B, N, D_d)
        
        # Append CLS token
        x = torch.cat([cls_token, x_], dim=1)  # (B, N+1, D_d)
        
        # Add positional encoding
        x = self.decoder_pos_encoding(x)
        
        # Forward through decoder blocks
        for block in self.decoder_blocks:
            x = block(x)
        
        # Apply final normalization
        x = self.decoder_norm(x)
        
        # Predict original patches (skip CLS token)
        x = self.decoder_pred(x[:, 1:, :])  # (B, N, patch_size^2*C)
        
        # Reshape to patch dimensions
        x = x.reshape(batch_size, self.img_size // self.patch_size, self.img_size // self.patch_size, 
                    self.patch_size, self.patch_size, self.in_channels)
        x = x.permute(0, 5, 1, 3, 2, 4)  # (B, C, H/p1, p1, W/p2, p2)
        x = x.reshape(batch_size, self.in_channels, self.img_size, self.img_size)  # (B, C, H, W)
        
        return x
    
    def forward(self, x, mask_ratio=None):
        """
        Forward pass through the full model.
        
        Args:
            x (torch.Tensor): Input distance maps, shape (B, C, H, W)
            mask_ratio (float, optional): Ratio of patches to mask. If None, use self.mask_ratio.
            
        Returns:
            pred (torch.Tensor): Reconstructed distance maps
            mask (torch.Tensor): Binary mask (1 = keep, 0 = mask)
        """
        if mask_ratio is None:
            mask_ratio = self.mask_ratio
            
        # Forward through encoder
        latent, mask, ids_restore = self.forward_encoder(x, mask_ratio)
        
        # Forward through decoder
        pred = self.forward_decoder(latent, ids_restore)
        
        return pred, mask