# -*- coding: utf-8 -*-
from einops import rearrange
from models.encoders.VIT.SAM.image_encoder import ImageEncoderViT
from models.encoders.VIT.vits_histo import VisionTransformer
import torch
import torch.nn as nn
from typing import Callable, Tuple, Type, List
class Conv2DBlock(nn.Module):
"""Conv2DBlock with convolution followed by batch-normalisation, ReLU activation and dropout
Args:
in_channels (int): Number of input channels for convolution
out_channels (int): Number of output channels for convolution
kernel_size (int, optional): Kernel size for convolution. Defaults to 3.
dropout (float, optional): Dropout. Defaults to 0.
"""
def __init__(
self,
in_channels: int,
out_channels: int,
kernel_size: int = 3,
dropout: float = 0,
) -> None:
super().__init__()
self.block = nn.Sequential(
nn.Conv2d(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=1,
padding=((kernel_size - 1) // 2),
),
nn.BatchNorm2d(out_channels),
nn.ReLU(True),
nn.Dropout(dropout),
)
def forward(self, x):
return self.block(x)
class Deconv2DBlock(nn.Module):
"""Deconvolution block with ConvTranspose2d followed by Conv2d, batch-normalisation, ReLU activation and dropout
Args:
in_channels (int): Number of input channels for deconv block
out_channels (int): Number of output channels for deconv and convolution.
kernel_size (int, optional): Kernel size for convolution. Defaults to 3.
dropout (float, optional): Dropout. Defaults to 0.
"""
def __init__(
self,
in_channels: int,
out_channels: int,
kernel_size: int = 3,
dropout: float = 0,
) -> None:
super().__init__()
self.block = nn.Sequential(
nn.ConvTranspose2d(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=2,
stride=2,
padding=0,
output_padding=0,
),
nn.Conv2d(
in_channels=out_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=1,
padding=((kernel_size - 1) // 2),
),
nn.BatchNorm2d(out_channels),
nn.ReLU(True),
nn.Dropout(dropout),
)
def forward(self, x):
return self.block(x)
class ViTCellViT(VisionTransformer):
def __init__(
self,
extract_layers: List[int],
img_size: List[int] = [224],
patch_size: int = 16,
in_chans: int = 3,
num_classes: int = 0,
embed_dim: int = 768,
depth: int = 12,
num_heads: int = 12,
mlp_ratio: float = 4,
qkv_bias: bool = False,
qk_scale: float = None,
drop_rate: float = 0,
attn_drop_rate: float = 0,
drop_path_rate: float = 0,
norm_layer: Callable = nn.LayerNorm,
**kwargs
):
"""Vision Transformer with 1D positional embedding
Args:
extract_layers: (List[int]): List of Transformer Blocks whose outputs should be returned in addition to the tokens. First blocks starts with 1, and maximum is N=depth.
img_size (int, optional): Input image size. Defaults to 224.
patch_size (int, optional): Patch Token size (one dimension only, cause tokens are squared). Defaults to 16.
in_chans (int, optional): Number of input channels. Defaults to 3.
num_classes (int, optional): Number of output classes. if num classes = 0, raw tokens are returned (nn.Identity).
Default to 0.
embed_dim (int, optional): Embedding dimension. Defaults to 768.
depth(int, optional): Number of Transformer Blocks. Defaults to 12.
num_heads (int, optional): Number of attention heads per Transformer Block. Defaults to 12.
mlp_ratio (float, optional): MLP ratio for hidden MLP dimension (Bottleneck = dim*mlp_ratio).
Defaults to 4.0.
qkv_bias (bool, optional): If bias should be used for query (q), key (k), and value (v). Defaults to False.
qk_scale (float, optional): Scaling parameter. Defaults to None.
drop_rate (float, optional): Dropout in MLP. Defaults to 0.0.
attn_drop_rate (float, optional): Dropout for attention layer. Defaults to 0.0.
drop_path_rate (float, optional): Dropout for skip connection. Defaults to 0.0.
norm_layer (Callable, optional): Normalization layer. Defaults to nn.LayerNorm.
"""
super().__init__(
img_size=img_size,
patch_size=patch_size,
in_chans=in_chans,
num_classes=num_classes,
embed_dim=embed_dim,
depth=depth,
num_heads=num_heads,
mlp_ratio=mlp_ratio,
qkv_bias=qkv_bias,
qk_scale=qk_scale,
drop_rate=drop_rate,
attn_drop_rate=attn_drop_rate,
drop_path_rate=drop_path_rate,
norm_layer=norm_layer,
)
self.extract_layers = extract_layers
def forward(
self, x: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Forward pass with returning intermediate outputs for skip connections
Args:
x (torch.Tensor): Input batch
Returns:
Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
torch.Tensor: Output of last layers (all tokens, without classification)
torch.Tensor: Classification output
torch.Tensor: Skip connection outputs from extract_layer selection
"""
extracted_layers = []
x = self.prepare_tokens(x)
for depth, blk in enumerate(self.blocks):
x = blk(x)
if depth + 1 in self.extract_layers:
extracted_layers.append(x)
x = self.norm(x)
output = self.head(x[:, 0])
return output, x[:, 0], extracted_layers
class ViTCellViTDeit(ImageEncoderViT):
def __init__(
self,
extract_layers: List[int],
img_size: int = 1024,
patch_size: int = 16,
in_chans: int = 3,
embed_dim: int = 768,
depth: int = 12,
num_heads: int = 12,
mlp_ratio: float = 4,
out_chans: int = 256,
qkv_bias: bool = True,
norm_layer: Type[nn.Module] = nn.LayerNorm,
act_layer: Type[nn.Module] = nn.GELU,
use_abs_pos: bool = True,
use_rel_pos: bool = False,
rel_pos_zero_init: bool = True,
window_size: int = 0,
global_attn_indexes: Tuple[int, ...] = (),
) -> None:
super().__init__(
img_size,
patch_size,
in_chans,
embed_dim,
depth,
num_heads,
mlp_ratio,
out_chans,
qkv_bias,
norm_layer,
act_layer,
use_abs_pos,
use_rel_pos,
rel_pos_zero_init,
window_size,
global_attn_indexes,
)
self.extract_layers = extract_layers
def forward(self, x: torch.Tensor) -> torch.Tensor:
extracted_layers = []
x = self.patch_embed(x)
if self.pos_embed is not None:
token_size = x.shape[1]
x = x + self.pos_embed[:, :token_size, :token_size, :]
for depth, blk in enumerate(self.blocks):
x = blk(x)
if depth + 1 in self.extract_layers:
extracted_layers.append(x)
output = self.neck(x.permute(0, 3, 1, 2))
_output = rearrange(output, "b c h w -> b c (h w)")
return torch.mean(_output, axis=-1), output, extracted_layers