model/maskformer.py

import random
from typing import Tuple, Union, List

import torch.nn as nn
import torch.nn.functional as F
import torch 
from einops import rearrange, repeat, reduce
from positional_encodings.torch_encodings import PositionalEncoding3D
from dynamic_network_architectures.architectures.unet import PlainConvUNet, ResidualEncoderUNet
from dynamic_network_architectures.building_blocks.simple_conv_blocks import StackedConvBlocks
from dynamic_network_architectures.initialization.weight_init import InitWeights_He

from .transformer_decoder import TransformerDecoder,TransformerDecoderLayer
from .SwinUNETR import SwinUNETR

class Maskformer(nn.Module):
    def __init__(self, vision_backbone='UNET', input_channels=1, image_size=[288, 288, 96], patch_size=[32, 32, 32], deep_supervision=False):
        """
        Args:
            vision_backbone (str, optional): visual backbone. Defaults to UNET.
            image_size (list, optional): image size. Defaults to [288, 288, 96].
            patch_size (list, optional): maxium downsample ratio of the bottleneck feature map. Defaults to [32, 32, 32].
            deep_supervision (bool, optional): seg results from mid layers of decoder. Defaults to False.
        """
        super().__init__()
        image_height, image_width, frames = image_size
        self.hw_patch_size = patch_size[0] 
        self.frame_patch_size = patch_size[-1]
        
        self.deep_supervision = deep_supervision
        
        # backbone can be any multi-scale enc-dec vision backbone
        # the enc outputs multi-scale latent features
        # the dec outputs multi-scale per-pixel features
        self.backbone = {
            'SwinUNETR' : SwinUNETR(
                            img_size=[288, 288, 96],    # 48, 48, 96, 192, 384, 768
                            in_channels=3,
                            feature_size=48,  
                            drop_rate=0.0,
                            attn_drop_rate=0.0,
                            dropout_path_rate=0.0,
                            use_checkpoint=False,
                            ),
            'UNET' : PlainConvUNet(input_channels=input_channels, #3, 
                                   n_stages=6, 
                                   features_per_stage=(64, 64, 128, 256, 512, 768), 
                                   conv_op=nn.Conv3d, 
                                   kernel_sizes=3, 
                                   strides=(1, 2, 2, 2, 2, 2), 
                                   n_conv_per_stage=(2, 2, 2, 2, 2, 2), 
                                   n_conv_per_stage_decoder=(2, 2, 2, 2, 2), 
                                   conv_bias=True, 
                                   norm_op=nn.InstanceNorm3d,
                                   norm_op_kwargs={'eps': 1e-5, 'affine': True}, 
                                   dropout_op=None,
                                   dropout_op_kwargs=None,
                                   nonlin=nn.LeakyReLU, 
                                   nonlin_kwargs=None,
                                   deep_supervision=deep_supervision,
                                   nonlin_first=False
                                   ),
            'UNET-L' : PlainConvUNet(input_channels=3, 
                                   n_stages=6, 
                                   features_per_stage=(128, 128, 256, 512, 1024, 1536), 
                                   conv_op=nn.Conv3d, 
                                   kernel_sizes=3, 
                                   strides=(1, 2, 2, 2, 2, 2), 
                                   n_conv_per_stage=(3, 3, 3, 3, 3, 3), 
                                   n_conv_per_stage_decoder=(3, 3, 3, 3, 3), 
                                   conv_bias=True, 
                                   norm_op=nn.InstanceNorm3d,
                                   norm_op_kwargs={'eps': 1e-5, 'affine': True}, 
                                   dropout_op=None,
                                   dropout_op_kwargs=None,
                                   nonlin=nn.LeakyReLU, 
                                   nonlin_kwargs=None,
                                   deep_supervision=deep_supervision,
                                   nonlin_first=False
                                   ),
            'UNET-H' : PlainConvUNet(input_channels=3, 
                                   n_stages=6, 
                                   features_per_stage=(256, 256, 512, 1024, 1536, 2048), 
                                   conv_op=nn.Conv3d, 
                                   kernel_sizes=3, 
                                   strides=(1, 2, 2, 2, 2, 2), 
                                   n_conv_per_stage=(3, 3, 3, 3, 3, 3), 
                                   n_conv_per_stage_decoder=(3, 3, 3, 3, 3), 
                                   conv_bias=True, 
                                   norm_op=nn.InstanceNorm3d,
                                   norm_op_kwargs={'eps': 1e-5, 'affine': True}, 
                                   dropout_op=None,
                                   dropout_op_kwargs=None,
                                   nonlin=nn.LeakyReLU, 
                                   nonlin_kwargs=None,
                                   deep_supervision=deep_supervision,
                                   nonlin_first=False
                                   ),
            'UNET-Res' : ResidualEncoderUNet(
                            input_channels=input_channels,
                            n_stages=6,
                            features_per_stage=[32, 64, 128, 256, 320, 320],
                            conv_op=nn.Conv3d,
                            kernel_sizes=[[3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3]],
                            strides=[[1, 1, 1], [2, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2]],
                            n_blocks_per_stage=[1, 3, 4, 6, 6, 6],
                            n_conv_per_stage_decoder=[1, 1, 1, 1, 1],
                            conv_bias=True,
                            norm_op=nn.InstanceNorm3d,
                            norm_op_kwargs={"eps": 1e-5, "affine": True},
                            nonlin=nn.LeakyReLU,
                            nonlin_kwargs={"inplace": True},
                            deep_supervision=deep_supervision,
                        )
        }[vision_backbone]
        
        self.backbone.apply(InitWeights_He(1e-2))
        
        # fixed to text encoder out dim
        if vision_backbone == 'UNET-H':
            query_dim = 1536
        elif vision_backbone == 'UNET-Res':
            query_dim = 320
        else:
            query_dim = 768

        # all backbones are 6-depth, thus the first 5 scale latent feature outputs need to be down-sampled
        self.avg_pool_ls = [    
            nn.AvgPool3d(32, 32),
            nn.AvgPool3d(16, 16),
            nn.AvgPool3d(8, 8),
            nn.AvgPool3d(4, 4),
            nn.AvgPool3d(2, 2),
            ]
            
        # multi-scale latent feature are projected to query_dim before query decoder
        self.projection_layer = {
            'SwinUNETR' : nn.Sequential(
                        nn.Linear(1536, 768),
                        nn.GELU(),
                        nn.Linear(768, query_dim),
                        nn.GELU()
                    ),
            'UNET' : nn.Sequential(
                        nn.Linear(1792, 768),
                        nn.GELU(),
                        nn.Linear(768, query_dim),
                        nn.GELU()
                    ),
            'UNET-L' : nn.Sequential(
                        nn.Linear(3584, 1536),  # 128, 128, 256, 512, 1024, 1536 --> 3584 --> 768
                        nn.GELU(),
                        nn.Linear(1536, query_dim),
                        nn.GELU()
                    ),
            'UNET-H' : nn.Sequential(
                        nn.Linear(5632, 3072),
                        nn.GELU(),
                        nn.Linear(3072, query_dim),
                        nn.GELU()
                    ),
            'UNET-Res' : nn.Sequential(
                        nn.Linear(1120, 320),
                        nn.GELU(),
                        nn.Linear(320, query_dim),
                        nn.GELU()
                    )
        }[vision_backbone]
        
        # positional encoding
        pos_embedding = PositionalEncoding3D(query_dim)(torch.zeros(1, (image_height//self.hw_patch_size), (image_width//self.hw_patch_size), (frames//self.frame_patch_size), query_dim)) # b h/p w/p d/p dim
        self.pos_embedding = rearrange(pos_embedding, 'b h w d c -> (h w d) b c')   # n b dim
        
        # (fused latent embeddings + pe) x query prompts
        decoder_layer = TransformerDecoderLayer(d_model=query_dim, nhead=8, normalize_before=True)
        decoder_norm = nn.LayerNorm(query_dim)
        self.transformer_decoder = TransformerDecoder(decoder_layer=decoder_layer, num_layers=6, norm=decoder_norm)
        
        if query_dim != 768:
            self.query_proj = nn.Sequential(
                        nn.Linear(768, query_dim), 
                        nn.GELU(),
                        nn.Linear(query_dim, query_dim),
                        nn.GELU()
                    )
        else:
            self.query_proj = nn.Identity()
        
        # mask embedding are projected to perpixel_dim
        # mid stage output (only consider the last 3 mid layers of decoder, i.e. feature maps with resolution /2 /4 /8)
        if self.deep_supervision:
            feature_per_stage = {
                'SwinUNETR':[48, 96, 192],
                'UNET':[64, 128, 256],
                'UNET-L':[128, 256, 512],
                'UNET-H':[256, 512, 1024],
                'UNET-Res':[64, 128, 256]
                }[vision_backbone]
            mid_dim = {
                'SwinUNETR':[256, 384, 512],
                'UNET':[256, 384, 512],
                'UNET-L':[384, 512, 512],
                'UNET-H':[768, 1024, 1024],
                'UNET-Res':[256, 320, 320]
                }[vision_backbone]
            self.mid_mask_embed_proj = []
            for hidden_dim, per_pixel_dim in zip(mid_dim, feature_per_stage):
                self.mid_mask_embed_proj.append(
                    nn.Sequential(
                        nn.Linear(query_dim, hidden_dim),
                        nn.GELU(),
                        nn.Linear(hidden_dim, per_pixel_dim),
                        nn.GELU(),
                        ),
                    )
                self.mid_mask_embed_proj = nn.ModuleList(self.mid_mask_embed_proj)
                
        # largest output        
        mid_dim, per_pixel_dim = {
            'SwinUNETR' : [256, 48],
            'UNET' : [256, 64],
            'UNET-L' : [384, 128],
            'UNET-H' : [768, 256],
            'UNET-Res' : [128, 32]
        }[vision_backbone]
        self.mask_embed_proj = nn.Sequential(
            nn.Linear(query_dim, mid_dim),
            nn.GELU(),
            nn.Linear(mid_dim, per_pixel_dim),
            nn.GELU(),
            )

        self.fusion_conv = StackedConvBlocks(
                1, nn.Conv3d, 2 * per_pixel_dim, per_pixel_dim,
                3, 1, True, nn.InstanceNorm3d, {'eps': 1e-5, 'affine': True},
                None, None, nn.LeakyReLU, None, False)

    def enhance_with_coarse_pred(self, pixel_emb, mask_emb, coarse_pred):
        """
        Enhance pixel embeddings with coarse prediction information
        
        Args:
            pixel_emb (torch.tensor): B,C,H,W,D per-pixel embeddings
            mask_emb (torch.tensor): B,N,C mask embeddings
            coarse_pred (torch.tensor): B,N,H,W,D coarse prediction probabilities
            
        Returns:
            torch.tensor: enhanced pixel embeddings B,C,H,W,D
        """
        
        # [B,N,C] -> [B,C,N]
        mask_emb_transposed = mask_emb.permute(0, 2, 1)  
        
        enhanced_emb = torch.einsum('bnhwd,bcn->bchwd', coarse_pred, mask_emb_transposed)
        
        combined = torch.cat([pixel_emb, enhanced_emb], dim=1)

        enhanced_pixel_emb = self.fusion_conv(combined)  # B,C,H,W,D
        
        return enhanced_pixel_emb
            
    def vision_backbone_forward(self, image_input):
        """
        Visual backbone forward

        Args:
            image_input (torch.tensor): C,H,W,D (C=1)

        Returns:
            image_embedding (torch.tensor): multiscale image features from encoder layers. N,B,d
            pos (torch.tensor): position encoding. N,B,d
            per_pixel_embedding_ls (List of torch.tensor): perpixel embeddings from decoder layers. B,d,H,W,D
        """

        # Image Encoder and Pixel Decoder
        latent_embedding_ls, per_pixel_embedding_ls = self.backbone(image_input) # B Dim H/P W/P D/P
        
        # avg pooling each multiscale feature to H/P W/P D/P
        image_embedding = []
        for latent_embedding, avg_pool in zip(latent_embedding_ls, self.avg_pool_ls):
            tmp = avg_pool(latent_embedding)
            image_embedding.append(tmp)   # B ? H/P W/P D/P
        image_embedding.append(latent_embedding_ls[-1])

        # aggregate multiscale features into image embedding (and proj to align with query dim)
        image_embedding = torch.cat(image_embedding, dim=1)
        image_embedding = rearrange(image_embedding, 'b d h w depth -> b h w depth d')
        image_embedding = self.projection_layer(image_embedding)   # B H/P W/P D/P Dim
        image_embedding = rearrange(image_embedding, 'b h w d dim -> (h w d) b dim') # (H/P W/P D/P) B Dim
            
        # add pe to image embedding
        pos = self.pos_embedding.to(latent_embedding_ls[-1].device)   # (H/P W/P D/P) B Dim
            
        return image_embedding, pos, per_pixel_embedding_ls 

    def infer_forward(self, q, image_embedding, pos, per_pixel_embedding_ls, simulated_lowres_mc_pred=None):
        """
        infer batches of queries (a list) on a batch of patches
        
        Args:
            q (List of torch.tensor): N,d
            simulated_lowres_mc_pred (torch.tensor, optional): B,N,H,W,D low-res multi-channel prediction

        Returns:
            logits (torch.tensor): concat seg output of all queries. B,N_all,H,W,D
        """
        _, B, _ = image_embedding.shape
        
        # query decoder
        N,_ = q.shape    # N is the num of query
        q = repeat(q, 'n dim -> n b dim', b=B) # N B Dim NOTE:By default, attention in torch is not batch_first
        q = self.query_proj(q)
        mask_embedding,_ = self.transformer_decoder(q, image_embedding, pos = pos) # N B Dim
        mask_embedding = rearrange(mask_embedding, 'n b dim -> (b n) dim') # (B N) Dim
        
        # Dot product
        mask_embedding = self.mask_embed_proj(mask_embedding)   # 768 -> 128/64/48
        mask_embedding = rearrange(mask_embedding, '(b n) dim -> b n dim', b=B, n=N)
        per_pixel_embedding = per_pixel_embedding_ls[0] # decoder最后一层的输出
        
        # Enhance features with low-res multi-channel prediction if available
        if simulated_lowres_mc_pred is not None:
            per_pixel_embedding = self.enhance_with_coarse_pred(
                per_pixel_embedding, 
                mask_embedding,
                simulated_lowres_mc_pred)
        
        logits = torch.einsum('bchwd,bnc->bnhwd', per_pixel_embedding, mask_embedding) # bnhwd
        
        return logits
    
    def train_forward(self, queries, image_embedding, pos, per_pixel_embedding_ls, simulated_lowres_mc_pred=None):
        """
        Args:
            queries (torch.tensor): B,N,d
            simulated_lowres_mc_pred (torch.tensor, optional): B,N,H,W,D low-res multi-channel prediction

        Returns:
            logits (List of torch.tensor): list of seg results. B,N,H,W,D
        """
        _, B, _ = image_embedding.shape
        
        # query decoder
        _, N, _ = queries.shape    # N is the num of query
        queries = rearrange(queries, 'b n dim -> n b dim') # N B Dim NOTE:By default, attention in torch is not batch_first
        queries = self.query_proj(queries)
        mask_embedding,_ = self.transformer_decoder(queries, image_embedding, pos = pos) # N B Dim
        mask_embedding = rearrange(mask_embedding, 'n b dim -> (b n) dim') # (B N) Dim
        
        # Dot product
        last_mask_embedding = self.mask_embed_proj(mask_embedding)   # 768 -> 128/64/48
        last_mask_embedding = rearrange(last_mask_embedding, '(b n) dim -> b n dim', b=B, n=N)
        per_pixel_embedding = per_pixel_embedding_ls[0] # decoder最后一层的输出
        
        # Enhance features with low-res multi-channel prediction if available
        if simulated_lowres_mc_pred is not None:
            per_pixel_embedding = self.enhance_with_coarse_pred(
                per_pixel_embedding,
                last_mask_embedding,
                simulated_lowres_mc_pred)
        
        logits = [torch.einsum('bchwd,bnc->bnhwd', per_pixel_embedding, last_mask_embedding)]
        
        # deep supervision
        if self.deep_supervision:
            for mask_embed_proj, per_pixel_embedding in zip(self.mid_mask_embed_proj, per_pixel_embedding_ls[1:]):  # H/2 --> H/16
                mid_mask_embedding = mask_embed_proj(mask_embedding)
                mid_mask_embedding = rearrange(mid_mask_embedding, '(b n) dim -> b n dim', b=B, n=N)
                
                logits.append(torch.einsum('bchwd,bnc->bnhwd', per_pixel_embedding, mid_mask_embedding))
                
        return logits

    def forward(self, queries, image_input, simulated_lowres_sc_pred=None, simulated_lowres_mc_pred=None, train_mode=True):
        # Handle single-channel low-res prediction if provided
        if simulated_lowres_sc_pred is not None:
            # concatenate image and simulated low-res single channel prediction
            image_input = torch.cat([image_input, simulated_lowres_sc_pred], dim=1)  # b2whd

        # get vision features
        image_embedding, pos, per_pixel_embedding_ls = self.vision_backbone_forward(image_input)
            
        # Train Forward -----------------------------------------------------------------------
        if train_mode:
            logits = self.train_forward(queries, image_embedding, pos, per_pixel_embedding_ls, simulated_lowres_mc_pred)
            
        # Infer / Evaluate Forward ------------------------------------------------------------
        else:
            del image_input
            torch.cuda.empty_cache()
            logits = self.infer_forward(queries, image_embedding, pos, per_pixel_embedding_ls, simulated_lowres_mc_pred)
        
        return logits

if __name__ == '__main__':
    model = Maskformer().cuda()    
    image = torch.rand((1, 3, 288, 288, 96)).cuda()
    query = torch.rand((2, 10, 768)).cuda()
    segmentations = model(query, image)
    print(segmentations.shape)