diff --git a/.gitattributes b/.gitattributes index ba87e40af32e1ed85fb59df2b0eb4265b9192572..b06e34daf4eccb9470b3e3673b8da0f9b69ee7d7 100644 --- a/.gitattributes +++ b/.gitattributes @@ -341,3 +341,4 @@ models/sapiens/depth/sapiens_1b_render_people_epoch_88_torchscript.pt2 filter=lf models/sapiens/normal/sapiens_1b_normal_render_people_epoch_115_torchscript.pt2 filter=lfs diff=lfs merge=lfs -text models/sapiens/pose/sapiens_1b_goliath_best_goliath_AP_639_torchscript.pt2 filter=lfs diff=lfs merge=lfs -text models/sapiens/seg/sapiens_1b_goliath_best_goliath_mIoU_7994_epoch_151_torchscript.pt2 filter=lfs diff=lfs merge=lfs -text +models/RMBG/SAM/Moonlit[[:space:]]Serenade.mp3 filter=lfs diff=lfs merge=lfs -text diff --git a/models/RMBG/BEN/BEN_Base.pth b/models/RMBG/BEN/BEN_Base.pth new file mode 100644 index 0000000000000000000000000000000000000000..f50f1c77d6c450652b3342df8a3613c7b3da7a29 --- /dev/null +++ b/models/RMBG/BEN/BEN_Base.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9f00a5804c96afed6e19be09dbbbe56ccaf82cff3d751f44aadb9626b77facfa +size 1134588350 diff --git a/models/RMBG/BEN/gitattributes b/models/RMBG/BEN/gitattributes new file mode 100644 index 0000000000000000000000000000000000000000..a6344aac8c09253b3b630fb776ae94478aa0275b --- /dev/null +++ b/models/RMBG/BEN/gitattributes @@ -0,0 +1,35 @@ +*.7z filter=lfs diff=lfs merge=lfs -text +*.arrow filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +*.bz2 filter=lfs diff=lfs merge=lfs -text +*.ckpt filter=lfs diff=lfs merge=lfs -text +*.ftz filter=lfs diff=lfs merge=lfs -text +*.gz filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text +*.joblib filter=lfs diff=lfs merge=lfs -text +*.lfs.* filter=lfs diff=lfs merge=lfs -text +*.mlmodel filter=lfs diff=lfs merge=lfs -text +*.model filter=lfs diff=lfs merge=lfs -text +*.msgpack filter=lfs diff=lfs merge=lfs -text +*.npy filter=lfs diff=lfs merge=lfs -text +*.npz filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.ot filter=lfs diff=lfs merge=lfs -text +*.parquet filter=lfs diff=lfs merge=lfs -text +*.pb filter=lfs diff=lfs merge=lfs -text +*.pickle filter=lfs diff=lfs merge=lfs -text +*.pkl filter=lfs diff=lfs merge=lfs -text +*.pt filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.rar filter=lfs diff=lfs merge=lfs -text +*.safetensors filter=lfs diff=lfs merge=lfs -text +saved_model/**/* filter=lfs diff=lfs merge=lfs -text +*.tar.* filter=lfs diff=lfs merge=lfs -text +*.tar filter=lfs diff=lfs merge=lfs -text +*.tflite filter=lfs diff=lfs merge=lfs -text +*.tgz filter=lfs diff=lfs merge=lfs -text +*.wasm filter=lfs diff=lfs merge=lfs -text +*.xz filter=lfs diff=lfs merge=lfs -text +*.zip filter=lfs diff=lfs merge=lfs -text +*.zst filter=lfs diff=lfs merge=lfs -text +*tfevents* filter=lfs diff=lfs merge=lfs -text diff --git a/models/RMBG/BEN/model.py b/models/RMBG/BEN/model.py new file mode 100644 index 0000000000000000000000000000000000000000..502c44898b1d786466c6bc801349cee5c15ee5ae --- /dev/null +++ b/models/RMBG/BEN/model.py @@ -0,0 +1,951 @@ +import math +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as checkpoint +from einops import rearrange +from PIL import Image, ImageFilter, ImageOps +from timm.layers import DropPath, to_2tuple, trunc_normal_ +from torchvision import transforms + +class Mlp(nn.Module): + """ Multilayer perceptron.""" + + def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +def window_partition(x, window_size): + """ + Args: + x: (B, H, W, C) + window_size (int): window size + Returns: + windows: (num_windows*B, window_size, window_size, C) + """ + B, H, W, C = x.shape + x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) + windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) + return windows + + +def window_reverse(windows, window_size, H, W): + """ + Args: + windows: (num_windows*B, window_size, window_size, C) + window_size (int): Window size + H (int): Height of image + W (int): Width of image + Returns: + x: (B, H, W, C) + """ + B = int(windows.shape[0] / (H * W / window_size / window_size)) + x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1) + x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) + return x + + +class WindowAttention(nn.Module): + """ Window based multi-head self attention (W-MSA) module with relative position bias. + It supports both of shifted and non-shifted window. + Args: + dim (int): Number of input channels. + window_size (tuple[int]): The height and width of the window. + num_heads (int): Number of attention heads. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set + attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0 + proj_drop (float, optional): Dropout ratio of output. Default: 0.0 + """ + + def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.): + + super().__init__() + self.dim = dim + self.window_size = window_size # Wh, Ww + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = qk_scale or head_dim ** -0.5 + + # define a parameter table of relative position bias + self.relative_position_bias_table = nn.Parameter( + torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)) # 2*Wh-1 * 2*Ww-1, nH + + # get pair-wise relative position index for each token inside the window + coords_h = torch.arange(self.window_size[0]) + coords_w = torch.arange(self.window_size[1]) + coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww + coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww + relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww + relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2 + relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0 + relative_coords[:, :, 1] += self.window_size[1] - 1 + relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1 + relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww + self.register_buffer("relative_position_index", relative_position_index) + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + trunc_normal_(self.relative_position_bias_table, std=.02) + self.softmax = nn.Softmax(dim=-1) + + def forward(self, x, mask=None): + """ Forward function. + Args: + x: input features with shape of (num_windows*B, N, C) + mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None + """ + B_, N, C = x.shape + qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple) + + q = q * self.scale + attn = (q @ k.transpose(-2, -1)) + + relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view( + self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1) # Wh*Ww,Wh*Ww,nH + relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww + attn = attn + relative_position_bias.unsqueeze(0) + + if mask is not None: + nW = mask.shape[0] + attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0) + attn = attn.view(-1, self.num_heads, N, N) + attn = self.softmax(attn) + else: + attn = self.softmax(attn) + + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B_, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class SwinTransformerBlock(nn.Module): + """ Swin Transformer Block. + Args: + dim (int): Number of input channels. + num_heads (int): Number of attention heads. + window_size (int): Window size. + shift_size (int): Shift size for SW-MSA. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float, optional): Stochastic depth rate. Default: 0.0 + act_layer (nn.Module, optional): Activation layer. Default: nn.GELU + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + """ + + def __init__(self, dim, num_heads, window_size=7, shift_size=0, + mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0., + act_layer=nn.GELU, norm_layer=nn.LayerNorm): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.window_size = window_size + self.shift_size = shift_size + self.mlp_ratio = mlp_ratio + assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size" + + self.norm1 = norm_layer(dim) + self.attn = WindowAttention( + dim, window_size=to_2tuple(self.window_size), num_heads=num_heads, + qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop) + + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + + self.H = None + self.W = None + + def forward(self, x, mask_matrix): + """ Forward function. + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + mask_matrix: Attention mask for cyclic shift. + """ + B, L, C = x.shape + H, W = self.H, self.W + assert L == H * W, "input feature has wrong size" + + shortcut = x + x = self.norm1(x) + x = x.view(B, H, W, C) + + # pad feature maps to multiples of window size + pad_l = pad_t = 0 + pad_r = (self.window_size - W % self.window_size) % self.window_size + pad_b = (self.window_size - H % self.window_size) % self.window_size + x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b)) + _, Hp, Wp, _ = x.shape + + # cyclic shift + if self.shift_size > 0: + shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)) + attn_mask = mask_matrix + else: + shifted_x = x + attn_mask = None + + # partition windows + x_windows = window_partition(shifted_x, self.window_size) # nW*B, window_size, window_size, C + x_windows = x_windows.view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C + + # W-MSA/SW-MSA + attn_windows = self.attn(x_windows, mask=attn_mask) # nW*B, window_size*window_size, C + + # merge windows + attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C) + shifted_x = window_reverse(attn_windows, self.window_size, Hp, Wp) # B H' W' C + + # reverse cyclic shift + if self.shift_size > 0: + x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2)) + else: + x = shifted_x + + if pad_r > 0 or pad_b > 0: + x = x[:, :H, :W, :].contiguous() + + x = x.view(B, H * W, C) + + # FFN + x = shortcut + self.drop_path(x) + x = x + self.drop_path(self.mlp(self.norm2(x))) + + return x + + +class PatchMerging(nn.Module): + """ Patch Merging Layer + Args: + dim (int): Number of input channels. + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + """ + def __init__(self, dim, norm_layer=nn.LayerNorm): + super().__init__() + self.dim = dim + self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False) + self.norm = norm_layer(4 * dim) + + def forward(self, x, H, W): + """ Forward function. + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + """ + B, L, C = x.shape + assert L == H * W, "input feature has wrong size" + + x = x.view(B, H, W, C) + + # padding + pad_input = (H % 2 == 1) or (W % 2 == 1) + if pad_input: + x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2)) + + x0 = x[:, 0::2, 0::2, :] # B H/2 W/2 C + x1 = x[:, 1::2, 0::2, :] # B H/2 W/2 C + x2 = x[:, 0::2, 1::2, :] # B H/2 W/2 C + x3 = x[:, 1::2, 1::2, :] # B H/2 W/2 C + x = torch.cat([x0, x1, x2, x3], -1) # B H/2 W/2 4*C + x = x.view(B, -1, 4 * C) # B H/2*W/2 4*C + + x = self.norm(x) + x = self.reduction(x) + + return x + + +class BasicLayer(nn.Module): + """ A basic Swin Transformer layer for one stage. + Args: + dim (int): Number of feature channels + depth (int): Depths of this stage. + num_heads (int): Number of attention head. + window_size (int): Local window size. Default: 7. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. + """ + + def __init__(self, + dim, + depth, + num_heads, + window_size=7, + mlp_ratio=4., + qkv_bias=True, + qk_scale=None, + drop=0., + attn_drop=0., + drop_path=0., + norm_layer=nn.LayerNorm, + downsample=None, + use_checkpoint=False): + super().__init__() + self.window_size = window_size + self.shift_size = window_size // 2 + self.depth = depth + self.use_checkpoint = use_checkpoint + + # build blocks + self.blocks = nn.ModuleList([ + SwinTransformerBlock( + dim=dim, + num_heads=num_heads, + window_size=window_size, + shift_size=0 if (i % 2 == 0) else window_size // 2, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop, + attn_drop=attn_drop, + drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path, + norm_layer=norm_layer) + for i in range(depth)]) + + # patch merging layer + if downsample is not None: + self.downsample = downsample(dim=dim, norm_layer=norm_layer) + else: + self.downsample = None + + def forward(self, x, H, W): + """ Forward function. + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + """ + + # calculate attention mask for SW-MSA + Hp = int(np.ceil(H / self.window_size)) * self.window_size + Wp = int(np.ceil(W / self.window_size)) * self.window_size + img_mask = torch.zeros((1, Hp, Wp, 1), device=x.device) # 1 Hp Wp 1 + h_slices = (slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None)) + w_slices = (slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None)) + cnt = 0 + for h in h_slices: + for w in w_slices: + img_mask[:, h, w, :] = cnt + cnt += 1 + + mask_windows = window_partition(img_mask, self.window_size) # nW, window_size, window_size, 1 + mask_windows = mask_windows.view(-1, self.window_size * self.window_size) + attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) + attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0)) + + for blk in self.blocks: + blk.H, blk.W = H, W + if self.use_checkpoint: + x = checkpoint.checkpoint(blk, x, attn_mask) + else: + x = blk(x, attn_mask) + if self.downsample is not None: + x_down = self.downsample(x, H, W) + Wh, Ww = (H + 1) // 2, (W + 1) // 2 + return x, H, W, x_down, Wh, Ww + else: + return x, H, W, x, H, W + + +class PatchEmbed(nn.Module): + """ Image to Patch Embedding + Args: + patch_size (int): Patch token size. Default: 4. + in_chans (int): Number of input image channels. Default: 3. + embed_dim (int): Number of linear projection output channels. Default: 96. + norm_layer (nn.Module, optional): Normalization layer. Default: None + """ + + def __init__(self, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None): + super().__init__() + patch_size = to_2tuple(patch_size) + self.patch_size = patch_size + + self.in_chans = in_chans + self.embed_dim = embed_dim + + self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) + if norm_layer is not None: + self.norm = norm_layer(embed_dim) + else: + self.norm = None + + def forward(self, x): + """Forward function.""" + # padding + _, _, H, W = x.size() + if W % self.patch_size[1] != 0: + x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1])) + if H % self.patch_size[0] != 0: + x = F.pad(x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0])) + + x = self.proj(x) # B C Wh Ww + if self.norm is not None: + Wh, Ww = x.size(2), x.size(3) + x = x.flatten(2).transpose(1, 2) + x = self.norm(x) + x = x.transpose(1, 2).view(-1, self.embed_dim, Wh, Ww) + + return x + + +class SwinTransformer(nn.Module): + """ Swin Transformer backbone. + A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows` - + https://arxiv.org/pdf/2103.14030 + Args: + pretrain_img_size (int): Input image size for training the pretrained model, + used in absolute postion embedding. Default 224. + patch_size (int | tuple(int)): Patch size. Default: 4. + in_chans (int): Number of input image channels. Default: 3. + embed_dim (int): Number of linear projection output channels. Default: 96. + depths (tuple[int]): Depths of each Swin Transformer stage. + num_heads (tuple[int]): Number of attention head of each stage. + window_size (int): Window size. Default: 7. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4. + qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. + drop_rate (float): Dropout rate. + attn_drop_rate (float): Attention dropout rate. Default: 0. + drop_path_rate (float): Stochastic depth rate. Default: 0.2. + norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm. + ape (bool): If True, add absolute position embedding to the patch embedding. Default: False. + patch_norm (bool): If True, add normalization after patch embedding. Default: True. + out_indices (Sequence[int]): Output from which stages. + frozen_stages (int): Stages to be frozen (stop grad and set eval mode). + -1 means not freezing any parameters. + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. + """ + + def __init__(self, + pretrain_img_size=224, + patch_size=4, + in_chans=3, + embed_dim=96, + depths=[2, 2, 6, 2], + num_heads=[3, 6, 12, 24], + window_size=7, + mlp_ratio=4., + qkv_bias=True, + qk_scale=None, + drop_rate=0., + attn_drop_rate=0., + drop_path_rate=0.2, + norm_layer=nn.LayerNorm, + ape=False, + patch_norm=True, + out_indices=(0, 1, 2, 3), + frozen_stages=-1, + use_checkpoint=False): + super().__init__() + + self.pretrain_img_size = pretrain_img_size + self.num_layers = len(depths) + self.embed_dim = embed_dim + self.ape = ape + self.patch_norm = patch_norm + self.out_indices = out_indices + self.frozen_stages = frozen_stages + + # split image into non-overlapping patches + self.patch_embed = PatchEmbed( + patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, + norm_layer=norm_layer if self.patch_norm else None) + + # absolute position embedding + if self.ape: + pretrain_img_size = to_2tuple(pretrain_img_size) + patch_size = to_2tuple(patch_size) + patches_resolution = [pretrain_img_size[0] // patch_size[0], pretrain_img_size[1] // patch_size[1]] + + self.absolute_pos_embed = nn.Parameter(torch.zeros(1, embed_dim, patches_resolution[0], patches_resolution[1])) + trunc_normal_(self.absolute_pos_embed, std=.02) + + self.pos_drop = nn.Dropout(p=drop_rate) + + # stochastic depth + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule + + # build layers + self.layers = nn.ModuleList() + for i_layer in range(self.num_layers): + layer = BasicLayer( + dim=int(embed_dim * 2 ** i_layer), + depth=depths[i_layer], + num_heads=num_heads[i_layer], + window_size=window_size, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])], + norm_layer=norm_layer, + downsample=PatchMerging if (i_layer < self.num_layers - 1) else None, + use_checkpoint=use_checkpoint) + self.layers.append(layer) + + num_features = [int(embed_dim * 2 ** i) for i in range(self.num_layers)] + self.num_features = num_features + + # add a norm layer for each output + for i_layer in out_indices: + layer = norm_layer(num_features[i_layer]) + layer_name = f'norm{i_layer}' + self.add_module(layer_name, layer) + + self._freeze_stages() + + def _freeze_stages(self): + if self.frozen_stages >= 0: + self.patch_embed.eval() + for param in self.patch_embed.parameters(): + param.requires_grad = False + + if self.frozen_stages >= 1 and self.ape: + self.absolute_pos_embed.requires_grad = False + + if self.frozen_stages >= 2: + self.pos_drop.eval() + for i in range(0, self.frozen_stages - 1): + m = self.layers[i] + m.eval() + for param in m.parameters(): + param.requires_grad = False + + + def forward(self, x): + + x = self.patch_embed(x) + + Wh, Ww = x.size(2), x.size(3) + if self.ape: + # interpolate the position embedding to the corresponding size + absolute_pos_embed = F.interpolate(self.absolute_pos_embed, size=(Wh, Ww), mode='bicubic') + x = (x + absolute_pos_embed) # B Wh*Ww C + + outs = [x.contiguous()] + x = x.flatten(2).transpose(1, 2) + x = self.pos_drop(x) + + + for i in range(self.num_layers): + layer = self.layers[i] + x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww) + + + if i in self.out_indices: + norm_layer = getattr(self, f'norm{i}') + x_out = norm_layer(x_out) + + out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous() + outs.append(out) + + + + return tuple(outs) + + + + + + + +def get_activation_fn(activation): + """Return an activation function given a string""" + if activation == "gelu": + return F.gelu + + raise RuntimeError(F"activation should be gelu, not {activation}.") + + +def make_cbr(in_dim, out_dim): + return nn.Sequential(nn.Conv2d(in_dim, out_dim, kernel_size=3, padding=1), nn.InstanceNorm2d(out_dim), nn.GELU()) + + +def make_cbg(in_dim, out_dim): + return nn.Sequential(nn.Conv2d(in_dim, out_dim, kernel_size=3, padding=1), nn.InstanceNorm2d(out_dim), nn.GELU()) + + +def rescale_to(x, scale_factor: float = 2, interpolation='nearest'): + return F.interpolate(x, scale_factor=scale_factor, mode=interpolation) + + +def resize_as(x, y, interpolation='bilinear'): + return F.interpolate(x, size=y.shape[-2:], mode=interpolation) + + +def image2patches(x): + """b c (hg h) (wg w) -> (hg wg b) c h w""" + x = rearrange(x, 'b c (hg h) (wg w) -> (hg wg b) c h w', hg=2, wg=2) + return x + + +def patches2image(x): + """(hg wg b) c h w -> b c (hg h) (wg w)""" + x = rearrange(x, '(hg wg b) c h w -> b c (hg h) (wg w)', hg=2, wg=2) + return x +class PositionEmbeddingSine: + def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None): + super().__init__() + self.num_pos_feats = num_pos_feats + self.temperature = temperature + self.normalize = normalize + if scale is not None and normalize is False: + raise ValueError("normalize should be True if scale is passed") + if scale is None: + scale = 2 * math.pi + self.scale = scale + self.dim_t = torch.arange(0, self.num_pos_feats, dtype=torch.float32) + + def __call__(self, b, h, w): + device = self.dim_t.device + mask = torch.zeros([b, h, w], dtype=torch.bool, device=device) + assert mask is not None + not_mask = ~mask + y_embed = not_mask.cumsum(dim=1, dtype=torch.float32) + x_embed = not_mask.cumsum(dim=2, dtype=torch.float32) + if self.normalize: + eps = 1e-6 + y_embed = (y_embed - 0.5) / (y_embed[:, -1:, :] + eps) * self.scale + x_embed = (x_embed - 0.5) / (x_embed[:, :, -1:] + eps) * self.scale + + dim_t = self.temperature ** (2 * (self.dim_t.to(device) // 2) / self.num_pos_feats) + pos_x = x_embed[:, :, :, None] / dim_t + pos_y = y_embed[:, :, :, None] / dim_t + + pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3) + pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3) + + return torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) + + +class MCLM(nn.Module): + def __init__(self, d_model, num_heads, pool_ratios=[1, 4, 8]): + super(MCLM, self).__init__() + self.attention = nn.ModuleList([ + nn.MultiheadAttention(d_model, num_heads, dropout=0.1), + nn.MultiheadAttention(d_model, num_heads, dropout=0.1), + nn.MultiheadAttention(d_model, num_heads, dropout=0.1), + nn.MultiheadAttention(d_model, num_heads, dropout=0.1), + nn.MultiheadAttention(d_model, num_heads, dropout=0.1) + ]) + + self.linear1 = nn.Linear(d_model, d_model * 2) + self.linear2 = nn.Linear(d_model * 2, d_model) + self.linear3 = nn.Linear(d_model, d_model * 2) + self.linear4 = nn.Linear(d_model * 2, d_model) + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + self.dropout = nn.Dropout(0.1) + self.dropout1 = nn.Dropout(0.1) + self.dropout2 = nn.Dropout(0.1) + self.activation = get_activation_fn('gelu') + self.pool_ratios = pool_ratios + self.p_poses = [] + self.g_pos = None + self.positional_encoding = PositionEmbeddingSine(num_pos_feats=d_model // 2, normalize=True) + + def forward(self, l, g): + """ + l: 4,c,h,w + g: 1,c,h,w + """ + b, c, h, w = l.size() + # 4,c,h,w -> 1,c,2h,2w + concated_locs = rearrange(l, '(hg wg b) c h w -> b c (hg h) (wg w)', hg=2, wg=2) + + pools = [] + for pool_ratio in self.pool_ratios: + # b,c,h,w + tgt_hw = (round(h / pool_ratio), round(w / pool_ratio)) + pool = F.adaptive_avg_pool2d(concated_locs, tgt_hw) + pools.append(rearrange(pool, 'b c h w -> (h w) b c')) + if self.g_pos is None: + pos_emb = self.positional_encoding(pool.shape[0], pool.shape[2], pool.shape[3]) + pos_emb = rearrange(pos_emb, 'b c h w -> (h w) b c') + self.p_poses.append(pos_emb) + pools = torch.cat(pools, 0) + if self.g_pos is None: + self.p_poses = torch.cat(self.p_poses, dim=0) + pos_emb = self.positional_encoding(g.shape[0], g.shape[2], g.shape[3]) + self.g_pos = rearrange(pos_emb, 'b c h w -> (h w) b c') + + device = pools.device + self.p_poses = self.p_poses.to(device) + self.g_pos = self.g_pos.to(device) + + + # attention between glb (q) & multisensory concated-locs (k,v) + g_hw_b_c = rearrange(g, 'b c h w -> (h w) b c') + + + g_hw_b_c = g_hw_b_c + self.dropout1(self.attention[0](g_hw_b_c + self.g_pos, pools + self.p_poses, pools)[0]) + g_hw_b_c = self.norm1(g_hw_b_c) + g_hw_b_c = g_hw_b_c + self.dropout2(self.linear2(self.dropout(self.activation(self.linear1(g_hw_b_c)).clone()))) + g_hw_b_c = self.norm2(g_hw_b_c) + + # attention between origin locs (q) & freashed glb (k,v) + l_hw_b_c = rearrange(l, "b c h w -> (h w) b c") + _g_hw_b_c = rearrange(g_hw_b_c, '(h w) b c -> h w b c', h=h, w=w) + _g_hw_b_c = rearrange(_g_hw_b_c, "(ng h) (nw w) b c -> (h w) (ng nw b) c", ng=2, nw=2) + outputs_re = [] + for i, (_l, _g) in enumerate(zip(l_hw_b_c.chunk(4, dim=1), _g_hw_b_c.chunk(4, dim=1))): + outputs_re.append(self.attention[i + 1](_l, _g, _g)[0]) # (h w) 1 c + outputs_re = torch.cat(outputs_re, 1) # (h w) 4 c + + l_hw_b_c = l_hw_b_c + self.dropout1(outputs_re) + l_hw_b_c = self.norm1(l_hw_b_c) + l_hw_b_c = l_hw_b_c + self.dropout2(self.linear4(self.dropout(self.activation(self.linear3(l_hw_b_c)).clone()))) + l_hw_b_c = self.norm2(l_hw_b_c) + + l = torch.cat((l_hw_b_c, g_hw_b_c), 1) # hw,b(5),c + return rearrange(l, "(h w) b c -> b c h w", h=h, w=w) ## (5,c,h*w) + + + + + + + + + +class MCRM(nn.Module): + def __init__(self, d_model, num_heads, pool_ratios=[4, 8, 16], h=None): + super(MCRM, self).__init__() + self.attention = nn.ModuleList([ + nn.MultiheadAttention(d_model, num_heads, dropout=0.1), + nn.MultiheadAttention(d_model, num_heads, dropout=0.1), + nn.MultiheadAttention(d_model, num_heads, dropout=0.1), + nn.MultiheadAttention(d_model, num_heads, dropout=0.1) + ]) + self.linear3 = nn.Linear(d_model, d_model * 2) + self.linear4 = nn.Linear(d_model * 2, d_model) + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + self.dropout = nn.Dropout(0.1) + self.dropout1 = nn.Dropout(0.1) + self.dropout2 = nn.Dropout(0.1) + self.sigmoid = nn.Sigmoid() + self.activation = get_activation_fn('gelu') + self.sal_conv = nn.Conv2d(d_model, 1, 1) + self.pool_ratios = pool_ratios + + def forward(self, x): + device = x.device + b, c, h, w = x.size() + loc, glb = x.split([4, 1], dim=0) # 4,c,h,w; 1,c,h,w + + patched_glb = rearrange(glb, 'b c (hg h) (wg w) -> (hg wg b) c h w', hg=2, wg=2) + + token_attention_map = self.sigmoid(self.sal_conv(glb)) + token_attention_map = F.interpolate(token_attention_map, size=patches2image(loc).shape[-2:], mode='nearest') + loc = loc * rearrange(token_attention_map, 'b c (hg h) (wg w) -> (hg wg b) c h w', hg=2, wg=2) + + pools = [] + for pool_ratio in self.pool_ratios: + tgt_hw = (round(h / pool_ratio), round(w / pool_ratio)) + pool = F.adaptive_avg_pool2d(patched_glb, tgt_hw) + pools.append(rearrange(pool, 'nl c h w -> nl c (h w)')) # nl(4),c,hw + + pools = rearrange(torch.cat(pools, 2), "nl c nphw -> nl nphw 1 c") + loc_ = rearrange(loc, 'nl c h w -> nl (h w) 1 c') + + outputs = [] + for i, q in enumerate(loc_.unbind(dim=0)): # traverse all local patches + v = pools[i] + k = v + outputs.append(self.attention[i](q, k, v)[0]) + + outputs = torch.cat(outputs, 1) + src = loc.view(4, c, -1).permute(2, 0, 1) + self.dropout1(outputs) + src = self.norm1(src) + src = src + self.dropout2(self.linear4(self.dropout(self.activation(self.linear3(src)).clone()))) + src = self.norm2(src) + src = src.permute(1, 2, 0).reshape(4, c, h, w) # freshed loc + glb = glb + F.interpolate(patches2image(src), size=glb.shape[-2:], mode='nearest') # freshed glb + + return torch.cat((src, glb), 0), token_attention_map + + +class BEN_Base(nn.Module): + def __init__(self): + super().__init__() + + self.backbone = SwinTransformer(embed_dim=128, depths=[2, 2, 18, 2], num_heads=[4, 8, 16, 32], window_size=12) + emb_dim = 128 + self.sideout5 = nn.Sequential(nn.Conv2d(emb_dim, 1, kernel_size=3, padding=1)) + self.sideout4 = nn.Sequential(nn.Conv2d(emb_dim, 1, kernel_size=3, padding=1)) + self.sideout3 = nn.Sequential(nn.Conv2d(emb_dim, 1, kernel_size=3, padding=1)) + self.sideout2 = nn.Sequential(nn.Conv2d(emb_dim, 1, kernel_size=3, padding=1)) + self.sideout1 = nn.Sequential(nn.Conv2d(emb_dim, 1, kernel_size=3, padding=1)) + + self.output5 = make_cbr(1024, emb_dim) + self.output4 = make_cbr(512, emb_dim) + self.output3 = make_cbr(256, emb_dim) + self.output2 = make_cbr(128, emb_dim) + self.output1 = make_cbr(128, emb_dim) + + self.multifieldcrossatt = MCLM(emb_dim, 1, [1, 4, 8]) + self.conv1 = make_cbr(emb_dim, emb_dim) + self.conv2 = make_cbr(emb_dim, emb_dim) + self.conv3 = make_cbr(emb_dim, emb_dim) + self.conv4 = make_cbr(emb_dim, emb_dim) + self.dec_blk1 = MCRM(emb_dim, 1, [2, 4, 8]) + self.dec_blk2 = MCRM(emb_dim, 1, [2, 4, 8]) + self.dec_blk3 = MCRM(emb_dim, 1, [2, 4, 8]) + self.dec_blk4 = MCRM(emb_dim, 1, [2, 4, 8]) + + self.insmask_head = nn.Sequential( + nn.Conv2d(emb_dim, 384, kernel_size=3, padding=1), + nn.InstanceNorm2d(384), + nn.GELU(), + nn.Conv2d(384, 384, kernel_size=3, padding=1), + nn.InstanceNorm2d(384), + nn.GELU(), + nn.Conv2d(384, emb_dim, kernel_size=3, padding=1) + ) + + self.shallow = nn.Sequential(nn.Conv2d(3, emb_dim, kernel_size=3, padding=1)) + self.upsample1 = make_cbg(emb_dim, emb_dim) + self.upsample2 = make_cbg(emb_dim, emb_dim) + self.output = nn.Sequential(nn.Conv2d(emb_dim, 1, kernel_size=3, padding=1)) + + for m in self.modules(): + if isinstance(m, nn.GELU) or isinstance(m, nn.Dropout): + m.inplace = True + + def forward(self, x): + device = x.device + shallow = self.shallow(x) + glb = rescale_to(x, scale_factor=0.5, interpolation='bilinear') + loc = image2patches(x) + input = torch.cat((loc, glb), dim=0) + feature = self.backbone(input) + e5 = self.output5(feature[4]) # (5,128,16,16) + e4 = self.output4(feature[3]) # (5,128,32,32) + e3 = self.output3(feature[2]) # (5,128,64,64) + e2 = self.output2(feature[1]) # (5,128,128,128) + e1 = self.output1(feature[0]) # (5,128,128,128) + loc_e5, glb_e5 = e5.split([4, 1], dim=0) + e5 = self.multifieldcrossatt(loc_e5, glb_e5) # (4,128,16,16) + + e4, tokenattmap4 = self.dec_blk4(e4 + resize_as(e5, e4)) + e4 = self.conv4(e4) + e3, tokenattmap3 = self.dec_blk3(e3 + resize_as(e4, e3)) + e3 = self.conv3(e3) + e2, tokenattmap2 = self.dec_blk2(e2 + resize_as(e3, e2)) + e2 = self.conv2(e2) + e1, tokenattmap1 = self.dec_blk1(e1 + resize_as(e2, e1)) + e1 = self.conv1(e1) + loc_e1, glb_e1 = e1.split([4, 1], dim=0) + output1_cat = patches2image(loc_e1) # (1,128,256,256) + output1_cat = output1_cat + resize_as(glb_e1, output1_cat) + final_output = self.insmask_head(output1_cat) # (1,128,256,256) + final_output = final_output + resize_as(shallow, final_output) + final_output = self.upsample1(rescale_to(final_output)) + final_output = rescale_to(final_output + resize_as(shallow, final_output)) + final_output = self.upsample2(final_output) + final_output = self.output(final_output) + + return final_output.sigmoid() + + @torch.no_grad() + def inference(self,image): + image, h, w,original_image = rgb_loader_refiner(image) + + img_tensor = img_transform(image).unsqueeze(0).to(next(self.parameters()).device) + + res = self.forward(img_tensor) + + pred_array = postprocess_image(res, im_size=[w, h]) + + mask_image = Image.fromarray(pred_array, mode='L') + + blurred_mask = mask_image.filter(ImageFilter.GaussianBlur(radius=1)) + + original_image_rgba = original_image.convert("RGBA") + + foreground = original_image_rgba.copy() + + foreground.putalpha(blurred_mask) + + return blurred_mask, foreground + + def loadcheckpoints(self,model_path): + model_dict = torch.load(model_path, map_location="cpu", weights_only=True) + self.load_state_dict(model_dict['model_state_dict'], strict=True) + del model_path + + + + +def rgb_loader_refiner( original_image): + h, w = original_image.size + # # Apply EXIF orientation + image = ImageOps.exif_transpose(original_image) + # Convert to RGB if necessary + if image.mode != 'RGB': + image = image.convert('RGB') + + # Resize the image + image = image.resize((1024, 1024), resample=Image.LANCZOS) + + return image.convert('RGB'), h, w,original_image + +# Define the image transformation +img_transform = transforms.Compose([ + transforms.ToTensor(), + transforms.ConvertImageDtype(torch.float32), + transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) +]) + +def postprocess_image(result: torch.Tensor, im_size: list) -> np.ndarray: + result = torch.squeeze(F.interpolate(result, size=im_size, mode='bilinear'), 0) + ma = torch.max(result) + mi = torch.min(result) + result = (result - mi) / (ma - mi) + im_array = (result * 255).permute(1, 2, 0).cpu().data.numpy().astype(np.uint8) + im_array = np.squeeze(im_array) + return im_array + + + + diff --git a/models/RMBG/BEN2/BEN2.py b/models/RMBG/BEN2/BEN2.py new file mode 100644 index 0000000000000000000000000000000000000000..6413d77174f96dbd6bd39ac4632ceee94edbda3c --- /dev/null +++ b/models/RMBG/BEN2/BEN2.py @@ -0,0 +1,1401 @@ + +import math +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +import torch.utils.checkpoint as checkpoint +import numpy as np +from timm.models.layers import DropPath, to_2tuple, trunc_normal_ +from PIL import Image, ImageOps +from torchvision import transforms +import numpy as np +import random +import cv2 +import os +import subprocess +import time +import tempfile + + + + +def set_random_seed(seed): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False +set_random_seed(9) + + +torch.set_float32_matmul_precision('highest') + + + +class Mlp(nn.Module): + """ Multilayer perceptron.""" + + def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +def window_partition(x, window_size): + """ + Args: + x: (B, H, W, C) + window_size (int): window size + Returns: + windows: (num_windows*B, window_size, window_size, C) + """ + B, H, W, C = x.shape + x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) + windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) + return windows + + +def window_reverse(windows, window_size, H, W): + """ + Args: + windows: (num_windows*B, window_size, window_size, C) + window_size (int): Window size + H (int): Height of image + W (int): Width of image + Returns: + x: (B, H, W, C) + """ + B = int(windows.shape[0] / (H * W / window_size / window_size)) + x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1) + x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) + return x + + +class WindowAttention(nn.Module): + """ Window based multi-head self attention (W-MSA) module with relative position bias. + It supports both of shifted and non-shifted window. + Args: + dim (int): Number of input channels. + window_size (tuple[int]): The height and width of the window. + num_heads (int): Number of attention heads. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set + attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0 + proj_drop (float, optional): Dropout ratio of output. Default: 0.0 + """ + + def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.): + + super().__init__() + self.dim = dim + self.window_size = window_size # Wh, Ww + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = qk_scale or head_dim ** -0.5 + + # define a parameter table of relative position bias + self.relative_position_bias_table = nn.Parameter( + torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)) # 2*Wh-1 * 2*Ww-1, nH + + # get pair-wise relative position index for each token inside the window + coords_h = torch.arange(self.window_size[0]) + coords_w = torch.arange(self.window_size[1]) + coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww + coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww + relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww + relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2 + relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0 + relative_coords[:, :, 1] += self.window_size[1] - 1 + relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1 + relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww + self.register_buffer("relative_position_index", relative_position_index) + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + trunc_normal_(self.relative_position_bias_table, std=.02) + self.softmax = nn.Softmax(dim=-1) + + def forward(self, x, mask=None): + """ Forward function. + Args: + x: input features with shape of (num_windows*B, N, C) + mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None + """ + B_, N, C = x.shape + qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple) + + q = q * self.scale + attn = (q @ k.transpose(-2, -1)) + + relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view( + self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1) # Wh*Ww,Wh*Ww,nH + relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww + attn = attn + relative_position_bias.unsqueeze(0) + + if mask is not None: + nW = mask.shape[0] + attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0) + attn = attn.view(-1, self.num_heads, N, N) + attn = self.softmax(attn) + else: + attn = self.softmax(attn) + + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B_, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class SwinTransformerBlock(nn.Module): + """ Swin Transformer Block. + Args: + dim (int): Number of input channels. + num_heads (int): Number of attention heads. + window_size (int): Window size. + shift_size (int): Shift size for SW-MSA. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float, optional): Stochastic depth rate. Default: 0.0 + act_layer (nn.Module, optional): Activation layer. Default: nn.GELU + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + """ + + def __init__(self, dim, num_heads, window_size=7, shift_size=0, + mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0., + act_layer=nn.GELU, norm_layer=nn.LayerNorm): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.window_size = window_size + self.shift_size = shift_size + self.mlp_ratio = mlp_ratio + assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size" + + self.norm1 = norm_layer(dim) + self.attn = WindowAttention( + dim, window_size=to_2tuple(self.window_size), num_heads=num_heads, + qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop) + + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + + self.H = None + self.W = None + + def forward(self, x, mask_matrix): + """ Forward function. + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + mask_matrix: Attention mask for cyclic shift. + """ + B, L, C = x.shape + H, W = self.H, self.W + assert L == H * W, "input feature has wrong size" + + shortcut = x + x = self.norm1(x) + x = x.view(B, H, W, C) + + # pad feature maps to multiples of window size + pad_l = pad_t = 0 + pad_r = (self.window_size - W % self.window_size) % self.window_size + pad_b = (self.window_size - H % self.window_size) % self.window_size + x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b)) + _, Hp, Wp, _ = x.shape + + # cyclic shift + if self.shift_size > 0: + shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)) + attn_mask = mask_matrix + else: + shifted_x = x + attn_mask = None + + # partition windows + x_windows = window_partition(shifted_x, self.window_size) # nW*B, window_size, window_size, C + x_windows = x_windows.view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C + + # W-MSA/SW-MSA + attn_windows = self.attn(x_windows, mask=attn_mask) # nW*B, window_size*window_size, C + + # merge windows + attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C) + shifted_x = window_reverse(attn_windows, self.window_size, Hp, Wp) # B H' W' C + + # reverse cyclic shift + if self.shift_size > 0: + x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2)) + else: + x = shifted_x + + if pad_r > 0 or pad_b > 0: + x = x[:, :H, :W, :].contiguous() + + x = x.view(B, H * W, C) + + # FFN + x = shortcut + self.drop_path(x) + x = x + self.drop_path(self.mlp(self.norm2(x))) + + return x + + +class PatchMerging(nn.Module): + """ Patch Merging Layer + Args: + dim (int): Number of input channels. + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + """ + def __init__(self, dim, norm_layer=nn.LayerNorm): + super().__init__() + self.dim = dim + self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False) + self.norm = norm_layer(4 * dim) + + def forward(self, x, H, W): + """ Forward function. + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + """ + B, L, C = x.shape + assert L == H * W, "input feature has wrong size" + + x = x.view(B, H, W, C) + + # padding + pad_input = (H % 2 == 1) or (W % 2 == 1) + if pad_input: + x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2)) + + x0 = x[:, 0::2, 0::2, :] # B H/2 W/2 C + x1 = x[:, 1::2, 0::2, :] # B H/2 W/2 C + x2 = x[:, 0::2, 1::2, :] # B H/2 W/2 C + x3 = x[:, 1::2, 1::2, :] # B H/2 W/2 C + x = torch.cat([x0, x1, x2, x3], -1) # B H/2 W/2 4*C + x = x.view(B, -1, 4 * C) # B H/2*W/2 4*C + + x = self.norm(x) + x = self.reduction(x) + + return x + + +class BasicLayer(nn.Module): + """ A basic Swin Transformer layer for one stage. + Args: + dim (int): Number of feature channels + depth (int): Depths of this stage. + num_heads (int): Number of attention head. + window_size (int): Local window size. Default: 7. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. + """ + + def __init__(self, + dim, + depth, + num_heads, + window_size=7, + mlp_ratio=4., + qkv_bias=True, + qk_scale=None, + drop=0., + attn_drop=0., + drop_path=0., + norm_layer=nn.LayerNorm, + downsample=None, + use_checkpoint=False): + super().__init__() + self.window_size = window_size + self.shift_size = window_size // 2 + self.depth = depth + self.use_checkpoint = use_checkpoint + + # build blocks + self.blocks = nn.ModuleList([ + SwinTransformerBlock( + dim=dim, + num_heads=num_heads, + window_size=window_size, + shift_size=0 if (i % 2 == 0) else window_size // 2, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop, + attn_drop=attn_drop, + drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path, + norm_layer=norm_layer) + for i in range(depth)]) + + # patch merging layer + if downsample is not None: + self.downsample = downsample(dim=dim, norm_layer=norm_layer) + else: + self.downsample = None + + def forward(self, x, H, W): + """ Forward function. + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + """ + + # calculate attention mask for SW-MSA + Hp = int(np.ceil(H / self.window_size)) * self.window_size + Wp = int(np.ceil(W / self.window_size)) * self.window_size + img_mask = torch.zeros((1, Hp, Wp, 1), device=x.device) # 1 Hp Wp 1 + h_slices = (slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None)) + w_slices = (slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None)) + cnt = 0 + for h in h_slices: + for w in w_slices: + img_mask[:, h, w, :] = cnt + cnt += 1 + + mask_windows = window_partition(img_mask, self.window_size) # nW, window_size, window_size, 1 + mask_windows = mask_windows.view(-1, self.window_size * self.window_size) + attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) + attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0)) + + for blk in self.blocks: + blk.H, blk.W = H, W + if self.use_checkpoint: + x = checkpoint.checkpoint(blk, x, attn_mask) + else: + x = blk(x, attn_mask) + if self.downsample is not None: + x_down = self.downsample(x, H, W) + Wh, Ww = (H + 1) // 2, (W + 1) // 2 + return x, H, W, x_down, Wh, Ww + else: + return x, H, W, x, H, W + + +class PatchEmbed(nn.Module): + """ Image to Patch Embedding + Args: + patch_size (int): Patch token size. Default: 4. + in_chans (int): Number of input image channels. Default: 3. + embed_dim (int): Number of linear projection output channels. Default: 96. + norm_layer (nn.Module, optional): Normalization layer. Default: None + """ + + def __init__(self, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None): + super().__init__() + patch_size = to_2tuple(patch_size) + self.patch_size = patch_size + + self.in_chans = in_chans + self.embed_dim = embed_dim + + self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) + if norm_layer is not None: + self.norm = norm_layer(embed_dim) + else: + self.norm = None + + def forward(self, x): + """Forward function.""" + # padding + _, _, H, W = x.size() + if W % self.patch_size[1] != 0: + x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1])) + if H % self.patch_size[0] != 0: + x = F.pad(x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0])) + + x = self.proj(x) # B C Wh Ww + if self.norm is not None: + Wh, Ww = x.size(2), x.size(3) + x = x.flatten(2).transpose(1, 2) + x = self.norm(x) + x = x.transpose(1, 2).view(-1, self.embed_dim, Wh, Ww) + + return x + + +class SwinTransformer(nn.Module): + """ Swin Transformer backbone. + A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows` - + https://arxiv.org/pdf/2103.14030 + Args: + pretrain_img_size (int): Input image size for training the pretrained model, + used in absolute postion embedding. Default 224. + patch_size (int | tuple(int)): Patch size. Default: 4. + in_chans (int): Number of input image channels. Default: 3. + embed_dim (int): Number of linear projection output channels. Default: 96. + depths (tuple[int]): Depths of each Swin Transformer stage. + num_heads (tuple[int]): Number of attention head of each stage. + window_size (int): Window size. Default: 7. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4. + qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. + drop_rate (float): Dropout rate. + attn_drop_rate (float): Attention dropout rate. Default: 0. + drop_path_rate (float): Stochastic depth rate. Default: 0.2. + norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm. + ape (bool): If True, add absolute position embedding to the patch embedding. Default: False. + patch_norm (bool): If True, add normalization after patch embedding. Default: True. + out_indices (Sequence[int]): Output from which stages. + frozen_stages (int): Stages to be frozen (stop grad and set eval mode). + -1 means not freezing any parameters. + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. + """ + + def __init__(self, + pretrain_img_size=224, + patch_size=4, + in_chans=3, + embed_dim=96, + depths=[2, 2, 6, 2], + num_heads=[3, 6, 12, 24], + window_size=7, + mlp_ratio=4., + qkv_bias=True, + qk_scale=None, + drop_rate=0., + attn_drop_rate=0., + drop_path_rate=0.2, + norm_layer=nn.LayerNorm, + ape=False, + patch_norm=True, + out_indices=(0, 1, 2, 3), + frozen_stages=-1, + use_checkpoint=False): + super().__init__() + + self.pretrain_img_size = pretrain_img_size + self.num_layers = len(depths) + self.embed_dim = embed_dim + self.ape = ape + self.patch_norm = patch_norm + self.out_indices = out_indices + self.frozen_stages = frozen_stages + + # split image into non-overlapping patches + self.patch_embed = PatchEmbed( + patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, + norm_layer=norm_layer if self.patch_norm else None) + + # absolute position embedding + if self.ape: + pretrain_img_size = to_2tuple(pretrain_img_size) + patch_size = to_2tuple(patch_size) + patches_resolution = [pretrain_img_size[0] // patch_size[0], pretrain_img_size[1] // patch_size[1]] + + self.absolute_pos_embed = nn.Parameter(torch.zeros(1, embed_dim, patches_resolution[0], patches_resolution[1])) + trunc_normal_(self.absolute_pos_embed, std=.02) + + self.pos_drop = nn.Dropout(p=drop_rate) + + # stochastic depth + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule + + # build layers + self.layers = nn.ModuleList() + for i_layer in range(self.num_layers): + layer = BasicLayer( + dim=int(embed_dim * 2 ** i_layer), + depth=depths[i_layer], + num_heads=num_heads[i_layer], + window_size=window_size, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])], + norm_layer=norm_layer, + downsample=PatchMerging if (i_layer < self.num_layers - 1) else None, + use_checkpoint=use_checkpoint) + self.layers.append(layer) + + num_features = [int(embed_dim * 2 ** i) for i in range(self.num_layers)] + self.num_features = num_features + + # add a norm layer for each output + for i_layer in out_indices: + layer = norm_layer(num_features[i_layer]) + layer_name = f'norm{i_layer}' + self.add_module(layer_name, layer) + + self._freeze_stages() + + def _freeze_stages(self): + if self.frozen_stages >= 0: + self.patch_embed.eval() + for param in self.patch_embed.parameters(): + param.requires_grad = False + + if self.frozen_stages >= 1 and self.ape: + self.absolute_pos_embed.requires_grad = False + + if self.frozen_stages >= 2: + self.pos_drop.eval() + for i in range(0, self.frozen_stages - 1): + m = self.layers[i] + m.eval() + for param in m.parameters(): + param.requires_grad = False + + + def forward(self, x): + + x = self.patch_embed(x) + + Wh, Ww = x.size(2), x.size(3) + if self.ape: + # interpolate the position embedding to the corresponding size + absolute_pos_embed = F.interpolate(self.absolute_pos_embed, size=(Wh, Ww), mode='bicubic') + x = (x + absolute_pos_embed) # B Wh*Ww C + + outs = [x.contiguous()] + x = x.flatten(2).transpose(1, 2) + x = self.pos_drop(x) + + + for i in range(self.num_layers): + layer = self.layers[i] + x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww) + + + if i in self.out_indices: + norm_layer = getattr(self, f'norm{i}') + x_out = norm_layer(x_out) + + out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous() + outs.append(out) + + + + return tuple(outs) + + + + + + + + +def get_activation_fn(activation): + """Return an activation function given a string""" + if activation == "gelu": + return F.gelu + + raise RuntimeError(F"activation should be gelu, not {activation}.") + + +def make_cbr(in_dim, out_dim): + return nn.Sequential(nn.Conv2d(in_dim, out_dim, kernel_size=3, padding=1), nn.InstanceNorm2d(out_dim), nn.GELU()) + + +def make_cbg(in_dim, out_dim): + return nn.Sequential(nn.Conv2d(in_dim, out_dim, kernel_size=3, padding=1), nn.InstanceNorm2d(out_dim), nn.GELU()) + + +def rescale_to(x, scale_factor: float = 2, interpolation='nearest'): + return F.interpolate(x, scale_factor=scale_factor, mode=interpolation) + + +def resize_as(x, y, interpolation='bilinear'): + return F.interpolate(x, size=y.shape[-2:], mode=interpolation) + + +def image2patches(x): + """b c (hg h) (wg w) -> (hg wg b) c h w""" + x = rearrange(x, 'b c (hg h) (wg w) -> (hg wg b) c h w', hg=2, wg=2 ) + return x + + +def patches2image(x): + """(hg wg b) c h w -> b c (hg h) (wg w)""" + x = rearrange(x, '(hg wg b) c h w -> b c (hg h) (wg w)', hg=2, wg=2) + return x + + + +class PositionEmbeddingSine: + def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None): + super().__init__() + self.num_pos_feats = num_pos_feats + self.temperature = temperature + self.normalize = normalize + if scale is not None and normalize is False: + raise ValueError("normalize should be True if scale is passed") + if scale is None: + scale = 2 * math.pi + self.scale = scale + self.dim_t = torch.arange(0, self.num_pos_feats, dtype=torch.float32) + + def __call__(self, b, h, w): + device = self.dim_t.device + mask = torch.zeros([b, h, w], dtype=torch.bool, device=device) + assert mask is not None + not_mask = ~mask + y_embed = not_mask.cumsum(dim=1, dtype=torch.float32) + x_embed = not_mask.cumsum(dim=2, dtype=torch.float32) + if self.normalize: + eps = 1e-6 + y_embed = (y_embed - 0.5) / (y_embed[:, -1:, :] + eps) * self.scale + x_embed = (x_embed - 0.5) / (x_embed[:, :, -1:] + eps) * self.scale + + dim_t = self.temperature ** (2 * (self.dim_t.to(device) // 2) / self.num_pos_feats) + pos_x = x_embed[:, :, :, None] / dim_t + pos_y = y_embed[:, :, :, None] / dim_t + + pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3) + pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3) + + return torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) + + + +class PositionEmbeddingSine: + def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None): + super().__init__() + self.num_pos_feats = num_pos_feats + self.temperature = temperature + self.normalize = normalize + if scale is not None and normalize is False: + raise ValueError("normalize should be True if scale is passed") + if scale is None: + scale = 2 * math.pi + self.scale = scale + self.dim_t = torch.arange(0, self.num_pos_feats, dtype=torch.float32) + + def __call__(self, b, h, w): + device = self.dim_t.device + mask = torch.zeros([b, h, w], dtype=torch.bool, device=device) + assert mask is not None + not_mask = ~mask + y_embed = not_mask.cumsum(dim=1, dtype=torch.float32) + x_embed = not_mask.cumsum(dim=2, dtype=torch.float32) + if self.normalize: + eps = 1e-6 + y_embed = (y_embed - 0.5) / (y_embed[:, -1:, :] + eps) * self.scale + x_embed = (x_embed - 0.5) / (x_embed[:, :, -1:] + eps) * self.scale + + dim_t = self.temperature ** (2 * (self.dim_t.to(device) // 2) / self.num_pos_feats) + pos_x = x_embed[:, :, :, None] / dim_t + pos_y = y_embed[:, :, :, None] / dim_t + + pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3) + pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3) + + return torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) + + +class MCLM(nn.Module): + def __init__(self, d_model, num_heads, pool_ratios=[1, 4, 8]): + super(MCLM, self).__init__() + self.attention = nn.ModuleList([ + nn.MultiheadAttention(d_model, num_heads, dropout=0.1), + nn.MultiheadAttention(d_model, num_heads, dropout=0.1), + nn.MultiheadAttention(d_model, num_heads, dropout=0.1), + nn.MultiheadAttention(d_model, num_heads, dropout=0.1), + nn.MultiheadAttention(d_model, num_heads, dropout=0.1) + ]) + + self.linear1 = nn.Linear(d_model, d_model * 2) + self.linear2 = nn.Linear(d_model * 2, d_model) + self.linear3 = nn.Linear(d_model, d_model * 2) + self.linear4 = nn.Linear(d_model * 2, d_model) + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + self.dropout = nn.Dropout(0.1) + self.dropout1 = nn.Dropout(0.1) + self.dropout2 = nn.Dropout(0.1) + self.activation = get_activation_fn('gelu') + self.pool_ratios = pool_ratios + self.p_poses = [] + self.g_pos = None + self.positional_encoding = PositionEmbeddingSine(num_pos_feats=d_model // 2, normalize=True) + + def forward(self, l, g): + """ + l: 4,c,h,w + g: 1,c,h,w + """ + self.p_poses = [] + self.g_pos = None + b, c, h, w = l.size() + # 4,c,h,w -> 1,c,2h,2w + concated_locs = rearrange(l, '(hg wg b) c h w -> b c (hg h) (wg w)', hg=2, wg=2) + + pools = [] + for pool_ratio in self.pool_ratios: + # b,c,h,w + tgt_hw = (round(h / pool_ratio), round(w / pool_ratio)) + pool = F.adaptive_avg_pool2d(concated_locs, tgt_hw) + pools.append(rearrange(pool, 'b c h w -> (h w) b c')) + if self.g_pos is None: + pos_emb = self.positional_encoding(pool.shape[0], pool.shape[2], pool.shape[3]) + pos_emb = rearrange(pos_emb, 'b c h w -> (h w) b c') + self.p_poses.append(pos_emb) + pools = torch.cat(pools, 0) + if self.g_pos is None: + self.p_poses = torch.cat(self.p_poses, dim=0) + pos_emb = self.positional_encoding(g.shape[0], g.shape[2], g.shape[3]) + self.g_pos = rearrange(pos_emb, 'b c h w -> (h w) b c') + + device = pools.device + self.p_poses = self.p_poses.to(device) + self.g_pos = self.g_pos.to(device) + + + # attention between glb (q) & multisensory concated-locs (k,v) + g_hw_b_c = rearrange(g, 'b c h w -> (h w) b c') + + + g_hw_b_c = g_hw_b_c + self.dropout1(self.attention[0](g_hw_b_c + self.g_pos, pools + self.p_poses, pools)[0]) + g_hw_b_c = self.norm1(g_hw_b_c) + g_hw_b_c = g_hw_b_c + self.dropout2(self.linear2(self.dropout(self.activation(self.linear1(g_hw_b_c)).clone()))) + g_hw_b_c = self.norm2(g_hw_b_c) + + # attention between origin locs (q) & freashed glb (k,v) + l_hw_b_c = rearrange(l, "b c h w -> (h w) b c") + _g_hw_b_c = rearrange(g_hw_b_c, '(h w) b c -> h w b c', h=h, w=w) + _g_hw_b_c = rearrange(_g_hw_b_c, "(ng h) (nw w) b c -> (h w) (ng nw b) c", ng=2, nw=2) + outputs_re = [] + for i, (_l, _g) in enumerate(zip(l_hw_b_c.chunk(4, dim=1), _g_hw_b_c.chunk(4, dim=1))): + outputs_re.append(self.attention[i + 1](_l, _g, _g)[0]) # (h w) 1 c + outputs_re = torch.cat(outputs_re, 1) # (h w) 4 c + + l_hw_b_c = l_hw_b_c + self.dropout1(outputs_re) + l_hw_b_c = self.norm1(l_hw_b_c) + l_hw_b_c = l_hw_b_c + self.dropout2(self.linear4(self.dropout(self.activation(self.linear3(l_hw_b_c)).clone()))) + l_hw_b_c = self.norm2(l_hw_b_c) + + l = torch.cat((l_hw_b_c, g_hw_b_c), 1) # hw,b(5),c + return rearrange(l, "(h w) b c -> b c h w", h=h, w=w) ## (5,c,h*w) + + + + + + + + + +class MCRM(nn.Module): + def __init__(self, d_model, num_heads, pool_ratios=[4, 8, 16], h=None): + super(MCRM, self).__init__() + self.attention = nn.ModuleList([ + nn.MultiheadAttention(d_model, num_heads, dropout=0.1), + nn.MultiheadAttention(d_model, num_heads, dropout=0.1), + nn.MultiheadAttention(d_model, num_heads, dropout=0.1), + nn.MultiheadAttention(d_model, num_heads, dropout=0.1) + ]) + self.linear3 = nn.Linear(d_model, d_model * 2) + self.linear4 = nn.Linear(d_model * 2, d_model) + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + self.dropout = nn.Dropout(0.1) + self.dropout1 = nn.Dropout(0.1) + self.dropout2 = nn.Dropout(0.1) + self.sigmoid = nn.Sigmoid() + self.activation = get_activation_fn('gelu') + self.sal_conv = nn.Conv2d(d_model, 1, 1) + self.pool_ratios = pool_ratios + + def forward(self, x): + device = x.device + b, c, h, w = x.size() + loc, glb = x.split([4, 1], dim=0) # 4,c,h,w; 1,c,h,w + + patched_glb = rearrange(glb, 'b c (hg h) (wg w) -> (hg wg b) c h w', hg=2, wg=2) + + token_attention_map = self.sigmoid(self.sal_conv(glb)) + token_attention_map = F.interpolate(token_attention_map, size=patches2image(loc).shape[-2:], mode='nearest') + loc = loc * rearrange(token_attention_map, 'b c (hg h) (wg w) -> (hg wg b) c h w', hg=2, wg=2) + + pools = [] + for pool_ratio in self.pool_ratios: + tgt_hw = (round(h / pool_ratio), round(w / pool_ratio)) + pool = F.adaptive_avg_pool2d(patched_glb, tgt_hw) + pools.append(rearrange(pool, 'nl c h w -> nl c (h w)')) # nl(4),c,hw + + pools = rearrange(torch.cat(pools, 2), "nl c nphw -> nl nphw 1 c") + loc_ = rearrange(loc, 'nl c h w -> nl (h w) 1 c') + + outputs = [] + for i, q in enumerate(loc_.unbind(dim=0)): # traverse all local patches + v = pools[i] + k = v + outputs.append(self.attention[i](q, k, v)[0]) + + outputs = torch.cat(outputs, 1) + src = loc.view(4, c, -1).permute(2, 0, 1) + self.dropout1(outputs) + src = self.norm1(src) + src = src + self.dropout2(self.linear4(self.dropout(self.activation(self.linear3(src)).clone()))) + src = self.norm2(src) + src = src.permute(1, 2, 0).reshape(4, c, h, w) # freshed loc + glb = glb + F.interpolate(patches2image(src), size=glb.shape[-2:], mode='nearest') # freshed glb + + return torch.cat((src, glb), 0), token_attention_map + + + +class BEN_Base(nn.Module): + def __init__(self): + super().__init__() + + self.backbone = SwinTransformer(embed_dim=128, depths=[2, 2, 18, 2], num_heads=[4, 8, 16, 32], window_size=12) + emb_dim = 128 + self.sideout5 = nn.Sequential(nn.Conv2d(emb_dim, 1, kernel_size=3, padding=1)) + self.sideout4 = nn.Sequential(nn.Conv2d(emb_dim, 1, kernel_size=3, padding=1)) + self.sideout3 = nn.Sequential(nn.Conv2d(emb_dim, 1, kernel_size=3, padding=1)) + self.sideout2 = nn.Sequential(nn.Conv2d(emb_dim, 1, kernel_size=3, padding=1)) + self.sideout1 = nn.Sequential(nn.Conv2d(emb_dim, 1, kernel_size=3, padding=1)) + + self.output5 = make_cbr(1024, emb_dim) + self.output4 = make_cbr(512, emb_dim) + self.output3 = make_cbr(256, emb_dim) + self.output2 = make_cbr(128, emb_dim) + self.output1 = make_cbr(128, emb_dim) + + self.multifieldcrossatt = MCLM(emb_dim, 1, [1, 4, 8]) + self.conv1 = make_cbr(emb_dim, emb_dim) + self.conv2 = make_cbr(emb_dim, emb_dim) + self.conv3 = make_cbr(emb_dim, emb_dim) + self.conv4 = make_cbr(emb_dim, emb_dim) + self.dec_blk1 = MCRM(emb_dim, 1, [2, 4, 8]) + self.dec_blk2 = MCRM(emb_dim, 1, [2, 4, 8]) + self.dec_blk3 = MCRM(emb_dim, 1, [2, 4, 8]) + self.dec_blk4 = MCRM(emb_dim, 1, [2, 4, 8]) + + self.insmask_head = nn.Sequential( + nn.Conv2d(emb_dim, 384, kernel_size=3, padding=1), + nn.InstanceNorm2d(384), + nn.GELU(), + nn.Conv2d(384, 384, kernel_size=3, padding=1), + nn.InstanceNorm2d(384), + nn.GELU(), + nn.Conv2d(384, emb_dim, kernel_size=3, padding=1) + ) + + self.shallow = nn.Sequential(nn.Conv2d(3, emb_dim, kernel_size=3, padding=1)) + self.upsample1 = make_cbg(emb_dim, emb_dim) + self.upsample2 = make_cbg(emb_dim, emb_dim) + self.output = nn.Sequential(nn.Conv2d(emb_dim, 1, kernel_size=3, padding=1)) + + for m in self.modules(): + if isinstance(m, nn.GELU) or isinstance(m, nn.Dropout): + m.inplace = True + + + + @torch.inference_mode() + @torch.autocast(device_type="cuda",dtype=torch.float16) + def forward(self, x): + real_batch = x.size(0) + + shallow_batch = self.shallow(x) + glb_batch = rescale_to(x, scale_factor=0.5, interpolation='bilinear') + + + + final_input = None + for i in range(real_batch): + start = i * 4 + end = (i + 1) * 4 + loc_batch = image2patches(x[i,:,:,:].unsqueeze(dim=0)) + input_ = torch.cat((loc_batch, glb_batch[i,:,:,:].unsqueeze(dim=0)), dim=0) + + + if final_input == None: + final_input= input_ + else: final_input = torch.cat((final_input, input_), dim=0) + + features = self.backbone(final_input) + outputs = [] + + for i in range(real_batch): + + start = i * 5 + end = (i + 1) * 5 + + f4 = features[4][start:end, :, :, :] # shape: [5, C, H, W] + f3 = features[3][start:end, :, :, :] + f2 = features[2][start:end, :, :, :] + f1 = features[1][start:end, :, :, :] + f0 = features[0][start:end, :, :, :] + e5 = self.output5(f4) + e4 = self.output4(f3) + e3 = self.output3(f2) + e2 = self.output2(f1) + e1 = self.output1(f0) + loc_e5, glb_e5 = e5.split([4, 1], dim=0) + e5 = self.multifieldcrossatt(loc_e5, glb_e5) # (4,128,16,16) + + + e4, tokenattmap4 = self.dec_blk4(e4 + resize_as(e5, e4)) + e4 = self.conv4(e4) + e3, tokenattmap3 = self.dec_blk3(e3 + resize_as(e4, e3)) + e3 = self.conv3(e3) + e2, tokenattmap2 = self.dec_blk2(e2 + resize_as(e3, e2)) + e2 = self.conv2(e2) + e1, tokenattmap1 = self.dec_blk1(e1 + resize_as(e2, e1)) + e1 = self.conv1(e1) + + loc_e1, glb_e1 = e1.split([4, 1], dim=0) + + output1_cat = patches2image(loc_e1) # (1,128,256,256) + + # add glb feat in + output1_cat = output1_cat + resize_as(glb_e1, output1_cat) + # merge + final_output = self.insmask_head(output1_cat) # (1,128,256,256) + # shallow feature merge + shallow = shallow_batch[i,:,:,:].unsqueeze(dim=0) + final_output = final_output + resize_as(shallow, final_output) + final_output = self.upsample1(rescale_to(final_output)) + final_output = rescale_to(final_output + resize_as(shallow, final_output)) + final_output = self.upsample2(final_output) + final_output = self.output(final_output) + mask = final_output.sigmoid() + outputs.append(mask) + + return torch.cat(outputs, dim=0) + + + + + def loadcheckpoints(self,model_path): + model_dict = torch.load(model_path, map_location="cpu", weights_only=True) + self.load_state_dict(model_dict['model_state_dict'], strict=True) + del model_path + + def inference(self,image,refine_foreground=False): + + set_random_seed(9) + # image = ImageOps.exif_transpose(image) + if isinstance(image, Image.Image): + image, h, w,original_image = rgb_loader_refiner(image) + if torch.cuda.is_available(): + + img_tensor = img_transform(image).unsqueeze(0).to(next(self.parameters()).device) + else: + img_tensor = img_transform32(image).unsqueeze(0).to(next(self.parameters()).device) + + + with torch.no_grad(): + res = self.forward(img_tensor) + + # Show Results + if refine_foreground == True: + + pred_pil = transforms.ToPILImage()(res.squeeze()) + image_masked = refine_foreground_process(original_image, pred_pil) + + image_masked.putalpha(pred_pil.resize(original_image.size)) + return image_masked + + else: + alpha = postprocess_image(res, im_size=[w,h]) + pred_pil = transforms.ToPILImage()(alpha) + mask = pred_pil.resize(original_image.size) + original_image.putalpha(mask) + # mask = Image.fromarray(alpha) + + return original_image + + + else: + foregrounds = [] + for batch in image: + image, h, w,original_image = rgb_loader_refiner(batch) + if torch.cuda.is_available(): + + img_tensor = img_transform(image).unsqueeze(0).to(next(self.parameters()).device) + else: + img_tensor = img_transform32(image).unsqueeze(0).to(next(self.parameters()).device) + + with torch.no_grad(): + res = self.forward(img_tensor) + + if refine_foreground == True: + + pred_pil = transforms.ToPILImage()(res.squeeze()) + image_masked = refine_foreground_process(original_image, pred_pil) + + image_masked.putalpha(pred_pil.resize(original_image.size)) + + foregrounds.append(image_masked) + else: + alpha = postprocess_image(res, im_size=[w,h]) + pred_pil = transforms.ToPILImage()(alpha) + mask = pred_pil.resize(original_image.size) + original_image.putalpha(mask) + # mask = Image.fromarray(alpha) + foregrounds.append(original_image) + + return foregrounds + + + + + def segment_video(self, video_path, output_path="./", fps=0, refine_foreground=False, batch=1, print_frames_processed=True, webm = False, rgb_value= (0, 255, 0)): + + """ + Segments the given video to extract the foreground (with alpha) from each frame + and saves the result as either a WebM video (with alpha channel) or MP4 (with a + color background). + + Args: + video_path (str): + Path to the input video file. + + output_path (str, optional): + Directory (or full path) where the output video and/or files will be saved. + Defaults to "./". + + fps (int, optional): + The frames per second (FPS) to use for the output video. If 0 (default), the + original FPS of the input video is used. Otherwise, overrides it. + + refine_foreground (bool, optional): + Whether to run an additional “refine foreground” process on each frame. + Defaults to False. + + batch (int, optional): + Number of frames to process at once (inference batch size). Large batch sizes + may require more GPU memory. Defaults to 1. + + print_frames_processed (bool, optional): + If True (default), prints progress (how many frames have been processed) to + the console. + + webm (bool, optional): + If True (default), exports a WebM video with alpha channel (VP9 / yuva420p). + If False, exports an MP4 video composited over a solid color background. + + rgb_value (tuple, optional): + The RGB background color (e.g., green screen) used to composite frames when + saving to MP4. Defaults to (0, 255, 0). + + Returns: + None. Writes the output video(s) to disk in the specified format. + """ + + + cap = cv2.VideoCapture(video_path) + if not cap.isOpened(): + raise IOError(f"Cannot open video: {video_path}") + + original_fps = cap.get(cv2.CAP_PROP_FPS) + original_fps = 30 if original_fps == 0 else original_fps + fps = original_fps if fps == 0 else fps + + ret, first_frame = cap.read() + if not ret: + raise ValueError("No frames found in the video.") + height, width = first_frame.shape[:2] + cap.set(cv2.CAP_PROP_POS_FRAMES, 0) + + foregrounds = [] + frame_idx = 0 + processed_count = 0 + batch_frames = [] + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + + while True: + ret, frame = cap.read() + if not ret: + if batch_frames: + batch_results = self.inference(batch_frames, refine_foreground) + if isinstance(batch_results, Image.Image): + foregrounds.append(batch_results) + else: + foregrounds.extend(batch_results) + if print_frames_processed: + print(f"Processed frames {frame_idx-len(batch_frames)+1} to {frame_idx} of {total_frames}") + break + + # Process every frame instead of using intervals + frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + pil_frame = Image.fromarray(frame_rgb) + batch_frames.append(pil_frame) + + if len(batch_frames) == batch: + batch_results = self.inference(batch_frames, refine_foreground) + if isinstance(batch_results, Image.Image): + foregrounds.append(batch_results) + else: + foregrounds.extend(batch_results) + if print_frames_processed: + print(f"Processed frames {frame_idx-batch+1} to {frame_idx} of {total_frames}") + batch_frames = [] + processed_count += batch + + frame_idx += 1 + + + if webm: + alpha_webm_path = os.path.join(output_path, "foreground.webm") + pil_images_to_webm_alpha(foregrounds, alpha_webm_path, fps=original_fps) + + else: + cap.release() + fg_output = os.path.join(output_path, 'foreground.mp4') + + pil_images_to_mp4(foregrounds, fg_output, fps=original_fps,rgb_value=rgb_value) + cv2.destroyAllWindows() + + try: + fg_audio_output = os.path.join(output_path, 'foreground_output_with_audio.mp4') + add_audio_to_video(fg_output, video_path, fg_audio_output) + except Exception as e: + print("No audio found in the original video") + print(e) + + + + + +def rgb_loader_refiner( original_image): + h, w = original_image.size + + image = original_image + # Convert to RGB if necessary + if image.mode != 'RGB': + image = image.convert('RGB') + + # Resize the image + image = image.resize((1024, 1024), resample=Image.LANCZOS) + + return image.convert('RGB'), h, w,original_image + +# Define the image transformation +img_transform = transforms.Compose([ + transforms.ToTensor(), + transforms.ConvertImageDtype(torch.float16), + transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) +]) + +img_transform32 = transforms.Compose([ + transforms.ToTensor(), + transforms.ConvertImageDtype(torch.float32), + transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) +]) + + + + + +def pil_images_to_mp4(images, output_path, fps=24, rgb_value=(0, 255, 0)): + """ + Converts an array of PIL images to an MP4 video. + + Args: + images: List of PIL images + output_path: Path to save the MP4 file + fps: Frames per second (default: 24) + rgb_value: Background RGB color tuple (default: green (0, 255, 0)) + """ + if not images: + raise ValueError("No images provided to convert to MP4.") + + width, height = images[0].size + fourcc = cv2.VideoWriter_fourcc(*'mp4v') + video_writer = cv2.VideoWriter(output_path, fourcc, fps, (width, height)) + + for image in images: + # If image has alpha channel, composite onto the specified background color + if image.mode == 'RGBA': + # Create background image with specified RGB color + background = Image.new('RGB', image.size, rgb_value) + background = background.convert('RGBA') + # Composite the image onto the background + image = Image.alpha_composite(background, image) + image = image.convert('RGB') + else: + # Ensure RGB format for non-alpha images + image = image.convert('RGB') + + # Convert to OpenCV format and write + open_cv_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) + video_writer.write(open_cv_image) + + video_writer.release() + +def pil_images_to_webm_alpha(images, output_path, fps=30): + """ + Converts a list of PIL RGBA images to a VP9 .webm video with alpha channel. + + NOTE: Not all players will display alpha in WebM. + Browsers like Chrome/Firefox typically do support VP9 alpha. + """ + if not images: + raise ValueError("No images provided for WebM with alpha.") + + # Ensure output directory exists + os.makedirs(os.path.dirname(output_path), exist_ok=True) + + with tempfile.TemporaryDirectory() as tmpdir: + # Save frames as PNG (with alpha) + for idx, img in enumerate(images): + if img.mode != "RGBA": + img = img.convert("RGBA") + out_path = os.path.join(tmpdir, f"{idx:06d}.png") + img.save(out_path, "PNG") + + # Construct ffmpeg command + # -c:v libvpx-vp9 => VP9 encoder + # -pix_fmt yuva420p => alpha-enabled pixel format + # -auto-alt-ref 0 => helps preserve alpha frames (libvpx quirk) + ffmpeg_cmd = [ + "ffmpeg", "-y", + "-framerate", str(fps), + "-i", os.path.join(tmpdir, "%06d.png"), + "-c:v", "libvpx-vp9", + "-pix_fmt", "yuva420p", + "-auto-alt-ref", "0", + output_path + ] + + subprocess.run(ffmpeg_cmd, check=True) + + print(f"WebM with alpha saved to {output_path}") + +def add_audio_to_video(video_without_audio_path, original_video_path, output_path): + """ + Check if the original video has an audio stream. If yes, add it. If not, skip. + """ + # 1) Probe original video for audio streams + probe_command = [ + 'ffprobe', '-v', 'error', + '-select_streams', 'a:0', + '-show_entries', 'stream=index', + '-of', 'csv=p=0', + original_video_path + ] + result = subprocess.run(probe_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) + + # result.stdout is empty if no audio stream found + if not result.stdout.strip(): + print("No audio track found in original video, skipping audio addition.") + return + + print("Audio track detected; proceeding to mux audio.") + # 2) If audio found, run ffmpeg to add it + command = [ + 'ffmpeg', '-y', + '-i', video_without_audio_path, + '-i', original_video_path, + '-c', 'copy', + '-map', '0:v:0', + '-map', '1:a:0', # we know there's an audio track now + output_path + ] + subprocess.run(command, check=True) + print(f"Audio added successfully => {output_path}") + + + + + +### Thanks to the source: https://huggingface.co/ZhengPeng7/BiRefNet/blob/main/handler.py +def refine_foreground_process(image, mask, r=90): + if mask.size != image.size: + mask = mask.resize(image.size) + image = np.array(image) / 255.0 + mask = np.array(mask) / 255.0 + estimated_foreground = FB_blur_fusion_foreground_estimator_2(image, mask, r=r) + image_masked = Image.fromarray((estimated_foreground * 255.0).astype(np.uint8)) + return image_masked + + +def FB_blur_fusion_foreground_estimator_2(image, alpha, r=90): + # Thanks to the source: https://github.com/Photoroom/fast-foreground-estimation + alpha = alpha[:, :, None] + F, blur_B = FB_blur_fusion_foreground_estimator(image, image, image, alpha, r) + return FB_blur_fusion_foreground_estimator(image, F, blur_B, alpha, r=6)[0] + + +def FB_blur_fusion_foreground_estimator(image, F, B, alpha, r=90): + if isinstance(image, Image.Image): + image = np.array(image) / 255.0 + blurred_alpha = cv2.blur(alpha, (r, r))[:, :, None] + + blurred_FA = cv2.blur(F * alpha, (r, r)) + blurred_F = blurred_FA / (blurred_alpha + 1e-5) + + blurred_B1A = cv2.blur(B * (1 - alpha), (r, r)) + blurred_B = blurred_B1A / ((1 - blurred_alpha) + 1e-5) + F = blurred_F + alpha * \ + (image - alpha * blurred_F - (1 - alpha) * blurred_B) + F = np.clip(F, 0, 1) + return F, blurred_B + + + +def postprocess_image(result: torch.Tensor, im_size: list) -> np.ndarray: + result = torch.squeeze(F.interpolate(result, size=im_size, mode='bilinear'), 0) + ma = torch.max(result) + mi = torch.min(result) + result = (result - mi) / (ma - mi) + im_array = (result * 255).permute(1, 2, 0).cpu().data.numpy().astype(np.uint8) + im_array = np.squeeze(im_array) + return im_array + + + + +def rgb_loader_refiner( original_image): + h, w = original_image.size + # # Apply EXIF orientation + + image = ImageOps.exif_transpose(original_image) + + if original_image.mode != 'RGB': + original_image = original_image.convert('RGB') + + image = original_image + # Convert to RGB if necessary + + # Resize the image + image = image.resize((1024, 1024), resample=Image.LANCZOS) + + return image, h, w,original_image + + + diff --git a/models/RMBG/BEN2/BEN2_Base.pth b/models/RMBG/BEN2/BEN2_Base.pth new file mode 100644 index 0000000000000000000000000000000000000000..d7549b2f1582f8662dd7b195bbdc426fcc985156 --- /dev/null +++ b/models/RMBG/BEN2/BEN2_Base.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:926144a876bda06f125555b4f5a239ece89dc6eb838a863700ca9bf192161a1c +size 1134584206 diff --git a/models/RMBG/BEN2/gitattributes b/models/RMBG/BEN2/gitattributes new file mode 100644 index 0000000000000000000000000000000000000000..a6344aac8c09253b3b630fb776ae94478aa0275b --- /dev/null +++ b/models/RMBG/BEN2/gitattributes @@ -0,0 +1,35 @@ +*.7z filter=lfs diff=lfs merge=lfs -text +*.arrow filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +*.bz2 filter=lfs diff=lfs merge=lfs -text +*.ckpt filter=lfs diff=lfs merge=lfs -text +*.ftz filter=lfs diff=lfs merge=lfs -text +*.gz filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text +*.joblib filter=lfs diff=lfs merge=lfs -text +*.lfs.* filter=lfs diff=lfs merge=lfs -text +*.mlmodel filter=lfs diff=lfs merge=lfs -text +*.model filter=lfs diff=lfs merge=lfs -text +*.msgpack filter=lfs diff=lfs merge=lfs -text +*.npy filter=lfs diff=lfs merge=lfs -text +*.npz filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.ot filter=lfs diff=lfs merge=lfs -text +*.parquet filter=lfs diff=lfs merge=lfs -text +*.pb filter=lfs diff=lfs merge=lfs -text +*.pickle filter=lfs diff=lfs merge=lfs -text +*.pkl filter=lfs diff=lfs merge=lfs -text +*.pt filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.rar filter=lfs diff=lfs merge=lfs -text +*.safetensors filter=lfs diff=lfs merge=lfs -text +saved_model/**/* filter=lfs diff=lfs merge=lfs -text +*.tar.* filter=lfs diff=lfs merge=lfs -text +*.tar filter=lfs diff=lfs merge=lfs -text +*.tflite filter=lfs diff=lfs merge=lfs -text +*.tgz filter=lfs diff=lfs merge=lfs -text +*.wasm filter=lfs diff=lfs merge=lfs -text +*.xz filter=lfs diff=lfs merge=lfs -text +*.zip filter=lfs diff=lfs merge=lfs -text +*.zst filter=lfs diff=lfs merge=lfs -text +*tfevents* filter=lfs diff=lfs merge=lfs -text diff --git a/models/RMBG/BiRefNet-HR/BiRefNet_config.py b/models/RMBG/BiRefNet-HR/BiRefNet_config.py new file mode 100644 index 0000000000000000000000000000000000000000..37c8ac58bec2f52dac34204978a7b61b69e3da76 --- /dev/null +++ b/models/RMBG/BiRefNet-HR/BiRefNet_config.py @@ -0,0 +1,11 @@ +from transformers import PretrainedConfig + +class BiRefNetConfig(PretrainedConfig): + model_type = "SegformerForSemanticSegmentation" + def __init__( + self, + bb_pretrained=False, + **kwargs + ): + self.bb_pretrained = bb_pretrained + super().__init__(**kwargs) diff --git a/models/RMBG/BiRefNet-HR/birefnet.py b/models/RMBG/BiRefNet-HR/birefnet.py new file mode 100644 index 0000000000000000000000000000000000000000..ae1a3600184ff6c8fd09be2d0247092f5a16f57b --- /dev/null +++ b/models/RMBG/BiRefNet-HR/birefnet.py @@ -0,0 +1,2248 @@ +### config.py + +import os +import math + + +class Config(): + def __init__(self) -> None: + # PATH settings + self.sys_home_dir = os.path.expanduser('~') # Make up your file system as: SYS_HOME_DIR/codes/dis/BiRefNet, SYS_HOME_DIR/datasets/dis/xx, SYS_HOME_DIR/weights/xx + + # TASK settings + self.task = ['DIS5K', 'COD', 'HRSOD', 'DIS5K+HRSOD+HRS10K', 'P3M-10k'][0] + self.training_set = { + 'DIS5K': ['DIS-TR', 'DIS-TR+DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4'][0], + 'COD': 'TR-COD10K+TR-CAMO', + 'HRSOD': ['TR-DUTS', 'TR-HRSOD', 'TR-UHRSD', 'TR-DUTS+TR-HRSOD', 'TR-DUTS+TR-UHRSD', 'TR-HRSOD+TR-UHRSD', 'TR-DUTS+TR-HRSOD+TR-UHRSD'][5], + 'DIS5K+HRSOD+HRS10K': 'DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4+DIS-TR+TE-HRS10K+TE-HRSOD+TE-UHRSD+TR-HRS10K+TR-HRSOD+TR-UHRSD', # leave DIS-VD for evaluation. + 'P3M-10k': 'TR-P3M-10k', + }[self.task] + self.prompt4loc = ['dense', 'sparse'][0] + + # Faster-Training settings + self.load_all = True + self.compile = True # 1. Trigger CPU memory leak in some extend, which is an inherent problem of PyTorch. + # Machines with > 70GB CPU memory can run the whole training on DIS5K with default setting. + # 2. Higher PyTorch version may fix it: https://github.com/pytorch/pytorch/issues/119607. + # 3. But compile in Pytorch > 2.0.1 seems to bring no acceleration for training. + self.precisionHigh = True + + # MODEL settings + self.ms_supervision = True + self.out_ref = self.ms_supervision and True + self.dec_ipt = True + self.dec_ipt_split = True + self.cxt_num = [0, 3][1] # multi-scale skip connections from encoder + self.mul_scl_ipt = ['', 'add', 'cat'][2] + self.dec_att = ['', 'ASPP', 'ASPPDeformable'][2] + self.squeeze_block = ['', 'BasicDecBlk_x1', 'ResBlk_x4', 'ASPP_x3', 'ASPPDeformable_x3'][1] + self.dec_blk = ['BasicDecBlk', 'ResBlk', 'HierarAttDecBlk'][0] + + # TRAINING settings + self.batch_size = 4 + self.IoU_finetune_last_epochs = [ + 0, + { + 'DIS5K': -50, + 'COD': -20, + 'HRSOD': -20, + 'DIS5K+HRSOD+HRS10K': -20, + 'P3M-10k': -20, + }[self.task] + ][1] # choose 0 to skip + self.lr = (1e-4 if 'DIS5K' in self.task else 1e-5) * math.sqrt(self.batch_size / 4) # DIS needs high lr to converge faster. Adapt the lr linearly + self.size = 1024 + self.num_workers = max(4, self.batch_size) # will be decrease to min(it, batch_size) at the initialization of the data_loader + + # Backbone settings + self.bb = [ + 'vgg16', 'vgg16bn', 'resnet50', # 0, 1, 2 + 'swin_v1_t', 'swin_v1_s', # 3, 4 + 'swin_v1_b', 'swin_v1_l', # 5-bs9, 6-bs4 + 'pvt_v2_b0', 'pvt_v2_b1', # 7, 8 + 'pvt_v2_b2', 'pvt_v2_b5', # 9-bs10, 10-bs5 + ][6] + self.lateral_channels_in_collection = { + 'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64], + 'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64], + 'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192], + 'swin_v1_t': [768, 384, 192, 96], 'swin_v1_s': [768, 384, 192, 96], + 'pvt_v2_b0': [256, 160, 64, 32], 'pvt_v2_b1': [512, 320, 128, 64], + }[self.bb] + if self.mul_scl_ipt == 'cat': + self.lateral_channels_in_collection = [channel * 2 for channel in self.lateral_channels_in_collection] + self.cxt = self.lateral_channels_in_collection[1:][::-1][-self.cxt_num:] if self.cxt_num else [] + + # MODEL settings - inactive + self.lat_blk = ['BasicLatBlk'][0] + self.dec_channels_inter = ['fixed', 'adap'][0] + self.refine = ['', 'itself', 'RefUNet', 'Refiner', 'RefinerPVTInChannels4'][0] + self.progressive_ref = self.refine and True + self.ender = self.progressive_ref and False + self.scale = self.progressive_ref and 2 + self.auxiliary_classification = False # Only for DIS5K, where class labels are saved in `dataset.py`. + self.refine_iteration = 1 + self.freeze_bb = False + self.model = [ + 'BiRefNet', + ][0] + if self.dec_blk == 'HierarAttDecBlk': + self.batch_size = 2 ** [0, 1, 2, 3, 4][2] + + # TRAINING settings - inactive + self.preproc_methods = ['flip', 'enhance', 'rotate', 'pepper', 'crop'][:4] + self.optimizer = ['Adam', 'AdamW'][1] + self.lr_decay_epochs = [1e5] # Set to negative N to decay the lr in the last N-th epoch. + self.lr_decay_rate = 0.5 + # Loss + self.lambdas_pix_last = { + # not 0 means opening this loss + # original rate -- 1 : 30 : 1.5 : 0.2, bce x 30 + 'bce': 30 * 1, # high performance + 'iou': 0.5 * 1, # 0 / 255 + 'iou_patch': 0.5 * 0, # 0 / 255, win_size = (64, 64) + 'mse': 150 * 0, # can smooth the saliency map + 'triplet': 3 * 0, + 'reg': 100 * 0, + 'ssim': 10 * 1, # help contours, + 'cnt': 5 * 0, # help contours + 'structure': 5 * 0, # structure loss from codes of MVANet. A little improvement on DIS-TE[1,2,3], a bit more decrease on DIS-TE4. + } + self.lambdas_cls = { + 'ce': 5.0 + } + # Adv + self.lambda_adv_g = 10. * 0 # turn to 0 to avoid adv training + self.lambda_adv_d = 3. * (self.lambda_adv_g > 0) + + # PATH settings - inactive + self.data_root_dir = os.path.join(self.sys_home_dir, 'datasets/dis') + self.weights_root_dir = os.path.join(self.sys_home_dir, 'weights') + self.weights = { + 'pvt_v2_b2': os.path.join(self.weights_root_dir, 'pvt_v2_b2.pth'), + 'pvt_v2_b5': os.path.join(self.weights_root_dir, ['pvt_v2_b5.pth', 'pvt_v2_b5_22k.pth'][0]), + 'swin_v1_b': os.path.join(self.weights_root_dir, ['swin_base_patch4_window12_384_22kto1k.pth', 'swin_base_patch4_window12_384_22k.pth'][0]), + 'swin_v1_l': os.path.join(self.weights_root_dir, ['swin_large_patch4_window12_384_22kto1k.pth', 'swin_large_patch4_window12_384_22k.pth'][0]), + 'swin_v1_t': os.path.join(self.weights_root_dir, ['swin_tiny_patch4_window7_224_22kto1k_finetune.pth'][0]), + 'swin_v1_s': os.path.join(self.weights_root_dir, ['swin_small_patch4_window7_224_22kto1k_finetune.pth'][0]), + 'pvt_v2_b0': os.path.join(self.weights_root_dir, ['pvt_v2_b0.pth'][0]), + 'pvt_v2_b1': os.path.join(self.weights_root_dir, ['pvt_v2_b1.pth'][0]), + } + + # Callbacks - inactive + self.verbose_eval = True + self.only_S_MAE = False + self.use_fp16 = False # Bugs. It may cause nan in training. + self.SDPA_enabled = False # Bugs. Slower and errors occur in multi-GPUs + + # others + self.device = [0, 'cpu'][0] # .to(0) == .to('cuda:0') + + self.batch_size_valid = 1 + self.rand_seed = 7 + # run_sh_file = [f for f in os.listdir('.') if 'train.sh' == f] + [os.path.join('..', f) for f in os.listdir('..') if 'train.sh' == f] + # with open(run_sh_file[0], 'r') as f: + # lines = f.readlines() + # self.save_last = int([l.strip() for l in lines if '"{}")'.format(self.task) in l and 'val_last=' in l][0].split('val_last=')[-1].split()[0]) + # self.save_step = int([l.strip() for l in lines if '"{}")'.format(self.task) in l and 'step=' in l][0].split('step=')[-1].split()[0]) + # self.val_step = [0, self.save_step][0] + + def print_task(self) -> None: + # Return task for choosing settings in shell scripts. + print(self.task) + + + +### models/backbones/pvt_v2.py + +import torch +import torch.nn as nn +from functools import partial + +from timm.models.layers import DropPath, to_2tuple, trunc_normal_ +from timm.models.registry import register_model + +import math + +# from config import Config + +# config = Config() + +class Mlp(nn.Module): + def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.dwconv = DWConv(hidden_features) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def forward(self, x, H, W): + x = self.fc1(x) + x = self.dwconv(x, H, W) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +class Attention(nn.Module): + def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., sr_ratio=1): + super().__init__() + assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}." + + self.dim = dim + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = qk_scale or head_dim ** -0.5 + + self.q = nn.Linear(dim, dim, bias=qkv_bias) + self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias) + self.attn_drop_prob = attn_drop + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + self.sr_ratio = sr_ratio + if sr_ratio > 1: + self.sr = nn.Conv2d(dim, dim, kernel_size=sr_ratio, stride=sr_ratio) + self.norm = nn.LayerNorm(dim) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def forward(self, x, H, W): + B, N, C = x.shape + q = self.q(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) + + if self.sr_ratio > 1: + x_ = x.permute(0, 2, 1).reshape(B, C, H, W) + x_ = self.sr(x_).reshape(B, C, -1).permute(0, 2, 1) + x_ = self.norm(x_) + kv = self.kv(x_).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + else: + kv = self.kv(x).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + k, v = kv[0], kv[1] + + if config.SDPA_enabled: + x = torch.nn.functional.scaled_dot_product_attention( + q, k, v, + attn_mask=None, dropout_p=self.attn_drop_prob, is_causal=False + ).transpose(1, 2).reshape(B, N, C) + else: + attn = (q @ k.transpose(-2, -1)) * self.scale + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B, N, C) + x = self.proj(x) + x = self.proj_drop(x) + + return x + + +class Block(nn.Module): + + def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., + drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, sr_ratio=1): + super().__init__() + self.norm1 = norm_layer(dim) + self.attn = Attention( + dim, + num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, + attn_drop=attn_drop, proj_drop=drop, sr_ratio=sr_ratio) + # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def forward(self, x, H, W): + x = x + self.drop_path(self.attn(self.norm1(x), H, W)) + x = x + self.drop_path(self.mlp(self.norm2(x), H, W)) + + return x + + +class OverlapPatchEmbed(nn.Module): + """ Image to Patch Embedding + """ + + def __init__(self, img_size=224, patch_size=7, stride=4, in_channels=3, embed_dim=768): + super().__init__() + img_size = to_2tuple(img_size) + patch_size = to_2tuple(patch_size) + + self.img_size = img_size + self.patch_size = patch_size + self.H, self.W = img_size[0] // patch_size[0], img_size[1] // patch_size[1] + self.num_patches = self.H * self.W + self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=stride, + padding=(patch_size[0] // 2, patch_size[1] // 2)) + self.norm = nn.LayerNorm(embed_dim) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def forward(self, x): + x = self.proj(x) + _, _, H, W = x.shape + x = x.flatten(2).transpose(1, 2) + x = self.norm(x) + + return x, H, W + + +class PyramidVisionTransformerImpr(nn.Module): + def __init__(self, img_size=224, patch_size=16, in_channels=3, num_classes=1000, embed_dims=[64, 128, 256, 512], + num_heads=[1, 2, 4, 8], mlp_ratios=[4, 4, 4, 4], qkv_bias=False, qk_scale=None, drop_rate=0., + attn_drop_rate=0., drop_path_rate=0., norm_layer=nn.LayerNorm, + depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1]): + super().__init__() + self.num_classes = num_classes + self.depths = depths + + # patch_embed + self.patch_embed1 = OverlapPatchEmbed(img_size=img_size, patch_size=7, stride=4, in_channels=in_channels, + embed_dim=embed_dims[0]) + self.patch_embed2 = OverlapPatchEmbed(img_size=img_size // 4, patch_size=3, stride=2, in_channels=embed_dims[0], + embed_dim=embed_dims[1]) + self.patch_embed3 = OverlapPatchEmbed(img_size=img_size // 8, patch_size=3, stride=2, in_channels=embed_dims[1], + embed_dim=embed_dims[2]) + self.patch_embed4 = OverlapPatchEmbed(img_size=img_size // 16, patch_size=3, stride=2, in_channels=embed_dims[2], + embed_dim=embed_dims[3]) + + # transformer encoder + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule + cur = 0 + self.block1 = nn.ModuleList([Block( + dim=embed_dims[0], num_heads=num_heads[0], mlp_ratio=mlp_ratios[0], qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, + sr_ratio=sr_ratios[0]) + for i in range(depths[0])]) + self.norm1 = norm_layer(embed_dims[0]) + + cur += depths[0] + self.block2 = nn.ModuleList([Block( + dim=embed_dims[1], num_heads=num_heads[1], mlp_ratio=mlp_ratios[1], qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, + sr_ratio=sr_ratios[1]) + for i in range(depths[1])]) + self.norm2 = norm_layer(embed_dims[1]) + + cur += depths[1] + self.block3 = nn.ModuleList([Block( + dim=embed_dims[2], num_heads=num_heads[2], mlp_ratio=mlp_ratios[2], qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, + sr_ratio=sr_ratios[2]) + for i in range(depths[2])]) + self.norm3 = norm_layer(embed_dims[2]) + + cur += depths[2] + self.block4 = nn.ModuleList([Block( + dim=embed_dims[3], num_heads=num_heads[3], mlp_ratio=mlp_ratios[3], qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, + sr_ratio=sr_ratios[3]) + for i in range(depths[3])]) + self.norm4 = norm_layer(embed_dims[3]) + + # classification head + # self.head = nn.Linear(embed_dims[3], num_classes) if num_classes > 0 else nn.Identity() + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def init_weights(self, pretrained=None): + if isinstance(pretrained, str): + logger = 1 + #load_checkpoint(self, pretrained, map_location='cpu', strict=False, logger=logger) + + def reset_drop_path(self, drop_path_rate): + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(self.depths))] + cur = 0 + for i in range(self.depths[0]): + self.block1[i].drop_path.drop_prob = dpr[cur + i] + + cur += self.depths[0] + for i in range(self.depths[1]): + self.block2[i].drop_path.drop_prob = dpr[cur + i] + + cur += self.depths[1] + for i in range(self.depths[2]): + self.block3[i].drop_path.drop_prob = dpr[cur + i] + + cur += self.depths[2] + for i in range(self.depths[3]): + self.block4[i].drop_path.drop_prob = dpr[cur + i] + + def freeze_patch_emb(self): + self.patch_embed1.requires_grad = False + + @torch.jit.ignore + def no_weight_decay(self): + return {'pos_embed1', 'pos_embed2', 'pos_embed3', 'pos_embed4', 'cls_token'} # has pos_embed may be better + + def get_classifier(self): + return self.head + + def reset_classifier(self, num_classes, global_pool=''): + self.num_classes = num_classes + self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity() + + def forward_features(self, x): + B = x.shape[0] + outs = [] + + # stage 1 + x, H, W = self.patch_embed1(x) + for i, blk in enumerate(self.block1): + x = blk(x, H, W) + x = self.norm1(x) + x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() + outs.append(x) + + # stage 2 + x, H, W = self.patch_embed2(x) + for i, blk in enumerate(self.block2): + x = blk(x, H, W) + x = self.norm2(x) + x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() + outs.append(x) + + # stage 3 + x, H, W = self.patch_embed3(x) + for i, blk in enumerate(self.block3): + x = blk(x, H, W) + x = self.norm3(x) + x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() + outs.append(x) + + # stage 4 + x, H, W = self.patch_embed4(x) + for i, blk in enumerate(self.block4): + x = blk(x, H, W) + x = self.norm4(x) + x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() + outs.append(x) + + return outs + + # return x.mean(dim=1) + + def forward(self, x): + x = self.forward_features(x) + # x = self.head(x) + + return x + + +class DWConv(nn.Module): + def __init__(self, dim=768): + super(DWConv, self).__init__() + self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim) + + def forward(self, x, H, W): + B, N, C = x.shape + x = x.transpose(1, 2).view(B, C, H, W).contiguous() + x = self.dwconv(x) + x = x.flatten(2).transpose(1, 2) + + return x + + +def _conv_filter(state_dict, patch_size=16): + """ convert patch embedding weight from manual patchify + linear proj to conv""" + out_dict = {} + for k, v in state_dict.items(): + if 'patch_embed.proj.weight' in k: + v = v.reshape((v.shape[0], 3, patch_size, patch_size)) + out_dict[k] = v + + return out_dict + + +## @register_model +class pvt_v2_b0(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b0, self).__init__( + patch_size=4, embed_dims=[32, 64, 160, 256], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[2, 2, 2, 2], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + + + +## @register_model +class pvt_v2_b1(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b1, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[2, 2, 2, 2], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + +## @register_model +class pvt_v2_b2(PyramidVisionTransformerImpr): + def __init__(self, in_channels=3, **kwargs): + super(pvt_v2_b2, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1, in_channels=in_channels) + +## @register_model +class pvt_v2_b3(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b3, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 18, 3], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + +## @register_model +class pvt_v2_b4(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b4, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 8, 27, 3], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + + +## @register_model +class pvt_v2_b5(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b5, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[4, 4, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 6, 40, 3], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + + + +### models/backbones/swin_v1.py + +# -------------------------------------------------------- +# Swin Transformer +# Copyright (c) 2021 Microsoft +# Licensed under The MIT License [see LICENSE for details] +# Written by Ze Liu, Yutong Lin, Yixuan Wei +# -------------------------------------------------------- + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as checkpoint +import numpy as np +from timm.models.layers import DropPath, to_2tuple, trunc_normal_ + +# from config import Config + + +# config = Config() + + +class Mlp(nn.Module): + """ Multilayer perceptron.""" + + def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +def window_partition(x, window_size): + """ + Args: + x: (B, H, W, C) + window_size (int): window size + + Returns: + windows: (num_windows*B, window_size, window_size, C) + """ + B, H, W, C = x.shape + x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) + windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) + return windows + + +def window_reverse(windows, window_size, H, W): + """ + Args: + windows: (num_windows*B, window_size, window_size, C) + window_size (int): Window size + H (int): Height of image + W (int): Width of image + + Returns: + x: (B, H, W, C) + """ + B = int(windows.shape[0] / (H * W / window_size / window_size)) + x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1) + x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) + return x + + +class WindowAttention(nn.Module): + """ Window based multi-head self attention (W-MSA) module with relative position bias. + It supports both of shifted and non-shifted window. + + Args: + dim (int): Number of input channels. + window_size (tuple[int]): The height and width of the window. + num_heads (int): Number of attention heads. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set + attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0 + proj_drop (float, optional): Dropout ratio of output. Default: 0.0 + """ + + def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.): + + super().__init__() + self.dim = dim + self.window_size = window_size # Wh, Ww + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = qk_scale or head_dim ** -0.5 + + # define a parameter table of relative position bias + self.relative_position_bias_table = nn.Parameter( + torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)) # 2*Wh-1 * 2*Ww-1, nH + + # get pair-wise relative position index for each token inside the window + coords_h = torch.arange(self.window_size[0]) + coords_w = torch.arange(self.window_size[1]) + coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing='ij')) # 2, Wh, Ww + coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww + relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww + relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2 + relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0 + relative_coords[:, :, 1] += self.window_size[1] - 1 + relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1 + relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww + self.register_buffer("relative_position_index", relative_position_index) + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop_prob = attn_drop + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + trunc_normal_(self.relative_position_bias_table, std=.02) + self.softmax = nn.Softmax(dim=-1) + + def forward(self, x, mask=None): + """ Forward function. + + Args: + x: input features with shape of (num_windows*B, N, C) + mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None + """ + B_, N, C = x.shape + qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple) + + q = q * self.scale + + if config.SDPA_enabled: + x = torch.nn.functional.scaled_dot_product_attention( + q, k, v, + attn_mask=None, dropout_p=self.attn_drop_prob, is_causal=False + ).transpose(1, 2).reshape(B_, N, C) + else: + attn = (q @ k.transpose(-2, -1)) + + relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view( + self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1 + ) # Wh*Ww, Wh*Ww, nH + relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww + attn = attn + relative_position_bias.unsqueeze(0) + + if mask is not None: + nW = mask.shape[0] + attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0) + attn = attn.view(-1, self.num_heads, N, N) + attn = self.softmax(attn) + else: + attn = self.softmax(attn) + + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B_, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class SwinTransformerBlock(nn.Module): + """ Swin Transformer Block. + + Args: + dim (int): Number of input channels. + num_heads (int): Number of attention heads. + window_size (int): Window size. + shift_size (int): Shift size for SW-MSA. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float, optional): Stochastic depth rate. Default: 0.0 + act_layer (nn.Module, optional): Activation layer. Default: nn.GELU + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + """ + + def __init__(self, dim, num_heads, window_size=7, shift_size=0, + mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0., + act_layer=nn.GELU, norm_layer=nn.LayerNorm): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.window_size = window_size + self.shift_size = shift_size + self.mlp_ratio = mlp_ratio + assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size" + + self.norm1 = norm_layer(dim) + self.attn = WindowAttention( + dim, window_size=to_2tuple(self.window_size), num_heads=num_heads, + qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop) + + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + + self.H = None + self.W = None + + def forward(self, x, mask_matrix): + """ Forward function. + + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + mask_matrix: Attention mask for cyclic shift. + """ + B, L, C = x.shape + H, W = self.H, self.W + assert L == H * W, "input feature has wrong size" + + shortcut = x + x = self.norm1(x) + x = x.view(B, H, W, C) + + # pad feature maps to multiples of window size + pad_l = pad_t = 0 + pad_r = (self.window_size - W % self.window_size) % self.window_size + pad_b = (self.window_size - H % self.window_size) % self.window_size + x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b)) + _, Hp, Wp, _ = x.shape + + # cyclic shift + if self.shift_size > 0: + shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)) + attn_mask = mask_matrix + else: + shifted_x = x + attn_mask = None + + # partition windows + x_windows = window_partition(shifted_x, self.window_size) # nW*B, window_size, window_size, C + x_windows = x_windows.view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C + + # W-MSA/SW-MSA + attn_windows = self.attn(x_windows, mask=attn_mask) # nW*B, window_size*window_size, C + + # merge windows + attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C) + shifted_x = window_reverse(attn_windows, self.window_size, Hp, Wp) # B H' W' C + + # reverse cyclic shift + if self.shift_size > 0: + x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2)) + else: + x = shifted_x + + if pad_r > 0 or pad_b > 0: + x = x[:, :H, :W, :].contiguous() + + x = x.view(B, H * W, C) + + # FFN + x = shortcut + self.drop_path(x) + x = x + self.drop_path(self.mlp(self.norm2(x))) + + return x + + +class PatchMerging(nn.Module): + """ Patch Merging Layer + + Args: + dim (int): Number of input channels. + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + """ + def __init__(self, dim, norm_layer=nn.LayerNorm): + super().__init__() + self.dim = dim + self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False) + self.norm = norm_layer(4 * dim) + + def forward(self, x, H, W): + """ Forward function. + + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + """ + B, L, C = x.shape + assert L == H * W, "input feature has wrong size" + + x = x.view(B, H, W, C) + + # padding + pad_input = (H % 2 == 1) or (W % 2 == 1) + if pad_input: + x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2)) + + x0 = x[:, 0::2, 0::2, :] # B H/2 W/2 C + x1 = x[:, 1::2, 0::2, :] # B H/2 W/2 C + x2 = x[:, 0::2, 1::2, :] # B H/2 W/2 C + x3 = x[:, 1::2, 1::2, :] # B H/2 W/2 C + x = torch.cat([x0, x1, x2, x3], -1) # B H/2 W/2 4*C + x = x.view(B, -1, 4 * C) # B H/2*W/2 4*C + + x = self.norm(x) + x = self.reduction(x) + + return x + + +class BasicLayer(nn.Module): + """ A basic Swin Transformer layer for one stage. + + Args: + dim (int): Number of feature channels + depth (int): Depths of this stage. + num_heads (int): Number of attention head. + window_size (int): Local window size. Default: 7. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. + """ + + def __init__(self, + dim, + depth, + num_heads, + window_size=7, + mlp_ratio=4., + qkv_bias=True, + qk_scale=None, + drop=0., + attn_drop=0., + drop_path=0., + norm_layer=nn.LayerNorm, + downsample=None, + use_checkpoint=False): + super().__init__() + self.window_size = window_size + self.shift_size = window_size // 2 + self.depth = depth + self.use_checkpoint = use_checkpoint + + # build blocks + self.blocks = nn.ModuleList([ + SwinTransformerBlock( + dim=dim, + num_heads=num_heads, + window_size=window_size, + shift_size=0 if (i % 2 == 0) else window_size // 2, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop, + attn_drop=attn_drop, + drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path, + norm_layer=norm_layer) + for i in range(depth)]) + + # patch merging layer + if downsample is not None: + self.downsample = downsample(dim=dim, norm_layer=norm_layer) + else: + self.downsample = None + + def forward(self, x, H, W): + """ Forward function. + + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + """ + + # calculate attention mask for SW-MSA + # Turn int to torch.tensor for the compatiability with torch.compile in PyTorch 2.5. + Hp = torch.ceil(torch.tensor(H) / self.window_size).to(torch.int64) * self.window_size + Wp = torch.ceil(torch.tensor(W) / self.window_size).to(torch.int64) * self.window_size + img_mask = torch.zeros((1, Hp, Wp, 1), device=x.device) # 1 Hp Wp 1 + h_slices = (slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None)) + w_slices = (slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None)) + cnt = 0 + for h in h_slices: + for w in w_slices: + img_mask[:, h, w, :] = cnt + cnt += 1 + + mask_windows = window_partition(img_mask, self.window_size) # nW, window_size, window_size, 1 + mask_windows = mask_windows.view(-1, self.window_size * self.window_size) + attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) + attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0)).to(x.dtype) + + for blk in self.blocks: + blk.H, blk.W = H, W + if self.use_checkpoint: + x = checkpoint.checkpoint(blk, x, attn_mask) + else: + x = blk(x, attn_mask) + if self.downsample is not None: + x_down = self.downsample(x, H, W) + Wh, Ww = (H + 1) // 2, (W + 1) // 2 + return x, H, W, x_down, Wh, Ww + else: + return x, H, W, x, H, W + + +class PatchEmbed(nn.Module): + """ Image to Patch Embedding + + Args: + patch_size (int): Patch token size. Default: 4. + in_channels (int): Number of input image channels. Default: 3. + embed_dim (int): Number of linear projection output channels. Default: 96. + norm_layer (nn.Module, optional): Normalization layer. Default: None + """ + + def __init__(self, patch_size=4, in_channels=3, embed_dim=96, norm_layer=None): + super().__init__() + patch_size = to_2tuple(patch_size) + self.patch_size = patch_size + + self.in_channels = in_channels + self.embed_dim = embed_dim + + self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size) + if norm_layer is not None: + self.norm = norm_layer(embed_dim) + else: + self.norm = None + + def forward(self, x): + """Forward function.""" + # padding + _, _, H, W = x.size() + if W % self.patch_size[1] != 0: + x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1])) + if H % self.patch_size[0] != 0: + x = F.pad(x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0])) + + x = self.proj(x) # B C Wh Ww + if self.norm is not None: + Wh, Ww = x.size(2), x.size(3) + x = x.flatten(2).transpose(1, 2) + x = self.norm(x) + x = x.transpose(1, 2).view(-1, self.embed_dim, Wh, Ww) + + return x + + +class SwinTransformer(nn.Module): + """ Swin Transformer backbone. + A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows` - + https://arxiv.org/pdf/2103.14030 + + Args: + pretrain_img_size (int): Input image size for training the pretrained model, + used in absolute postion embedding. Default 224. + patch_size (int | tuple(int)): Patch size. Default: 4. + in_channels (int): Number of input image channels. Default: 3. + embed_dim (int): Number of linear projection output channels. Default: 96. + depths (tuple[int]): Depths of each Swin Transformer stage. + num_heads (tuple[int]): Number of attention head of each stage. + window_size (int): Window size. Default: 7. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4. + qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. + drop_rate (float): Dropout rate. + attn_drop_rate (float): Attention dropout rate. Default: 0. + drop_path_rate (float): Stochastic depth rate. Default: 0.2. + norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm. + ape (bool): If True, add absolute position embedding to the patch embedding. Default: False. + patch_norm (bool): If True, add normalization after patch embedding. Default: True. + out_indices (Sequence[int]): Output from which stages. + frozen_stages (int): Stages to be frozen (stop grad and set eval mode). + -1 means not freezing any parameters. + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. + """ + + def __init__(self, + pretrain_img_size=224, + patch_size=4, + in_channels=3, + embed_dim=96, + depths=[2, 2, 6, 2], + num_heads=[3, 6, 12, 24], + window_size=7, + mlp_ratio=4., + qkv_bias=True, + qk_scale=None, + drop_rate=0., + attn_drop_rate=0., + drop_path_rate=0.2, + norm_layer=nn.LayerNorm, + ape=False, + patch_norm=True, + out_indices=(0, 1, 2, 3), + frozen_stages=-1, + use_checkpoint=False): + super().__init__() + + self.pretrain_img_size = pretrain_img_size + self.num_layers = len(depths) + self.embed_dim = embed_dim + self.ape = ape + self.patch_norm = patch_norm + self.out_indices = out_indices + self.frozen_stages = frozen_stages + + # split image into non-overlapping patches + self.patch_embed = PatchEmbed( + patch_size=patch_size, in_channels=in_channels, embed_dim=embed_dim, + norm_layer=norm_layer if self.patch_norm else None) + + # absolute position embedding + if self.ape: + pretrain_img_size = to_2tuple(pretrain_img_size) + patch_size = to_2tuple(patch_size) + patches_resolution = [pretrain_img_size[0] // patch_size[0], pretrain_img_size[1] // patch_size[1]] + + self.absolute_pos_embed = nn.Parameter(torch.zeros(1, embed_dim, patches_resolution[0], patches_resolution[1])) + trunc_normal_(self.absolute_pos_embed, std=.02) + + self.pos_drop = nn.Dropout(p=drop_rate) + + # stochastic depth + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule + + # build layers + self.layers = nn.ModuleList() + for i_layer in range(self.num_layers): + layer = BasicLayer( + dim=int(embed_dim * 2 ** i_layer), + depth=depths[i_layer], + num_heads=num_heads[i_layer], + window_size=window_size, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])], + norm_layer=norm_layer, + downsample=PatchMerging if (i_layer < self.num_layers - 1) else None, + use_checkpoint=use_checkpoint) + self.layers.append(layer) + + num_features = [int(embed_dim * 2 ** i) for i in range(self.num_layers)] + self.num_features = num_features + + # add a norm layer for each output + for i_layer in out_indices: + layer = norm_layer(num_features[i_layer]) + layer_name = f'norm{i_layer}' + self.add_module(layer_name, layer) + + self._freeze_stages() + + def _freeze_stages(self): + if self.frozen_stages >= 0: + self.patch_embed.eval() + for param in self.patch_embed.parameters(): + param.requires_grad = False + + if self.frozen_stages >= 1 and self.ape: + self.absolute_pos_embed.requires_grad = False + + if self.frozen_stages >= 2: + self.pos_drop.eval() + for i in range(0, self.frozen_stages - 1): + m = self.layers[i] + m.eval() + for param in m.parameters(): + param.requires_grad = False + + + def forward(self, x): + """Forward function.""" + x = self.patch_embed(x) + + Wh, Ww = x.size(2), x.size(3) + if self.ape: + # interpolate the position embedding to the corresponding size + absolute_pos_embed = F.interpolate(self.absolute_pos_embed, size=(Wh, Ww), mode='bicubic') + x = (x + absolute_pos_embed) # B Wh*Ww C + + outs = []#x.contiguous()] + x = x.flatten(2).transpose(1, 2) + x = self.pos_drop(x) + for i in range(self.num_layers): + layer = self.layers[i] + x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww) + + if i in self.out_indices: + norm_layer = getattr(self, f'norm{i}') + x_out = norm_layer(x_out) + + out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous() + outs.append(out) + + return tuple(outs) + + def train(self, mode=True): + """Convert the model into training mode while keep layers freezed.""" + super(SwinTransformer, self).train(mode) + self._freeze_stages() + +def swin_v1_t(): + model = SwinTransformer(embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24], window_size=7) + return model + +def swin_v1_s(): + model = SwinTransformer(embed_dim=96, depths=[2, 2, 18, 2], num_heads=[3, 6, 12, 24], window_size=7) + return model + +def swin_v1_b(): + model = SwinTransformer(embed_dim=128, depths=[2, 2, 18, 2], num_heads=[4, 8, 16, 32], window_size=12) + return model + +def swin_v1_l(): + model = SwinTransformer(embed_dim=192, depths=[2, 2, 18, 2], num_heads=[6, 12, 24, 48], window_size=12) + return model + + + +### models/modules/deform_conv.py + +import torch +import torch.nn as nn +from torchvision.ops import deform_conv2d + + +class DeformableConv2d(nn.Module): + def __init__(self, + in_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1, + bias=False): + + super(DeformableConv2d, self).__init__() + + assert type(kernel_size) == tuple or type(kernel_size) == int + + kernel_size = kernel_size if type(kernel_size) == tuple else (kernel_size, kernel_size) + self.stride = stride if type(stride) == tuple else (stride, stride) + self.padding = padding + + self.offset_conv = nn.Conv2d(in_channels, + 2 * kernel_size[0] * kernel_size[1], + kernel_size=kernel_size, + stride=stride, + padding=self.padding, + bias=True) + + nn.init.constant_(self.offset_conv.weight, 0.) + nn.init.constant_(self.offset_conv.bias, 0.) + + self.modulator_conv = nn.Conv2d(in_channels, + 1 * kernel_size[0] * kernel_size[1], + kernel_size=kernel_size, + stride=stride, + padding=self.padding, + bias=True) + + nn.init.constant_(self.modulator_conv.weight, 0.) + nn.init.constant_(self.modulator_conv.bias, 0.) + + self.regular_conv = nn.Conv2d(in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=self.padding, + bias=bias) + + def forward(self, x): + #h, w = x.shape[2:] + #max_offset = max(h, w)/4. + + offset = self.offset_conv(x)#.clamp(-max_offset, max_offset) + modulator = 2. * torch.sigmoid(self.modulator_conv(x)) + + x = deform_conv2d( + input=x, + offset=offset, + weight=self.regular_conv.weight, + bias=self.regular_conv.bias, + padding=self.padding, + mask=modulator, + stride=self.stride, + ) + return x + + + + +### utils.py + +import torch.nn as nn + + +def build_act_layer(act_layer): + if act_layer == 'ReLU': + return nn.ReLU(inplace=True) + elif act_layer == 'SiLU': + return nn.SiLU(inplace=True) + elif act_layer == 'GELU': + return nn.GELU() + + raise NotImplementedError(f'build_act_layer does not support {act_layer}') + + +def build_norm_layer(dim, + norm_layer, + in_format='channels_last', + out_format='channels_last', + eps=1e-6): + layers = [] + if norm_layer == 'BN': + if in_format == 'channels_last': + layers.append(to_channels_first()) + layers.append(nn.BatchNorm2d(dim)) + if out_format == 'channels_last': + layers.append(to_channels_last()) + elif norm_layer == 'LN': + if in_format == 'channels_first': + layers.append(to_channels_last()) + layers.append(nn.LayerNorm(dim, eps=eps)) + if out_format == 'channels_first': + layers.append(to_channels_first()) + else: + raise NotImplementedError( + f'build_norm_layer does not support {norm_layer}') + return nn.Sequential(*layers) + + +class to_channels_first(nn.Module): + + def __init__(self): + super().__init__() + + def forward(self, x): + return x.permute(0, 3, 1, 2) + + +class to_channels_last(nn.Module): + + def __init__(self): + super().__init__() + + def forward(self, x): + return x.permute(0, 2, 3, 1) + + + +### dataset.py + +_class_labels_TR_sorted = ( + 'Airplane, Ant, Antenna, Archery, Axe, BabyCarriage, Bag, BalanceBeam, Balcony, Balloon, Basket, BasketballHoop, Beatle, Bed, Bee, Bench, Bicycle, ' + 'BicycleFrame, BicycleStand, Boat, Bonsai, BoomLift, Bridge, BunkBed, Butterfly, Button, Cable, CableLift, Cage, Camcorder, Cannon, Canoe, Car, ' + 'CarParkDropArm, Carriage, Cart, Caterpillar, CeilingLamp, Centipede, Chair, Clip, Clock, Clothes, CoatHanger, Comb, ConcretePumpTruck, Crack, Crane, ' + 'Cup, DentalChair, Desk, DeskChair, Diagram, DishRack, DoorHandle, Dragonfish, Dragonfly, Drum, Earphone, Easel, ElectricIron, Excavator, Eyeglasses, ' + 'Fan, Fence, Fencing, FerrisWheel, FireExtinguisher, Fishing, Flag, FloorLamp, Forklift, GasStation, Gate, Gear, Goal, Golf, GymEquipment, Hammock, ' + 'Handcart, Handcraft, Handrail, HangGlider, Harp, Harvester, Headset, Helicopter, Helmet, Hook, HorizontalBar, Hydrovalve, IroningTable, Jewelry, Key, ' + 'KidsPlayground, Kitchenware, Kite, Knife, Ladder, LaundryRack, Lightning, Lobster, Locust, Machine, MachineGun, MagazineRack, Mantis, Medal, MemorialArchway, ' + 'Microphone, Missile, MobileHolder, Monitor, Mosquito, Motorcycle, MovingTrolley, Mower, MusicPlayer, MusicStand, ObservationTower, Octopus, OilWell, ' + 'OlympicLogo, OperatingTable, OutdoorFitnessEquipment, Parachute, Pavilion, Piano, Pipe, PlowHarrow, PoleVault, Punchbag, Rack, Racket, Rifle, Ring, Robot, ' + 'RockClimbing, Rope, Sailboat, Satellite, Scaffold, Scale, Scissor, Scooter, Sculpture, Seadragon, Seahorse, Seal, SewingMachine, Ship, Shoe, ShoppingCart, ' + 'ShoppingTrolley, Shower, Shrimp, Signboard, Skateboarding, Skeleton, Skiing, Spade, SpeedBoat, Spider, Spoon, Stair, Stand, Stationary, SteeringWheel, ' + 'Stethoscope, Stool, Stove, StreetLamp, SweetStand, Swing, Sword, TV, Table, TableChair, TableLamp, TableTennis, Tank, Tapeline, Teapot, Telescope, Tent, ' + 'TobaccoPipe, Toy, Tractor, TrafficLight, TrafficSign, Trampoline, TransmissionTower, Tree, Tricycle, TrimmerCover, Tripod, Trombone, Truck, Trumpet, Tuba, ' + 'UAV, Umbrella, UnevenBars, UtilityPole, VacuumCleaner, Violin, Wakesurfing, Watch, WaterTower, WateringPot, Well, WellLid, Wheel, Wheelchair, WindTurbine, Windmill, WineGlass, WireWhisk, Yacht' +) +class_labels_TR_sorted = _class_labels_TR_sorted.split(', ') + + +### models/backbones/build_backbones.py + +import torch +import torch.nn as nn +from collections import OrderedDict +from torchvision.models import vgg16, vgg16_bn, VGG16_Weights, VGG16_BN_Weights, resnet50, ResNet50_Weights +# from models.pvt_v2 import pvt_v2_b0, pvt_v2_b1, pvt_v2_b2, pvt_v2_b5 +# from models.swin_v1 import swin_v1_t, swin_v1_s, swin_v1_b, swin_v1_l +# from config import Config + + +config = Config() + +def build_backbone(bb_name, pretrained=True, params_settings=''): + if bb_name == 'vgg16': + bb_net = list(vgg16(pretrained=VGG16_Weights.DEFAULT if pretrained else None).children())[0] + bb = nn.Sequential(OrderedDict({'conv1': bb_net[:4], 'conv2': bb_net[4:9], 'conv3': bb_net[9:16], 'conv4': bb_net[16:23]})) + elif bb_name == 'vgg16bn': + bb_net = list(vgg16_bn(pretrained=VGG16_BN_Weights.DEFAULT if pretrained else None).children())[0] + bb = nn.Sequential(OrderedDict({'conv1': bb_net[:6], 'conv2': bb_net[6:13], 'conv3': bb_net[13:23], 'conv4': bb_net[23:33]})) + elif bb_name == 'resnet50': + bb_net = list(resnet50(pretrained=ResNet50_Weights.DEFAULT if pretrained else None).children()) + bb = nn.Sequential(OrderedDict({'conv1': nn.Sequential(*bb_net[0:3]), 'conv2': bb_net[4], 'conv3': bb_net[5], 'conv4': bb_net[6]})) + else: + bb = eval('{}({})'.format(bb_name, params_settings)) + if pretrained: + bb = load_weights(bb, bb_name) + return bb + +def load_weights(model, model_name): + save_model = torch.load(config.weights[model_name], map_location='cpu') + model_dict = model.state_dict() + state_dict = {k: v if v.size() == model_dict[k].size() else model_dict[k] for k, v in save_model.items() if k in model_dict.keys()} + # to ignore the weights with mismatched size when I modify the backbone itself. + if not state_dict: + save_model_keys = list(save_model.keys()) + sub_item = save_model_keys[0] if len(save_model_keys) == 1 else None + state_dict = {k: v if v.size() == model_dict[k].size() else model_dict[k] for k, v in save_model[sub_item].items() if k in model_dict.keys()} + if not state_dict or not sub_item: + print('Weights are not successully loaded. Check the state dict of weights file.') + return None + else: + print('Found correct weights in the "{}" item of loaded state_dict.'.format(sub_item)) + model_dict.update(state_dict) + model.load_state_dict(model_dict) + return model + + + +### models/modules/decoder_blocks.py + +import torch +import torch.nn as nn +# from models.aspp import ASPP, ASPPDeformable +# from config import Config + + +# config = Config() + + +class BasicDecBlk(nn.Module): + def __init__(self, in_channels=64, out_channels=64, inter_channels=64): + super(BasicDecBlk, self).__init__() + inter_channels = in_channels // 4 if config.dec_channels_inter == 'adap' else 64 + self.conv_in = nn.Conv2d(in_channels, inter_channels, 3, 1, padding=1) + self.relu_in = nn.ReLU(inplace=True) + if config.dec_att == 'ASPP': + self.dec_att = ASPP(in_channels=inter_channels) + elif config.dec_att == 'ASPPDeformable': + self.dec_att = ASPPDeformable(in_channels=inter_channels) + self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, padding=1) + self.bn_in = nn.BatchNorm2d(inter_channels) if config.batch_size > 1 else nn.Identity() + self.bn_out = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity() + + def forward(self, x): + x = self.conv_in(x) + x = self.bn_in(x) + x = self.relu_in(x) + if hasattr(self, 'dec_att'): + x = self.dec_att(x) + x = self.conv_out(x) + x = self.bn_out(x) + return x + + +class ResBlk(nn.Module): + def __init__(self, in_channels=64, out_channels=None, inter_channels=64): + super(ResBlk, self).__init__() + if out_channels is None: + out_channels = in_channels + inter_channels = in_channels // 4 if config.dec_channels_inter == 'adap' else 64 + + self.conv_in = nn.Conv2d(in_channels, inter_channels, 3, 1, padding=1) + self.bn_in = nn.BatchNorm2d(inter_channels) if config.batch_size > 1 else nn.Identity() + self.relu_in = nn.ReLU(inplace=True) + + if config.dec_att == 'ASPP': + self.dec_att = ASPP(in_channels=inter_channels) + elif config.dec_att == 'ASPPDeformable': + self.dec_att = ASPPDeformable(in_channels=inter_channels) + + self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, padding=1) + self.bn_out = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity() + + self.conv_resi = nn.Conv2d(in_channels, out_channels, 1, 1, 0) + + def forward(self, x): + _x = self.conv_resi(x) + x = self.conv_in(x) + x = self.bn_in(x) + x = self.relu_in(x) + if hasattr(self, 'dec_att'): + x = self.dec_att(x) + x = self.conv_out(x) + x = self.bn_out(x) + return x + _x + + + +### models/modules/lateral_blocks.py + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from functools import partial + +# from config import Config + + +# config = Config() + + +class BasicLatBlk(nn.Module): + def __init__(self, in_channels=64, out_channels=64, inter_channels=64): + super(BasicLatBlk, self).__init__() + inter_channels = in_channels // 4 if config.dec_channels_inter == 'adap' else 64 + self.conv = nn.Conv2d(in_channels, out_channels, 1, 1, 0) + + def forward(self, x): + x = self.conv(x) + return x + + + +### models/modules/aspp.py + +import torch +import torch.nn as nn +import torch.nn.functional as F +# from models.deform_conv import DeformableConv2d +# from config import Config + + +# config = Config() + + +class _ASPPModule(nn.Module): + def __init__(self, in_channels, planes, kernel_size, padding, dilation): + super(_ASPPModule, self).__init__() + self.atrous_conv = nn.Conv2d(in_channels, planes, kernel_size=kernel_size, + stride=1, padding=padding, dilation=dilation, bias=False) + self.bn = nn.BatchNorm2d(planes) if config.batch_size > 1 else nn.Identity() + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + x = self.atrous_conv(x) + x = self.bn(x) + + return self.relu(x) + + +class ASPP(nn.Module): + def __init__(self, in_channels=64, out_channels=None, output_stride=16): + super(ASPP, self).__init__() + self.down_scale = 1 + if out_channels is None: + out_channels = in_channels + self.in_channelster = 256 // self.down_scale + if output_stride == 16: + dilations = [1, 6, 12, 18] + elif output_stride == 8: + dilations = [1, 12, 24, 36] + else: + raise NotImplementedError + + self.aspp1 = _ASPPModule(in_channels, self.in_channelster, 1, padding=0, dilation=dilations[0]) + self.aspp2 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[1], dilation=dilations[1]) + self.aspp3 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[2], dilation=dilations[2]) + self.aspp4 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[3], dilation=dilations[3]) + + self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)), + nn.Conv2d(in_channels, self.in_channelster, 1, stride=1, bias=False), + nn.BatchNorm2d(self.in_channelster) if config.batch_size > 1 else nn.Identity(), + nn.ReLU(inplace=True)) + self.conv1 = nn.Conv2d(self.in_channelster * 5, out_channels, 1, bias=False) + self.bn1 = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity() + self.relu = nn.ReLU(inplace=True) + self.dropout = nn.Dropout(0.5) + + def forward(self, x): + x1 = self.aspp1(x) + x2 = self.aspp2(x) + x3 = self.aspp3(x) + x4 = self.aspp4(x) + x5 = self.global_avg_pool(x) + x5 = F.interpolate(x5, size=x1.size()[2:], mode='bilinear', align_corners=True) + x = torch.cat((x1, x2, x3, x4, x5), dim=1) + + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + + return self.dropout(x) + + +##################### Deformable +class _ASPPModuleDeformable(nn.Module): + def __init__(self, in_channels, planes, kernel_size, padding): + super(_ASPPModuleDeformable, self).__init__() + self.atrous_conv = DeformableConv2d(in_channels, planes, kernel_size=kernel_size, + stride=1, padding=padding, bias=False) + self.bn = nn.BatchNorm2d(planes) if config.batch_size > 1 else nn.Identity() + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + x = self.atrous_conv(x) + x = self.bn(x) + + return self.relu(x) + + +class ASPPDeformable(nn.Module): + def __init__(self, in_channels, out_channels=None, parallel_block_sizes=[1, 3, 7]): + super(ASPPDeformable, self).__init__() + self.down_scale = 1 + if out_channels is None: + out_channels = in_channels + self.in_channelster = 256 // self.down_scale + + self.aspp1 = _ASPPModuleDeformable(in_channels, self.in_channelster, 1, padding=0) + self.aspp_deforms = nn.ModuleList([ + _ASPPModuleDeformable(in_channels, self.in_channelster, conv_size, padding=int(conv_size//2)) for conv_size in parallel_block_sizes + ]) + + self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)), + nn.Conv2d(in_channels, self.in_channelster, 1, stride=1, bias=False), + nn.BatchNorm2d(self.in_channelster) if config.batch_size > 1 else nn.Identity(), + nn.ReLU(inplace=True)) + self.conv1 = nn.Conv2d(self.in_channelster * (2 + len(self.aspp_deforms)), out_channels, 1, bias=False) + self.bn1 = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity() + self.relu = nn.ReLU(inplace=True) + self.dropout = nn.Dropout(0.5) + + def forward(self, x): + x1 = self.aspp1(x) + x_aspp_deforms = [aspp_deform(x) for aspp_deform in self.aspp_deforms] + x5 = self.global_avg_pool(x) + x5 = F.interpolate(x5, size=x1.size()[2:], mode='bilinear', align_corners=True) + x = torch.cat((x1, *x_aspp_deforms, x5), dim=1) + + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + + return self.dropout(x) + + + +### models/refinement/refiner.py + +import torch +import torch.nn as nn +from collections import OrderedDict +import torch +import torch.nn as nn +import torch.nn.functional as F +from torchvision.models import vgg16, vgg16_bn +from torchvision.models import resnet50 + +# from config import Config +# from dataset import class_labels_TR_sorted +# from models.build_backbone import build_backbone +# from models.decoder_blocks import BasicDecBlk +# from models.lateral_blocks import BasicLatBlk +# from models.ing import * +# from models.stem_layer import StemLayer + + +class RefinerPVTInChannels4(nn.Module): + def __init__(self, in_channels=3+1): + super(RefinerPVTInChannels4, self).__init__() + self.config = Config() + self.epoch = 1 + self.bb = build_backbone(self.config.bb, params_settings='in_channels=4') + + lateral_channels_in_collection = { + 'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64], + 'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64], + 'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192], + } + channels = lateral_channels_in_collection[self.config.bb] + self.squeeze_module = BasicDecBlk(channels[0], channels[0]) + + self.decoder = Decoder(channels) + + if 0: + for key, value in self.named_parameters(): + if 'bb.' in key: + value.requires_grad = False + + def forward(self, x): + if isinstance(x, list): + x = torch.cat(x, dim=1) + ########## Encoder ########## + if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']: + x1 = self.bb.conv1(x) + x2 = self.bb.conv2(x1) + x3 = self.bb.conv3(x2) + x4 = self.bb.conv4(x3) + else: + x1, x2, x3, x4 = self.bb(x) + + x4 = self.squeeze_module(x4) + + ########## Decoder ########## + + features = [x, x1, x2, x3, x4] + scaled_preds = self.decoder(features) + + return scaled_preds + + +class Refiner(nn.Module): + def __init__(self, in_channels=3+1): + super(Refiner, self).__init__() + self.config = Config() + self.epoch = 1 + self.stem_layer = StemLayer(in_channels=in_channels, inter_channels=48, out_channels=3, norm_layer='BN' if self.config.batch_size > 1 else 'LN') + self.bb = build_backbone(self.config.bb) + + lateral_channels_in_collection = { + 'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64], + 'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64], + 'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192], + } + channels = lateral_channels_in_collection[self.config.bb] + self.squeeze_module = BasicDecBlk(channels[0], channels[0]) + + self.decoder = Decoder(channels) + + if 0: + for key, value in self.named_parameters(): + if 'bb.' in key: + value.requires_grad = False + + def forward(self, x): + if isinstance(x, list): + x = torch.cat(x, dim=1) + x = self.stem_layer(x) + ########## Encoder ########## + if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']: + x1 = self.bb.conv1(x) + x2 = self.bb.conv2(x1) + x3 = self.bb.conv3(x2) + x4 = self.bb.conv4(x3) + else: + x1, x2, x3, x4 = self.bb(x) + + x4 = self.squeeze_module(x4) + + ########## Decoder ########## + + features = [x, x1, x2, x3, x4] + scaled_preds = self.decoder(features) + + return scaled_preds + + +class Decoder(nn.Module): + def __init__(self, channels): + super(Decoder, self).__init__() + self.config = Config() + DecoderBlock = eval('BasicDecBlk') + LateralBlock = eval('BasicLatBlk') + + self.decoder_block4 = DecoderBlock(channels[0], channels[1]) + self.decoder_block3 = DecoderBlock(channels[1], channels[2]) + self.decoder_block2 = DecoderBlock(channels[2], channels[3]) + self.decoder_block1 = DecoderBlock(channels[3], channels[3]//2) + + self.lateral_block4 = LateralBlock(channels[1], channels[1]) + self.lateral_block3 = LateralBlock(channels[2], channels[2]) + self.lateral_block2 = LateralBlock(channels[3], channels[3]) + + if self.config.ms_supervision: + self.conv_ms_spvn_4 = nn.Conv2d(channels[1], 1, 1, 1, 0) + self.conv_ms_spvn_3 = nn.Conv2d(channels[2], 1, 1, 1, 0) + self.conv_ms_spvn_2 = nn.Conv2d(channels[3], 1, 1, 1, 0) + self.conv_out1 = nn.Sequential(nn.Conv2d(channels[3]//2, 1, 1, 1, 0)) + + def forward(self, features): + x, x1, x2, x3, x4 = features + outs = [] + p4 = self.decoder_block4(x4) + _p4 = F.interpolate(p4, size=x3.shape[2:], mode='bilinear', align_corners=True) + _p3 = _p4 + self.lateral_block4(x3) + + p3 = self.decoder_block3(_p3) + _p3 = F.interpolate(p3, size=x2.shape[2:], mode='bilinear', align_corners=True) + _p2 = _p3 + self.lateral_block3(x2) + + p2 = self.decoder_block2(_p2) + _p2 = F.interpolate(p2, size=x1.shape[2:], mode='bilinear', align_corners=True) + _p1 = _p2 + self.lateral_block2(x1) + + _p1 = self.decoder_block1(_p1) + _p1 = F.interpolate(_p1, size=x.shape[2:], mode='bilinear', align_corners=True) + p1_out = self.conv_out1(_p1) + + if self.config.ms_supervision: + outs.append(self.conv_ms_spvn_4(p4)) + outs.append(self.conv_ms_spvn_3(p3)) + outs.append(self.conv_ms_spvn_2(p2)) + outs.append(p1_out) + return outs + + +class RefUNet(nn.Module): + # Refinement + def __init__(self, in_channels=3+1): + super(RefUNet, self).__init__() + self.encoder_1 = nn.Sequential( + nn.Conv2d(in_channels, 64, 3, 1, 1), + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.encoder_2 = nn.Sequential( + nn.MaxPool2d(2, 2, ceil_mode=True), + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.encoder_3 = nn.Sequential( + nn.MaxPool2d(2, 2, ceil_mode=True), + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.encoder_4 = nn.Sequential( + nn.MaxPool2d(2, 2, ceil_mode=True), + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.pool4 = nn.MaxPool2d(2, 2, ceil_mode=True) + ##### + self.decoder_5 = nn.Sequential( + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + ##### + self.decoder_4 = nn.Sequential( + nn.Conv2d(128, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.decoder_3 = nn.Sequential( + nn.Conv2d(128, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.decoder_2 = nn.Sequential( + nn.Conv2d(128, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.decoder_1 = nn.Sequential( + nn.Conv2d(128, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.conv_d0 = nn.Conv2d(64, 1, 3, 1, 1) + + self.upscore2 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) + + def forward(self, x): + outs = [] + if isinstance(x, list): + x = torch.cat(x, dim=1) + hx = x + + hx1 = self.encoder_1(hx) + hx2 = self.encoder_2(hx1) + hx3 = self.encoder_3(hx2) + hx4 = self.encoder_4(hx3) + + hx = self.decoder_5(self.pool4(hx4)) + hx = torch.cat((self.upscore2(hx), hx4), 1) + + d4 = self.decoder_4(hx) + hx = torch.cat((self.upscore2(d4), hx3), 1) + + d3 = self.decoder_3(hx) + hx = torch.cat((self.upscore2(d3), hx2), 1) + + d2 = self.decoder_2(hx) + hx = torch.cat((self.upscore2(d2), hx1), 1) + + d1 = self.decoder_1(hx) + + x = self.conv_d0(d1) + outs.append(x) + return outs + + + +### models/stem_layer.py + +import torch.nn as nn +# from utils import build_act_layer, build_norm_layer + + +class StemLayer(nn.Module): + r""" Stem layer of InternImage + Args: + in_channels (int): number of input channels + out_channels (int): number of output channels + act_layer (str): activation layer + norm_layer (str): normalization layer + """ + + def __init__(self, + in_channels=3+1, + inter_channels=48, + out_channels=96, + act_layer='GELU', + norm_layer='BN'): + super().__init__() + self.conv1 = nn.Conv2d(in_channels, + inter_channels, + kernel_size=3, + stride=1, + padding=1) + self.norm1 = build_norm_layer( + inter_channels, norm_layer, 'channels_first', 'channels_first' + ) + self.act = build_act_layer(act_layer) + self.conv2 = nn.Conv2d(inter_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1) + self.norm2 = build_norm_layer( + out_channels, norm_layer, 'channels_first', 'channels_first' + ) + + def forward(self, x): + x = self.conv1(x) + x = self.norm1(x) + x = self.act(x) + x = self.conv2(x) + x = self.norm2(x) + return x + + +### models/birefnet.py + +import torch +import torch.nn as nn +import torch.nn.functional as F +from kornia.filters import laplacian +from transformers import PreTrainedModel +from einops import rearrange + +# from config import Config +# from dataset import class_labels_TR_sorted +# from models.build_backbone import build_backbone +# from models.decoder_blocks import BasicDecBlk, ResBlk, HierarAttDecBlk +# from models.lateral_blocks import BasicLatBlk +# from models.aspp import ASPP, ASPPDeformable +# from models.ing import * +# from models.refiner import Refiner, RefinerPVTInChannels4, RefUNet +# from models.stem_layer import StemLayer +from .BiRefNet_config import BiRefNetConfig + + +def image2patches(image, grid_h=2, grid_w=2, patch_ref=None, transformation='b c (hg h) (wg w) -> (b hg wg) c h w'): + if patch_ref is not None: + grid_h, grid_w = image.shape[-2] // patch_ref.shape[-2], image.shape[-1] // patch_ref.shape[-1] + patches = rearrange(image, transformation, hg=grid_h, wg=grid_w) + return patches + +def patches2image(patches, grid_h=2, grid_w=2, patch_ref=None, transformation='(b hg wg) c h w -> b c (hg h) (wg w)'): + if patch_ref is not None: + grid_h, grid_w = patch_ref.shape[-2] // patches[0].shape[-2], patch_ref.shape[-1] // patches[0].shape[-1] + image = rearrange(patches, transformation, hg=grid_h, wg=grid_w) + return image + +class BiRefNet( + PreTrainedModel +): + config_class = BiRefNetConfig + def __init__(self, bb_pretrained=True, config=BiRefNetConfig()): + super(BiRefNet, self).__init__(config) + bb_pretrained = config.bb_pretrained + self.config = Config() + self.epoch = 1 + self.bb = build_backbone(self.config.bb, pretrained=bb_pretrained) + + channels = self.config.lateral_channels_in_collection + + if self.config.auxiliary_classification: + self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) + self.cls_head = nn.Sequential( + nn.Linear(channels[0], len(class_labels_TR_sorted)) + ) + + if self.config.squeeze_block: + self.squeeze_module = nn.Sequential(*[ + eval(self.config.squeeze_block.split('_x')[0])(channels[0]+sum(self.config.cxt), channels[0]) + for _ in range(eval(self.config.squeeze_block.split('_x')[1])) + ]) + + self.decoder = Decoder(channels) + + if self.config.ender: + self.dec_end = nn.Sequential( + nn.Conv2d(1, 16, 3, 1, 1), + nn.Conv2d(16, 1, 3, 1, 1), + nn.ReLU(inplace=True), + ) + + # refine patch-level segmentation + if self.config.refine: + if self.config.refine == 'itself': + self.stem_layer = StemLayer(in_channels=3+1, inter_channels=48, out_channels=3, norm_layer='BN' if self.config.batch_size > 1 else 'LN') + else: + self.refiner = eval('{}({})'.format(self.config.refine, 'in_channels=3+1')) + + if self.config.freeze_bb: + # Freeze the backbone... + print(self.named_parameters()) + for key, value in self.named_parameters(): + if 'bb.' in key and 'refiner.' not in key: + value.requires_grad = False + + def forward_enc(self, x): + if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']: + x1 = self.bb.conv1(x); x2 = self.bb.conv2(x1); x3 = self.bb.conv3(x2); x4 = self.bb.conv4(x3) + else: + x1, x2, x3, x4 = self.bb(x) + if self.config.mul_scl_ipt == 'cat': + B, C, H, W = x.shape + x1_, x2_, x3_, x4_ = self.bb(F.interpolate(x, size=(H//2, W//2), mode='bilinear', align_corners=True)) + x1 = torch.cat([x1, F.interpolate(x1_, size=x1.shape[2:], mode='bilinear', align_corners=True)], dim=1) + x2 = torch.cat([x2, F.interpolate(x2_, size=x2.shape[2:], mode='bilinear', align_corners=True)], dim=1) + x3 = torch.cat([x3, F.interpolate(x3_, size=x3.shape[2:], mode='bilinear', align_corners=True)], dim=1) + x4 = torch.cat([x4, F.interpolate(x4_, size=x4.shape[2:], mode='bilinear', align_corners=True)], dim=1) + elif self.config.mul_scl_ipt == 'add': + B, C, H, W = x.shape + x1_, x2_, x3_, x4_ = self.bb(F.interpolate(x, size=(H//2, W//2), mode='bilinear', align_corners=True)) + x1 = x1 + F.interpolate(x1_, size=x1.shape[2:], mode='bilinear', align_corners=True) + x2 = x2 + F.interpolate(x2_, size=x2.shape[2:], mode='bilinear', align_corners=True) + x3 = x3 + F.interpolate(x3_, size=x3.shape[2:], mode='bilinear', align_corners=True) + x4 = x4 + F.interpolate(x4_, size=x4.shape[2:], mode='bilinear', align_corners=True) + class_preds = self.cls_head(self.avgpool(x4).view(x4.shape[0], -1)) if self.training and self.config.auxiliary_classification else None + if self.config.cxt: + x4 = torch.cat( + ( + *[ + F.interpolate(x1, size=x4.shape[2:], mode='bilinear', align_corners=True), + F.interpolate(x2, size=x4.shape[2:], mode='bilinear', align_corners=True), + F.interpolate(x3, size=x4.shape[2:], mode='bilinear', align_corners=True), + ][-len(self.config.cxt):], + x4 + ), + dim=1 + ) + return (x1, x2, x3, x4), class_preds + + def forward_ori(self, x): + ########## Encoder ########## + (x1, x2, x3, x4), class_preds = self.forward_enc(x) + if self.config.squeeze_block: + x4 = self.squeeze_module(x4) + ########## Decoder ########## + features = [x, x1, x2, x3, x4] + if self.training and self.config.out_ref: + features.append(laplacian(torch.mean(x, dim=1).unsqueeze(1), kernel_size=5)) + scaled_preds = self.decoder(features) + return scaled_preds, class_preds + + def forward(self, x): + scaled_preds, class_preds = self.forward_ori(x) + class_preds_lst = [class_preds] + return [scaled_preds, class_preds_lst] if self.training else scaled_preds + + +class Decoder(nn.Module): + def __init__(self, channels): + super(Decoder, self).__init__() + self.config = Config() + DecoderBlock = eval(self.config.dec_blk) + LateralBlock = eval(self.config.lat_blk) + + if self.config.dec_ipt: + self.split = self.config.dec_ipt_split + N_dec_ipt = 64 + DBlock = SimpleConvs + ic = 64 + ipt_cha_opt = 1 + self.ipt_blk5 = DBlock(2**10*3 if self.split else 3, [N_dec_ipt, channels[0]//8][ipt_cha_opt], inter_channels=ic) + self.ipt_blk4 = DBlock(2**8*3 if self.split else 3, [N_dec_ipt, channels[0]//8][ipt_cha_opt], inter_channels=ic) + self.ipt_blk3 = DBlock(2**6*3 if self.split else 3, [N_dec_ipt, channels[1]//8][ipt_cha_opt], inter_channels=ic) + self.ipt_blk2 = DBlock(2**4*3 if self.split else 3, [N_dec_ipt, channels[2]//8][ipt_cha_opt], inter_channels=ic) + self.ipt_blk1 = DBlock(2**0*3 if self.split else 3, [N_dec_ipt, channels[3]//8][ipt_cha_opt], inter_channels=ic) + else: + self.split = None + + self.decoder_block4 = DecoderBlock(channels[0]+([N_dec_ipt, channels[0]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[1]) + self.decoder_block3 = DecoderBlock(channels[1]+([N_dec_ipt, channels[0]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[2]) + self.decoder_block2 = DecoderBlock(channels[2]+([N_dec_ipt, channels[1]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[3]) + self.decoder_block1 = DecoderBlock(channels[3]+([N_dec_ipt, channels[2]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[3]//2) + self.conv_out1 = nn.Sequential(nn.Conv2d(channels[3]//2+([N_dec_ipt, channels[3]//8][ipt_cha_opt] if self.config.dec_ipt else 0), 1, 1, 1, 0)) + + self.lateral_block4 = LateralBlock(channels[1], channels[1]) + self.lateral_block3 = LateralBlock(channels[2], channels[2]) + self.lateral_block2 = LateralBlock(channels[3], channels[3]) + + if self.config.ms_supervision: + self.conv_ms_spvn_4 = nn.Conv2d(channels[1], 1, 1, 1, 0) + self.conv_ms_spvn_3 = nn.Conv2d(channels[2], 1, 1, 1, 0) + self.conv_ms_spvn_2 = nn.Conv2d(channels[3], 1, 1, 1, 0) + + if self.config.out_ref: + _N = 16 + self.gdt_convs_4 = nn.Sequential(nn.Conv2d(channels[1], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True)) + self.gdt_convs_3 = nn.Sequential(nn.Conv2d(channels[2], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True)) + self.gdt_convs_2 = nn.Sequential(nn.Conv2d(channels[3], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True)) + + self.gdt_convs_pred_4 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + self.gdt_convs_pred_3 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + self.gdt_convs_pred_2 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + + self.gdt_convs_attn_4 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + self.gdt_convs_attn_3 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + self.gdt_convs_attn_2 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + + def forward(self, features): + if self.training and self.config.out_ref: + outs_gdt_pred = [] + outs_gdt_label = [] + x, x1, x2, x3, x4, gdt_gt = features + else: + x, x1, x2, x3, x4 = features + outs = [] + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=x4, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + x4 = torch.cat((x4, self.ipt_blk5(F.interpolate(patches_batch, size=x4.shape[2:], mode='bilinear', align_corners=True))), 1) + p4 = self.decoder_block4(x4) + m4 = self.conv_ms_spvn_4(p4) if self.config.ms_supervision and self.training else None + if self.config.out_ref: + p4_gdt = self.gdt_convs_4(p4) + if self.training: + # >> GT: + m4_dia = m4 + gdt_label_main_4 = gdt_gt * F.interpolate(m4_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True) + outs_gdt_label.append(gdt_label_main_4) + # >> Pred: + gdt_pred_4 = self.gdt_convs_pred_4(p4_gdt) + outs_gdt_pred.append(gdt_pred_4) + gdt_attn_4 = self.gdt_convs_attn_4(p4_gdt).sigmoid() + # >> Finally: + p4 = p4 * gdt_attn_4 + _p4 = F.interpolate(p4, size=x3.shape[2:], mode='bilinear', align_corners=True) + _p3 = _p4 + self.lateral_block4(x3) + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=_p3, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + _p3 = torch.cat((_p3, self.ipt_blk4(F.interpolate(patches_batch, size=x3.shape[2:], mode='bilinear', align_corners=True))), 1) + p3 = self.decoder_block3(_p3) + m3 = self.conv_ms_spvn_3(p3) if self.config.ms_supervision and self.training else None + if self.config.out_ref: + p3_gdt = self.gdt_convs_3(p3) + if self.training: + # >> GT: + # m3 --dilation--> m3_dia + # G_3^gt * m3_dia --> G_3^m, which is the label of gradient + m3_dia = m3 + gdt_label_main_3 = gdt_gt * F.interpolate(m3_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True) + outs_gdt_label.append(gdt_label_main_3) + # >> Pred: + # p3 --conv--BN--> F_3^G, where F_3^G predicts the \hat{G_3} with xx + # F_3^G --sigmoid--> A_3^G + gdt_pred_3 = self.gdt_convs_pred_3(p3_gdt) + outs_gdt_pred.append(gdt_pred_3) + gdt_attn_3 = self.gdt_convs_attn_3(p3_gdt).sigmoid() + # >> Finally: + # p3 = p3 * A_3^G + p3 = p3 * gdt_attn_3 + _p3 = F.interpolate(p3, size=x2.shape[2:], mode='bilinear', align_corners=True) + _p2 = _p3 + self.lateral_block3(x2) + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=_p2, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + _p2 = torch.cat((_p2, self.ipt_blk3(F.interpolate(patches_batch, size=x2.shape[2:], mode='bilinear', align_corners=True))), 1) + p2 = self.decoder_block2(_p2) + m2 = self.conv_ms_spvn_2(p2) if self.config.ms_supervision and self.training else None + if self.config.out_ref: + p2_gdt = self.gdt_convs_2(p2) + if self.training: + # >> GT: + m2_dia = m2 + gdt_label_main_2 = gdt_gt * F.interpolate(m2_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True) + outs_gdt_label.append(gdt_label_main_2) + # >> Pred: + gdt_pred_2 = self.gdt_convs_pred_2(p2_gdt) + outs_gdt_pred.append(gdt_pred_2) + gdt_attn_2 = self.gdt_convs_attn_2(p2_gdt).sigmoid() + # >> Finally: + p2 = p2 * gdt_attn_2 + _p2 = F.interpolate(p2, size=x1.shape[2:], mode='bilinear', align_corners=True) + _p1 = _p2 + self.lateral_block2(x1) + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=_p1, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + _p1 = torch.cat((_p1, self.ipt_blk2(F.interpolate(patches_batch, size=x1.shape[2:], mode='bilinear', align_corners=True))), 1) + _p1 = self.decoder_block1(_p1) + _p1 = F.interpolate(_p1, size=x.shape[2:], mode='bilinear', align_corners=True) + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=_p1, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + _p1 = torch.cat((_p1, self.ipt_blk1(F.interpolate(patches_batch, size=x.shape[2:], mode='bilinear', align_corners=True))), 1) + p1_out = self.conv_out1(_p1) + + if self.config.ms_supervision and self.training: + outs.append(m4) + outs.append(m3) + outs.append(m2) + outs.append(p1_out) + return outs if not (self.config.out_ref and self.training) else ([outs_gdt_pred, outs_gdt_label], outs) + + +class SimpleConvs(nn.Module): + def __init__( + self, in_channels: int, out_channels: int, inter_channels=64 + ) -> None: + super().__init__() + self.conv1 = nn.Conv2d(in_channels, inter_channels, 3, 1, 1) + self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, 1) + + def forward(self, x): + return self.conv_out(self.conv1(x)) diff --git a/models/RMBG/BiRefNet-HR/config.json b/models/RMBG/BiRefNet-HR/config.json new file mode 100644 index 0000000000000000000000000000000000000000..929f098ac42106c836917bc33c27426df14ee295 --- /dev/null +++ b/models/RMBG/BiRefNet-HR/config.json @@ -0,0 +1,20 @@ +{ + "_name_or_path": "ZhengPeng7/BiRefNet_HR", + "architectures": [ + "BiRefNet" + ], + "auto_map": { + "AutoConfig": "BiRefNet_config.BiRefNetConfig", + "AutoModelForImageSegmentation": "birefnet.BiRefNet" + }, + "custom_pipelines": { + "image-segmentation": { + "pt": [ + "AutoModelForImageSegmentation" + ], + "tf": [], + "type": "image" + } + }, + "bb_pretrained": false +} \ No newline at end of file diff --git a/models/RMBG/BiRefNet-HR/gitattributes b/models/RMBG/BiRefNet-HR/gitattributes new file mode 100644 index 0000000000000000000000000000000000000000..a6344aac8c09253b3b630fb776ae94478aa0275b --- /dev/null +++ b/models/RMBG/BiRefNet-HR/gitattributes @@ -0,0 +1,35 @@ +*.7z filter=lfs diff=lfs merge=lfs -text +*.arrow filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +*.bz2 filter=lfs diff=lfs merge=lfs -text +*.ckpt filter=lfs diff=lfs merge=lfs -text +*.ftz filter=lfs diff=lfs merge=lfs -text +*.gz filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text +*.joblib filter=lfs diff=lfs merge=lfs -text +*.lfs.* filter=lfs diff=lfs merge=lfs -text +*.mlmodel filter=lfs diff=lfs merge=lfs -text +*.model filter=lfs diff=lfs merge=lfs -text +*.msgpack filter=lfs diff=lfs merge=lfs -text +*.npy filter=lfs diff=lfs merge=lfs -text +*.npz filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.ot filter=lfs diff=lfs merge=lfs -text +*.parquet filter=lfs diff=lfs merge=lfs -text +*.pb filter=lfs diff=lfs merge=lfs -text +*.pickle filter=lfs diff=lfs merge=lfs -text +*.pkl filter=lfs diff=lfs merge=lfs -text +*.pt filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.rar filter=lfs diff=lfs merge=lfs -text +*.safetensors filter=lfs diff=lfs merge=lfs -text +saved_model/**/* filter=lfs diff=lfs merge=lfs -text +*.tar.* filter=lfs diff=lfs merge=lfs -text +*.tar filter=lfs diff=lfs merge=lfs -text +*.tflite filter=lfs diff=lfs merge=lfs -text +*.tgz filter=lfs diff=lfs merge=lfs -text +*.wasm filter=lfs diff=lfs merge=lfs -text +*.xz filter=lfs diff=lfs merge=lfs -text +*.zip filter=lfs diff=lfs merge=lfs -text +*.zst filter=lfs diff=lfs merge=lfs -text +*tfevents* filter=lfs diff=lfs merge=lfs -text diff --git a/models/RMBG/BiRefNet-HR/model.safetensors b/models/RMBG/BiRefNet-HR/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..29ac24d6c2be25975739412d7343cfb637152025 --- /dev/null +++ b/models/RMBG/BiRefNet-HR/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9d678bafec0b0019fbb073b7fd02f05ede25dc4b15254f23b2fb0be333200c0d +size 444473596 diff --git a/models/RMBG/BiRefNet/BiRefNet-HR-matting.safetensors b/models/RMBG/BiRefNet/BiRefNet-HR-matting.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..389840ee5d3ed950cbae97eb91d32815a77f5997 --- /dev/null +++ b/models/RMBG/BiRefNet/BiRefNet-HR-matting.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a5a4de698739ea5e0e8bbab28e1b293dde95092b87a442d566cbc585c53cef55 +size 444473596 diff --git a/models/RMBG/BiRefNet/BiRefNet-HR.safetensors b/models/RMBG/BiRefNet/BiRefNet-HR.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..29ac24d6c2be25975739412d7343cfb637152025 --- /dev/null +++ b/models/RMBG/BiRefNet/BiRefNet-HR.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9d678bafec0b0019fbb073b7fd02f05ede25dc4b15254f23b2fb0be333200c0d +size 444473596 diff --git a/models/RMBG/BiRefNet/BiRefNet-general.safetensors b/models/RMBG/BiRefNet/BiRefNet-general.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..f872c610a18fa09e7d2af3ce6e995a40ea91b17c --- /dev/null +++ b/models/RMBG/BiRefNet/BiRefNet-general.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:77277264c0e8c74149d3ff2fade4fd8176965b7108f3c5fc3b8c9c811edb4519 +size 884878856 diff --git a/models/RMBG/BiRefNet/BiRefNet-matting.safetensors b/models/RMBG/BiRefNet/BiRefNet-matting.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..0b751a199714732a4d17a4658728f1399e790314 --- /dev/null +++ b/models/RMBG/BiRefNet/BiRefNet-matting.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a9875de5b1e6c8eb5fdaa8c727a82927ce442cdc87ba3abee6a77e6fa46c25bb +size 884878856 diff --git a/models/RMBG/BiRefNet/BiRefNet-portrait.safetensors b/models/RMBG/BiRefNet/BiRefNet-portrait.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..fa2cb143e1454aa5e4328c07b5f006270f0dff15 --- /dev/null +++ b/models/RMBG/BiRefNet/BiRefNet-portrait.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4a4eb3a5469b75f0cccaec6772c22fc30e6c12ed429c2c0ba71b43e3d8d97182 +size 884878856 diff --git a/models/RMBG/BiRefNet/BiRefNet_512x512.safetensors b/models/RMBG/BiRefNet/BiRefNet_512x512.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..596ed763a1ecb5f043c1405b6c310515f4084c5a --- /dev/null +++ b/models/RMBG/BiRefNet/BiRefNet_512x512.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d94ae0eefb2d2020192001e984ecd6b367478118257a3132e6a484bbf18b0f41 +size 444473596 diff --git a/models/RMBG/BiRefNet/BiRefNet_config.py b/models/RMBG/BiRefNet/BiRefNet_config.py new file mode 100644 index 0000000000000000000000000000000000000000..37c8ac58bec2f52dac34204978a7b61b69e3da76 --- /dev/null +++ b/models/RMBG/BiRefNet/BiRefNet_config.py @@ -0,0 +1,11 @@ +from transformers import PretrainedConfig + +class BiRefNetConfig(PretrainedConfig): + model_type = "SegformerForSemanticSegmentation" + def __init__( + self, + bb_pretrained=False, + **kwargs + ): + self.bb_pretrained = bb_pretrained + super().__init__(**kwargs) diff --git a/models/RMBG/BiRefNet/BiRefNet_lite-2K.safetensors b/models/RMBG/BiRefNet/BiRefNet_lite-2K.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..d80bff5738ab92e6340967e06f71e633412601db --- /dev/null +++ b/models/RMBG/BiRefNet/BiRefNet_lite-2K.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:aa2e4a5af5eb3904694feb40f2b39ec5dd7cd9110906590cfeb982f09a46021d +size 177634392 diff --git a/models/RMBG/BiRefNet/BiRefNet_lite.safetensors b/models/RMBG/BiRefNet/BiRefNet_lite.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..fa51d7521bf93eaf4b1cb30b3ad78797a3bcd850 --- /dev/null +++ b/models/RMBG/BiRefNet/BiRefNet_lite.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4417d89795250e698c3cb0ae8df15743810065f646f48a694fdfa7ca052d0815 +size 177634392 diff --git a/models/RMBG/BiRefNet/birefnet.py b/models/RMBG/BiRefNet/birefnet.py new file mode 100644 index 0000000000000000000000000000000000000000..a6e6383ce92070f84152bd4ec93a7805cde75d6d --- /dev/null +++ b/models/RMBG/BiRefNet/birefnet.py @@ -0,0 +1,2248 @@ +### config.py + +import os +import math + + +class Config(): + def __init__(self) -> None: + # PATH settings + self.sys_home_dir = os.path.expanduser('~') # Make up your file system as: SYS_HOME_DIR/codes/dis/BiRefNet, SYS_HOME_DIR/datasets/dis/xx, SYS_HOME_DIR/weights/xx + + # TASK settings + self.task = ['DIS5K', 'COD', 'HRSOD', 'DIS5K+HRSOD+HRS10K', 'P3M-10k'][0] + self.training_set = { + 'DIS5K': ['DIS-TR', 'DIS-TR+DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4'][0], + 'COD': 'TR-COD10K+TR-CAMO', + 'HRSOD': ['TR-DUTS', 'TR-HRSOD', 'TR-UHRSD', 'TR-DUTS+TR-HRSOD', 'TR-DUTS+TR-UHRSD', 'TR-HRSOD+TR-UHRSD', 'TR-DUTS+TR-HRSOD+TR-UHRSD'][5], + 'DIS5K+HRSOD+HRS10K': 'DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4+DIS-TR+TE-HRS10K+TE-HRSOD+TE-UHRSD+TR-HRS10K+TR-HRSOD+TR-UHRSD', # leave DIS-VD for evaluation. + 'P3M-10k': 'TR-P3M-10k', + }[self.task] + self.prompt4loc = ['dense', 'sparse'][0] + + # Faster-Training settings + self.load_all = True + self.compile = True # 1. Trigger CPU memory leak in some extend, which is an inherent problem of PyTorch. + # Machines with > 70GB CPU memory can run the whole training on DIS5K with default setting. + # 2. Higher PyTorch version may fix it: https://github.com/pytorch/pytorch/issues/119607. + # 3. But compile in Pytorch > 2.0.1 seems to bring no acceleration for training. + self.precisionHigh = True + + # MODEL settings + self.ms_supervision = True + self.out_ref = self.ms_supervision and True + self.dec_ipt = True + self.dec_ipt_split = True + self.cxt_num = [0, 3][1] # multi-scale skip connections from encoder + self.mul_scl_ipt = ['', 'add', 'cat'][2] + self.dec_att = ['', 'ASPP', 'ASPPDeformable'][2] + self.squeeze_block = ['', 'BasicDecBlk_x1', 'ResBlk_x4', 'ASPP_x3', 'ASPPDeformable_x3'][1] + self.dec_blk = ['BasicDecBlk', 'ResBlk', 'HierarAttDecBlk'][0] + + # TRAINING settings + self.batch_size = 4 + self.IoU_finetune_last_epochs = [ + 0, + { + 'DIS5K': -50, + 'COD': -20, + 'HRSOD': -20, + 'DIS5K+HRSOD+HRS10K': -20, + 'P3M-10k': -20, + }[self.task] + ][1] # choose 0 to skip + self.lr = (1e-4 if 'DIS5K' in self.task else 1e-5) * math.sqrt(self.batch_size / 4) # DIS needs high lr to converge faster. Adapt the lr linearly + self.size = 1024 + self.num_workers = max(4, self.batch_size) # will be decrease to min(it, batch_size) at the initialization of the data_loader + + # Backbone settings + self.bb = [ + 'vgg16', 'vgg16bn', 'resnet50', # 0, 1, 2 + 'swin_v1_t', 'swin_v1_s', # 3, 4 + 'swin_v1_b', 'swin_v1_l', # 5-bs9, 6-bs4 + 'pvt_v2_b0', 'pvt_v2_b1', # 7, 8 + 'pvt_v2_b2', 'pvt_v2_b5', # 9-bs10, 10-bs5 + ][6] + self.lateral_channels_in_collection = { + 'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64], + 'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64], + 'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192], + 'swin_v1_t': [768, 384, 192, 96], 'swin_v1_s': [768, 384, 192, 96], + 'pvt_v2_b0': [256, 160, 64, 32], 'pvt_v2_b1': [512, 320, 128, 64], + }[self.bb] + if self.mul_scl_ipt == 'cat': + self.lateral_channels_in_collection = [channel * 2 for channel in self.lateral_channels_in_collection] + self.cxt = self.lateral_channels_in_collection[1:][::-1][-self.cxt_num:] if self.cxt_num else [] + + # MODEL settings - inactive + self.lat_blk = ['BasicLatBlk'][0] + self.dec_channels_inter = ['fixed', 'adap'][0] + self.refine = ['', 'itself', 'RefUNet', 'Refiner', 'RefinerPVTInChannels4'][0] + self.progressive_ref = self.refine and True + self.ender = self.progressive_ref and False + self.scale = self.progressive_ref and 2 + self.auxiliary_classification = False # Only for DIS5K, where class labels are saved in `dataset.py`. + self.refine_iteration = 1 + self.freeze_bb = False + self.model = [ + 'BiRefNet', + ][0] + if self.dec_blk == 'HierarAttDecBlk': + self.batch_size = 2 ** [0, 1, 2, 3, 4][2] + + # TRAINING settings - inactive + self.preproc_methods = ['flip', 'enhance', 'rotate', 'pepper', 'crop'][:4] + self.optimizer = ['Adam', 'AdamW'][1] + self.lr_decay_epochs = [1e5] # Set to negative N to decay the lr in the last N-th epoch. + self.lr_decay_rate = 0.5 + # Loss + self.lambdas_pix_last = { + # not 0 means opening this loss + # original rate -- 1 : 30 : 1.5 : 0.2, bce x 30 + 'bce': 30 * 1, # high performance + 'iou': 0.5 * 1, # 0 / 255 + 'iou_patch': 0.5 * 0, # 0 / 255, win_size = (64, 64) + 'mse': 150 * 0, # can smooth the saliency map + 'triplet': 3 * 0, + 'reg': 100 * 0, + 'ssim': 10 * 1, # help contours, + 'cnt': 5 * 0, # help contours + 'structure': 5 * 0, # structure loss from codes of MVANet. A little improvement on DIS-TE[1,2,3], a bit more decrease on DIS-TE4. + } + self.lambdas_cls = { + 'ce': 5.0 + } + # Adv + self.lambda_adv_g = 10. * 0 # turn to 0 to avoid adv training + self.lambda_adv_d = 3. * (self.lambda_adv_g > 0) + + # PATH settings - inactive + self.data_root_dir = os.path.join(self.sys_home_dir, 'datasets/dis') + self.weights_root_dir = os.path.join(self.sys_home_dir, 'weights') + self.weights = { + 'pvt_v2_b2': os.path.join(self.weights_root_dir, 'pvt_v2_b2.pth'), + 'pvt_v2_b5': os.path.join(self.weights_root_dir, ['pvt_v2_b5.pth', 'pvt_v2_b5_22k.pth'][0]), + 'swin_v1_b': os.path.join(self.weights_root_dir, ['swin_base_patch4_window12_384_22kto1k.pth', 'swin_base_patch4_window12_384_22k.pth'][0]), + 'swin_v1_l': os.path.join(self.weights_root_dir, ['swin_large_patch4_window12_384_22kto1k.pth', 'swin_large_patch4_window12_384_22k.pth'][0]), + 'swin_v1_t': os.path.join(self.weights_root_dir, ['swin_tiny_patch4_window7_224_22kto1k_finetune.pth'][0]), + 'swin_v1_s': os.path.join(self.weights_root_dir, ['swin_small_patch4_window7_224_22kto1k_finetune.pth'][0]), + 'pvt_v2_b0': os.path.join(self.weights_root_dir, ['pvt_v2_b0.pth'][0]), + 'pvt_v2_b1': os.path.join(self.weights_root_dir, ['pvt_v2_b1.pth'][0]), + } + + # Callbacks - inactive + self.verbose_eval = True + self.only_S_MAE = False + self.use_fp16 = False # Bugs. It may cause nan in training. + self.SDPA_enabled = False # Bugs. Slower and errors occur in multi-GPUs + + # others + self.device = [0, 'cpu'][0] # .to(0) == .to('cuda:0') + + self.batch_size_valid = 1 + self.rand_seed = 7 + # run_sh_file = [f for f in os.listdir('.') if 'train.sh' == f] + [os.path.join('..', f) for f in os.listdir('..') if 'train.sh' == f] + # with open(run_sh_file[0], 'r') as f: + # lines = f.readlines() + # self.save_last = int([l.strip() for l in lines if '"{}")'.format(self.task) in l and 'val_last=' in l][0].split('val_last=')[-1].split()[0]) + # self.save_step = int([l.strip() for l in lines if '"{}")'.format(self.task) in l and 'step=' in l][0].split('step=')[-1].split()[0]) + # self.val_step = [0, self.save_step][0] + + def print_task(self) -> None: + # Return task for choosing settings in shell scripts. + print(self.task) + + + +### models/backbones/pvt_v2.py + +import torch +import torch.nn as nn +from functools import partial + +from timm.models.layers import DropPath, to_2tuple, trunc_normal_ +from timm.models.registry import register_model + +import math + +# from config import Config + +# config = Config() + +class Mlp(nn.Module): + def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.dwconv = DWConv(hidden_features) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def forward(self, x, H, W): + x = self.fc1(x) + x = self.dwconv(x, H, W) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +class Attention(nn.Module): + def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., sr_ratio=1): + super().__init__() + assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}." + + self.dim = dim + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = qk_scale or head_dim ** -0.5 + + self.q = nn.Linear(dim, dim, bias=qkv_bias) + self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias) + self.attn_drop_prob = attn_drop + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + self.sr_ratio = sr_ratio + if sr_ratio > 1: + self.sr = nn.Conv2d(dim, dim, kernel_size=sr_ratio, stride=sr_ratio) + self.norm = nn.LayerNorm(dim) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def forward(self, x, H, W): + B, N, C = x.shape + q = self.q(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) + + if self.sr_ratio > 1: + x_ = x.permute(0, 2, 1).reshape(B, C, H, W) + x_ = self.sr(x_).reshape(B, C, -1).permute(0, 2, 1) + x_ = self.norm(x_) + kv = self.kv(x_).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + else: + kv = self.kv(x).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + k, v = kv[0], kv[1] + + if config.SDPA_enabled: + x = torch.nn.functional.scaled_dot_product_attention( + q, k, v, + attn_mask=None, dropout_p=self.attn_drop_prob, is_causal=False + ).transpose(1, 2).reshape(B, N, C) + else: + attn = (q @ k.transpose(-2, -1)) * self.scale + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B, N, C) + x = self.proj(x) + x = self.proj_drop(x) + + return x + + +class Block(nn.Module): + + def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., + drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, sr_ratio=1): + super().__init__() + self.norm1 = norm_layer(dim) + self.attn = Attention( + dim, + num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, + attn_drop=attn_drop, proj_drop=drop, sr_ratio=sr_ratio) + # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def forward(self, x, H, W): + x = x + self.drop_path(self.attn(self.norm1(x), H, W)) + x = x + self.drop_path(self.mlp(self.norm2(x), H, W)) + + return x + + +class OverlapPatchEmbed(nn.Module): + """ Image to Patch Embedding + """ + + def __init__(self, img_size=224, patch_size=7, stride=4, in_channels=3, embed_dim=768): + super().__init__() + img_size = to_2tuple(img_size) + patch_size = to_2tuple(patch_size) + + self.img_size = img_size + self.patch_size = patch_size + self.H, self.W = img_size[0] // patch_size[0], img_size[1] // patch_size[1] + self.num_patches = self.H * self.W + self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=stride, + padding=(patch_size[0] // 2, patch_size[1] // 2)) + self.norm = nn.LayerNorm(embed_dim) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def forward(self, x): + x = self.proj(x) + _, _, H, W = x.shape + x = x.flatten(2).transpose(1, 2) + x = self.norm(x) + + return x, H, W + + +class PyramidVisionTransformerImpr(nn.Module): + def __init__(self, img_size=224, patch_size=16, in_channels=3, num_classes=1000, embed_dims=[64, 128, 256, 512], + num_heads=[1, 2, 4, 8], mlp_ratios=[4, 4, 4, 4], qkv_bias=False, qk_scale=None, drop_rate=0., + attn_drop_rate=0., drop_path_rate=0., norm_layer=nn.LayerNorm, + depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1]): + super().__init__() + self.num_classes = num_classes + self.depths = depths + + # patch_embed + self.patch_embed1 = OverlapPatchEmbed(img_size=img_size, patch_size=7, stride=4, in_channels=in_channels, + embed_dim=embed_dims[0]) + self.patch_embed2 = OverlapPatchEmbed(img_size=img_size // 4, patch_size=3, stride=2, in_channels=embed_dims[0], + embed_dim=embed_dims[1]) + self.patch_embed3 = OverlapPatchEmbed(img_size=img_size // 8, patch_size=3, stride=2, in_channels=embed_dims[1], + embed_dim=embed_dims[2]) + self.patch_embed4 = OverlapPatchEmbed(img_size=img_size // 16, patch_size=3, stride=2, in_channels=embed_dims[2], + embed_dim=embed_dims[3]) + + # transformer encoder + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule + cur = 0 + self.block1 = nn.ModuleList([Block( + dim=embed_dims[0], num_heads=num_heads[0], mlp_ratio=mlp_ratios[0], qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, + sr_ratio=sr_ratios[0]) + for i in range(depths[0])]) + self.norm1 = norm_layer(embed_dims[0]) + + cur += depths[0] + self.block2 = nn.ModuleList([Block( + dim=embed_dims[1], num_heads=num_heads[1], mlp_ratio=mlp_ratios[1], qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, + sr_ratio=sr_ratios[1]) + for i in range(depths[1])]) + self.norm2 = norm_layer(embed_dims[1]) + + cur += depths[1] + self.block3 = nn.ModuleList([Block( + dim=embed_dims[2], num_heads=num_heads[2], mlp_ratio=mlp_ratios[2], qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, + sr_ratio=sr_ratios[2]) + for i in range(depths[2])]) + self.norm3 = norm_layer(embed_dims[2]) + + cur += depths[2] + self.block4 = nn.ModuleList([Block( + dim=embed_dims[3], num_heads=num_heads[3], mlp_ratio=mlp_ratios[3], qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, + sr_ratio=sr_ratios[3]) + for i in range(depths[3])]) + self.norm4 = norm_layer(embed_dims[3]) + + # classification head + # self.head = nn.Linear(embed_dims[3], num_classes) if num_classes > 0 else nn.Identity() + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def init_weights(self, pretrained=None): + if isinstance(pretrained, str): + logger = 1 + #load_checkpoint(self, pretrained, map_location='cpu', strict=False, logger=logger) + + def reset_drop_path(self, drop_path_rate): + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(self.depths))] + cur = 0 + for i in range(self.depths[0]): + self.block1[i].drop_path.drop_prob = dpr[cur + i] + + cur += self.depths[0] + for i in range(self.depths[1]): + self.block2[i].drop_path.drop_prob = dpr[cur + i] + + cur += self.depths[1] + for i in range(self.depths[2]): + self.block3[i].drop_path.drop_prob = dpr[cur + i] + + cur += self.depths[2] + for i in range(self.depths[3]): + self.block4[i].drop_path.drop_prob = dpr[cur + i] + + def freeze_patch_emb(self): + self.patch_embed1.requires_grad = False + + @torch.jit.ignore + def no_weight_decay(self): + return {'pos_embed1', 'pos_embed2', 'pos_embed3', 'pos_embed4', 'cls_token'} # has pos_embed may be better + + def get_classifier(self): + return self.head + + def reset_classifier(self, num_classes, global_pool=''): + self.num_classes = num_classes + self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity() + + def forward_features(self, x): + B = x.shape[0] + outs = [] + + # stage 1 + x, H, W = self.patch_embed1(x) + for i, blk in enumerate(self.block1): + x = blk(x, H, W) + x = self.norm1(x) + x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() + outs.append(x) + + # stage 2 + x, H, W = self.patch_embed2(x) + for i, blk in enumerate(self.block2): + x = blk(x, H, W) + x = self.norm2(x) + x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() + outs.append(x) + + # stage 3 + x, H, W = self.patch_embed3(x) + for i, blk in enumerate(self.block3): + x = blk(x, H, W) + x = self.norm3(x) + x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() + outs.append(x) + + # stage 4 + x, H, W = self.patch_embed4(x) + for i, blk in enumerate(self.block4): + x = blk(x, H, W) + x = self.norm4(x) + x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() + outs.append(x) + + return outs + + # return x.mean(dim=1) + + def forward(self, x): + x = self.forward_features(x) + # x = self.head(x) + + return x + + +class DWConv(nn.Module): + def __init__(self, dim=768): + super(DWConv, self).__init__() + self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim) + + def forward(self, x, H, W): + B, N, C = x.shape + x = x.transpose(1, 2).view(B, C, H, W).contiguous() + x = self.dwconv(x) + x = x.flatten(2).transpose(1, 2) + + return x + + +def _conv_filter(state_dict, patch_size=16): + """ convert patch embedding weight from manual patchify + linear proj to conv""" + out_dict = {} + for k, v in state_dict.items(): + if 'patch_embed.proj.weight' in k: + v = v.reshape((v.shape[0], 3, patch_size, patch_size)) + out_dict[k] = v + + return out_dict + + +## @register_model +class pvt_v2_b0(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b0, self).__init__( + patch_size=4, embed_dims=[32, 64, 160, 256], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[2, 2, 2, 2], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + + + +## @register_model +class pvt_v2_b1(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b1, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[2, 2, 2, 2], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + +## @register_model +class pvt_v2_b2(PyramidVisionTransformerImpr): + def __init__(self, in_channels=3, **kwargs): + super(pvt_v2_b2, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1, in_channels=in_channels) + +## @register_model +class pvt_v2_b3(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b3, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 18, 3], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + +## @register_model +class pvt_v2_b4(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b4, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 8, 27, 3], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + + +## @register_model +class pvt_v2_b5(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b5, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[4, 4, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 6, 40, 3], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + + + +### models/backbones/swin_v1.py + +# -------------------------------------------------------- +# Swin Transformer +# Copyright (c) 2021 Microsoft +# Licensed under The MIT License [see LICENSE for details] +# Written by Ze Liu, Yutong Lin, Yixuan Wei +# -------------------------------------------------------- + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as checkpoint +import numpy as np +from timm.models.layers import DropPath, to_2tuple, trunc_normal_ + +# from config import Config + + +# config = Config() + + +class Mlp(nn.Module): + """ Multilayer perceptron.""" + + def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +def window_partition(x, window_size): + """ + Args: + x: (B, H, W, C) + window_size (int): window size + + Returns: + windows: (num_windows*B, window_size, window_size, C) + """ + B, H, W, C = x.shape + x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) + windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) + return windows + + +def window_reverse(windows, window_size, H, W): + """ + Args: + windows: (num_windows*B, window_size, window_size, C) + window_size (int): Window size + H (int): Height of image + W (int): Width of image + + Returns: + x: (B, H, W, C) + """ + B = int(windows.shape[0] / (H * W / window_size / window_size)) + x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1) + x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) + return x + + +class WindowAttention(nn.Module): + """ Window based multi-head self attention (W-MSA) module with relative position bias. + It supports both of shifted and non-shifted window. + + Args: + dim (int): Number of input channels. + window_size (tuple[int]): The height and width of the window. + num_heads (int): Number of attention heads. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set + attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0 + proj_drop (float, optional): Dropout ratio of output. Default: 0.0 + """ + + def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.): + + super().__init__() + self.dim = dim + self.window_size = window_size # Wh, Ww + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = qk_scale or head_dim ** -0.5 + + # define a parameter table of relative position bias + self.relative_position_bias_table = nn.Parameter( + torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)) # 2*Wh-1 * 2*Ww-1, nH + + # get pair-wise relative position index for each token inside the window + coords_h = torch.arange(self.window_size[0]) + coords_w = torch.arange(self.window_size[1]) + coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing='ij')) # 2, Wh, Ww + coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww + relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww + relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2 + relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0 + relative_coords[:, :, 1] += self.window_size[1] - 1 + relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1 + relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww + self.register_buffer("relative_position_index", relative_position_index) + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop_prob = attn_drop + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + trunc_normal_(self.relative_position_bias_table, std=.02) + self.softmax = nn.Softmax(dim=-1) + + def forward(self, x, mask=None): + """ Forward function. + + Args: + x: input features with shape of (num_windows*B, N, C) + mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None + """ + B_, N, C = x.shape + qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple) + + q = q * self.scale + + if config.SDPA_enabled: + x = torch.nn.functional.scaled_dot_product_attention( + q, k, v, + attn_mask=None, dropout_p=self.attn_drop_prob, is_causal=False + ).transpose(1, 2).reshape(B_, N, C) + else: + attn = (q @ k.transpose(-2, -1)) + + relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view( + self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1 + ) # Wh*Ww, Wh*Ww, nH + relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww + attn = attn + relative_position_bias.unsqueeze(0) + + if mask is not None: + nW = mask.shape[0] + attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0) + attn = attn.view(-1, self.num_heads, N, N) + attn = self.softmax(attn) + else: + attn = self.softmax(attn) + + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B_, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class SwinTransformerBlock(nn.Module): + """ Swin Transformer Block. + + Args: + dim (int): Number of input channels. + num_heads (int): Number of attention heads. + window_size (int): Window size. + shift_size (int): Shift size for SW-MSA. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float, optional): Stochastic depth rate. Default: 0.0 + act_layer (nn.Module, optional): Activation layer. Default: nn.GELU + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + """ + + def __init__(self, dim, num_heads, window_size=7, shift_size=0, + mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0., + act_layer=nn.GELU, norm_layer=nn.LayerNorm): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.window_size = window_size + self.shift_size = shift_size + self.mlp_ratio = mlp_ratio + assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size" + + self.norm1 = norm_layer(dim) + self.attn = WindowAttention( + dim, window_size=to_2tuple(self.window_size), num_heads=num_heads, + qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop) + + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + + self.H = None + self.W = None + + def forward(self, x, mask_matrix): + """ Forward function. + + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + mask_matrix: Attention mask for cyclic shift. + """ + B, L, C = x.shape + H, W = self.H, self.W + assert L == H * W, "input feature has wrong size" + + shortcut = x + x = self.norm1(x) + x = x.view(B, H, W, C) + + # pad feature maps to multiples of window size + pad_l = pad_t = 0 + pad_r = (self.window_size - W % self.window_size) % self.window_size + pad_b = (self.window_size - H % self.window_size) % self.window_size + x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b)) + _, Hp, Wp, _ = x.shape + + # cyclic shift + if self.shift_size > 0: + shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)) + attn_mask = mask_matrix + else: + shifted_x = x + attn_mask = None + + # partition windows + x_windows = window_partition(shifted_x, self.window_size) # nW*B, window_size, window_size, C + x_windows = x_windows.view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C + + # W-MSA/SW-MSA + attn_windows = self.attn(x_windows, mask=attn_mask) # nW*B, window_size*window_size, C + + # merge windows + attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C) + shifted_x = window_reverse(attn_windows, self.window_size, Hp, Wp) # B H' W' C + + # reverse cyclic shift + if self.shift_size > 0: + x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2)) + else: + x = shifted_x + + if pad_r > 0 or pad_b > 0: + x = x[:, :H, :W, :].contiguous() + + x = x.view(B, H * W, C) + + # FFN + x = shortcut + self.drop_path(x) + x = x + self.drop_path(self.mlp(self.norm2(x))) + + return x + + +class PatchMerging(nn.Module): + """ Patch Merging Layer + + Args: + dim (int): Number of input channels. + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + """ + def __init__(self, dim, norm_layer=nn.LayerNorm): + super().__init__() + self.dim = dim + self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False) + self.norm = norm_layer(4 * dim) + + def forward(self, x, H, W): + """ Forward function. + + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + """ + B, L, C = x.shape + assert L == H * W, "input feature has wrong size" + + x = x.view(B, H, W, C) + + # padding + pad_input = (H % 2 == 1) or (W % 2 == 1) + if pad_input: + x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2)) + + x0 = x[:, 0::2, 0::2, :] # B H/2 W/2 C + x1 = x[:, 1::2, 0::2, :] # B H/2 W/2 C + x2 = x[:, 0::2, 1::2, :] # B H/2 W/2 C + x3 = x[:, 1::2, 1::2, :] # B H/2 W/2 C + x = torch.cat([x0, x1, x2, x3], -1) # B H/2 W/2 4*C + x = x.view(B, -1, 4 * C) # B H/2*W/2 4*C + + x = self.norm(x) + x = self.reduction(x) + + return x + + +class BasicLayer(nn.Module): + """ A basic Swin Transformer layer for one stage. + + Args: + dim (int): Number of feature channels + depth (int): Depths of this stage. + num_heads (int): Number of attention head. + window_size (int): Local window size. Default: 7. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. + """ + + def __init__(self, + dim, + depth, + num_heads, + window_size=7, + mlp_ratio=4., + qkv_bias=True, + qk_scale=None, + drop=0., + attn_drop=0., + drop_path=0., + norm_layer=nn.LayerNorm, + downsample=None, + use_checkpoint=False): + super().__init__() + self.window_size = window_size + self.shift_size = window_size // 2 + self.depth = depth + self.use_checkpoint = use_checkpoint + + # build blocks + self.blocks = nn.ModuleList([ + SwinTransformerBlock( + dim=dim, + num_heads=num_heads, + window_size=window_size, + shift_size=0 if (i % 2 == 0) else window_size // 2, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop, + attn_drop=attn_drop, + drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path, + norm_layer=norm_layer) + for i in range(depth)]) + + # patch merging layer + if downsample is not None: + self.downsample = downsample(dim=dim, norm_layer=norm_layer) + else: + self.downsample = None + + def forward(self, x, H, W): + """ Forward function. + + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + """ + + # calculate attention mask for SW-MSA + # Turn int to torch.tensor for the compatiability with torch.compile in PyTorch 2.5. + Hp = torch.ceil(torch.tensor(H) / self.window_size).to(torch.int64) * self.window_size + Wp = torch.ceil(torch.tensor(W) / self.window_size).to(torch.int64) * self.window_size + img_mask = torch.zeros((1, Hp, Wp, 1), device=x.device) # 1 Hp Wp 1 + h_slices = (slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None)) + w_slices = (slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None)) + cnt = 0 + for h in h_slices: + for w in w_slices: + img_mask[:, h, w, :] = cnt + cnt += 1 + + mask_windows = window_partition(img_mask, self.window_size) # nW, window_size, window_size, 1 + mask_windows = mask_windows.view(-1, self.window_size * self.window_size) + attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) + attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0)).to(x.dtype) + + for blk in self.blocks: + blk.H, blk.W = H, W + if self.use_checkpoint: + x = checkpoint.checkpoint(blk, x, attn_mask) + else: + x = blk(x, attn_mask) + if self.downsample is not None: + x_down = self.downsample(x, H, W) + Wh, Ww = (H + 1) // 2, (W + 1) // 2 + return x, H, W, x_down, Wh, Ww + else: + return x, H, W, x, H, W + + +class PatchEmbed(nn.Module): + """ Image to Patch Embedding + + Args: + patch_size (int): Patch token size. Default: 4. + in_channels (int): Number of input image channels. Default: 3. + embed_dim (int): Number of linear projection output channels. Default: 96. + norm_layer (nn.Module, optional): Normalization layer. Default: None + """ + + def __init__(self, patch_size=4, in_channels=3, embed_dim=96, norm_layer=None): + super().__init__() + patch_size = to_2tuple(patch_size) + self.patch_size = patch_size + + self.in_channels = in_channels + self.embed_dim = embed_dim + + self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size) + if norm_layer is not None: + self.norm = norm_layer(embed_dim) + else: + self.norm = None + + def forward(self, x): + """Forward function.""" + # padding + _, _, H, W = x.size() + if W % self.patch_size[1] != 0: + x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1])) + if H % self.patch_size[0] != 0: + x = F.pad(x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0])) + + x = self.proj(x) # B C Wh Ww + if self.norm is not None: + Wh, Ww = x.size(2), x.size(3) + x = x.flatten(2).transpose(1, 2) + x = self.norm(x) + x = x.transpose(1, 2).view(-1, self.embed_dim, Wh, Ww) + + return x + + +class SwinTransformer(nn.Module): + """ Swin Transformer backbone. + A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows` - + https://arxiv.org/pdf/2103.14030 + + Args: + pretrain_img_size (int): Input image size for training the pretrained model, + used in absolute postion embedding. Default 224. + patch_size (int | tuple(int)): Patch size. Default: 4. + in_channels (int): Number of input image channels. Default: 3. + embed_dim (int): Number of linear projection output channels. Default: 96. + depths (tuple[int]): Depths of each Swin Transformer stage. + num_heads (tuple[int]): Number of attention head of each stage. + window_size (int): Window size. Default: 7. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4. + qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. + drop_rate (float): Dropout rate. + attn_drop_rate (float): Attention dropout rate. Default: 0. + drop_path_rate (float): Stochastic depth rate. Default: 0.2. + norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm. + ape (bool): If True, add absolute position embedding to the patch embedding. Default: False. + patch_norm (bool): If True, add normalization after patch embedding. Default: True. + out_indices (Sequence[int]): Output from which stages. + frozen_stages (int): Stages to be frozen (stop grad and set eval mode). + -1 means not freezing any parameters. + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. + """ + + def __init__(self, + pretrain_img_size=224, + patch_size=4, + in_channels=3, + embed_dim=96, + depths=[2, 2, 6, 2], + num_heads=[3, 6, 12, 24], + window_size=7, + mlp_ratio=4., + qkv_bias=True, + qk_scale=None, + drop_rate=0., + attn_drop_rate=0., + drop_path_rate=0.2, + norm_layer=nn.LayerNorm, + ape=False, + patch_norm=True, + out_indices=(0, 1, 2, 3), + frozen_stages=-1, + use_checkpoint=False): + super().__init__() + + self.pretrain_img_size = pretrain_img_size + self.num_layers = len(depths) + self.embed_dim = embed_dim + self.ape = ape + self.patch_norm = patch_norm + self.out_indices = out_indices + self.frozen_stages = frozen_stages + + # split image into non-overlapping patches + self.patch_embed = PatchEmbed( + patch_size=patch_size, in_channels=in_channels, embed_dim=embed_dim, + norm_layer=norm_layer if self.patch_norm else None) + + # absolute position embedding + if self.ape: + pretrain_img_size = to_2tuple(pretrain_img_size) + patch_size = to_2tuple(patch_size) + patches_resolution = [pretrain_img_size[0] // patch_size[0], pretrain_img_size[1] // patch_size[1]] + + self.absolute_pos_embed = nn.Parameter(torch.zeros(1, embed_dim, patches_resolution[0], patches_resolution[1])) + trunc_normal_(self.absolute_pos_embed, std=.02) + + self.pos_drop = nn.Dropout(p=drop_rate) + + # stochastic depth + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule + + # build layers + self.layers = nn.ModuleList() + for i_layer in range(self.num_layers): + layer = BasicLayer( + dim=int(embed_dim * 2 ** i_layer), + depth=depths[i_layer], + num_heads=num_heads[i_layer], + window_size=window_size, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])], + norm_layer=norm_layer, + downsample=PatchMerging if (i_layer < self.num_layers - 1) else None, + use_checkpoint=use_checkpoint) + self.layers.append(layer) + + num_features = [int(embed_dim * 2 ** i) for i in range(self.num_layers)] + self.num_features = num_features + + # add a norm layer for each output + for i_layer in out_indices: + layer = norm_layer(num_features[i_layer]) + layer_name = f'norm{i_layer}' + self.add_module(layer_name, layer) + + self._freeze_stages() + + def _freeze_stages(self): + if self.frozen_stages >= 0: + self.patch_embed.eval() + for param in self.patch_embed.parameters(): + param.requires_grad = False + + if self.frozen_stages >= 1 and self.ape: + self.absolute_pos_embed.requires_grad = False + + if self.frozen_stages >= 2: + self.pos_drop.eval() + for i in range(0, self.frozen_stages - 1): + m = self.layers[i] + m.eval() + for param in m.parameters(): + param.requires_grad = False + + + def forward(self, x): + """Forward function.""" + x = self.patch_embed(x) + + Wh, Ww = x.size(2), x.size(3) + if self.ape: + # interpolate the position embedding to the corresponding size + absolute_pos_embed = F.interpolate(self.absolute_pos_embed, size=(Wh, Ww), mode='bicubic') + x = (x + absolute_pos_embed) # B Wh*Ww C + + outs = []#x.contiguous()] + x = x.flatten(2).transpose(1, 2) + x = self.pos_drop(x) + for i in range(self.num_layers): + layer = self.layers[i] + x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww) + + if i in self.out_indices: + norm_layer = getattr(self, f'norm{i}') + x_out = norm_layer(x_out) + + out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous() + outs.append(out) + + return tuple(outs) + + def train(self, mode=True): + """Convert the model into training mode while keep layers freezed.""" + super(SwinTransformer, self).train(mode) + self._freeze_stages() + +def swin_v1_t(): + model = SwinTransformer(embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24], window_size=7) + return model + +def swin_v1_s(): + model = SwinTransformer(embed_dim=96, depths=[2, 2, 18, 2], num_heads=[3, 6, 12, 24], window_size=7) + return model + +def swin_v1_b(): + model = SwinTransformer(embed_dim=128, depths=[2, 2, 18, 2], num_heads=[4, 8, 16, 32], window_size=12) + return model + +def swin_v1_l(): + model = SwinTransformer(embed_dim=192, depths=[2, 2, 18, 2], num_heads=[6, 12, 24, 48], window_size=12) + return model + + + +### models/modules/deform_conv.py + +import torch +import torch.nn as nn +from torchvision.ops import deform_conv2d + + +class DeformableConv2d(nn.Module): + def __init__(self, + in_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1, + bias=False): + + super(DeformableConv2d, self).__init__() + + assert type(kernel_size) == tuple or type(kernel_size) == int + + kernel_size = kernel_size if type(kernel_size) == tuple else (kernel_size, kernel_size) + self.stride = stride if type(stride) == tuple else (stride, stride) + self.padding = padding + + self.offset_conv = nn.Conv2d(in_channels, + 2 * kernel_size[0] * kernel_size[1], + kernel_size=kernel_size, + stride=stride, + padding=self.padding, + bias=True) + + nn.init.constant_(self.offset_conv.weight, 0.) + nn.init.constant_(self.offset_conv.bias, 0.) + + self.modulator_conv = nn.Conv2d(in_channels, + 1 * kernel_size[0] * kernel_size[1], + kernel_size=kernel_size, + stride=stride, + padding=self.padding, + bias=True) + + nn.init.constant_(self.modulator_conv.weight, 0.) + nn.init.constant_(self.modulator_conv.bias, 0.) + + self.regular_conv = nn.Conv2d(in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=self.padding, + bias=bias) + + def forward(self, x): + #h, w = x.shape[2:] + #max_offset = max(h, w)/4. + + offset = self.offset_conv(x)#.clamp(-max_offset, max_offset) + modulator = 2. * torch.sigmoid(self.modulator_conv(x)) + + x = deform_conv2d( + input=x, + offset=offset, + weight=self.regular_conv.weight, + bias=self.regular_conv.bias, + padding=self.padding, + mask=modulator, + stride=self.stride, + ) + return x + + + + +### utils.py + +import torch.nn as nn + + +def build_act_layer(act_layer): + if act_layer == 'ReLU': + return nn.ReLU(inplace=True) + elif act_layer == 'SiLU': + return nn.SiLU(inplace=True) + elif act_layer == 'GELU': + return nn.GELU() + + raise NotImplementedError(f'build_act_layer does not support {act_layer}') + + +def build_norm_layer(dim, + norm_layer, + in_format='channels_last', + out_format='channels_last', + eps=1e-6): + layers = [] + if norm_layer == 'BN': + if in_format == 'channels_last': + layers.append(to_channels_first()) + layers.append(nn.BatchNorm2d(dim)) + if out_format == 'channels_last': + layers.append(to_channels_last()) + elif norm_layer == 'LN': + if in_format == 'channels_first': + layers.append(to_channels_last()) + layers.append(nn.LayerNorm(dim, eps=eps)) + if out_format == 'channels_first': + layers.append(to_channels_first()) + else: + raise NotImplementedError( + f'build_norm_layer does not support {norm_layer}') + return nn.Sequential(*layers) + + +class to_channels_first(nn.Module): + + def __init__(self): + super().__init__() + + def forward(self, x): + return x.permute(0, 3, 1, 2) + + +class to_channels_last(nn.Module): + + def __init__(self): + super().__init__() + + def forward(self, x): + return x.permute(0, 2, 3, 1) + + + +### dataset.py + +_class_labels_TR_sorted = ( + 'Airplane, Ant, Antenna, Archery, Axe, BabyCarriage, Bag, BalanceBeam, Balcony, Balloon, Basket, BasketballHoop, Beatle, Bed, Bee, Bench, Bicycle, ' + 'BicycleFrame, BicycleStand, Boat, Bonsai, BoomLift, Bridge, BunkBed, Butterfly, Button, Cable, CableLift, Cage, Camcorder, Cannon, Canoe, Car, ' + 'CarParkDropArm, Carriage, Cart, Caterpillar, CeilingLamp, Centipede, Chair, Clip, Clock, Clothes, CoatHanger, Comb, ConcretePumpTruck, Crack, Crane, ' + 'Cup, DentalChair, Desk, DeskChair, Diagram, DishRack, DoorHandle, Dragonfish, Dragonfly, Drum, Earphone, Easel, ElectricIron, Excavator, Eyeglasses, ' + 'Fan, Fence, Fencing, FerrisWheel, FireExtinguisher, Fishing, Flag, FloorLamp, Forklift, GasStation, Gate, Gear, Goal, Golf, GymEquipment, Hammock, ' + 'Handcart, Handcraft, Handrail, HangGlider, Harp, Harvester, Headset, Helicopter, Helmet, Hook, HorizontalBar, Hydrovalve, IroningTable, Jewelry, Key, ' + 'KidsPlayground, Kitchenware, Kite, Knife, Ladder, LaundryRack, Lightning, Lobster, Locust, Machine, MachineGun, MagazineRack, Mantis, Medal, MemorialArchway, ' + 'Microphone, Missile, MobileHolder, Monitor, Mosquito, Motorcycle, MovingTrolley, Mower, MusicPlayer, MusicStand, ObservationTower, Octopus, OilWell, ' + 'OlympicLogo, OperatingTable, OutdoorFitnessEquipment, Parachute, Pavilion, Piano, Pipe, PlowHarrow, PoleVault, Punchbag, Rack, Racket, Rifle, Ring, Robot, ' + 'RockClimbing, Rope, Sailboat, Satellite, Scaffold, Scale, Scissor, Scooter, Sculpture, Seadragon, Seahorse, Seal, SewingMachine, Ship, Shoe, ShoppingCart, ' + 'ShoppingTrolley, Shower, Shrimp, Signboard, Skateboarding, Skeleton, Skiing, Spade, SpeedBoat, Spider, Spoon, Stair, Stand, Stationary, SteeringWheel, ' + 'Stethoscope, Stool, Stove, StreetLamp, SweetStand, Swing, Sword, TV, Table, TableChair, TableLamp, TableTennis, Tank, Tapeline, Teapot, Telescope, Tent, ' + 'TobaccoPipe, Toy, Tractor, TrafficLight, TrafficSign, Trampoline, TransmissionTower, Tree, Tricycle, TrimmerCover, Tripod, Trombone, Truck, Trumpet, Tuba, ' + 'UAV, Umbrella, UnevenBars, UtilityPole, VacuumCleaner, Violin, Wakesurfing, Watch, WaterTower, WateringPot, Well, WellLid, Wheel, Wheelchair, WindTurbine, Windmill, WineGlass, WireWhisk, Yacht' +) +class_labels_TR_sorted = _class_labels_TR_sorted.split(', ') + + +### models/backbones/build_backbones.py + +import torch +import torch.nn as nn +from collections import OrderedDict +from torchvision.models import vgg16, vgg16_bn, VGG16_Weights, VGG16_BN_Weights, resnet50, ResNet50_Weights +# from models.pvt_v2 import pvt_v2_b0, pvt_v2_b1, pvt_v2_b2, pvt_v2_b5 +# from models.swin_v1 import swin_v1_t, swin_v1_s, swin_v1_b, swin_v1_l +# from config import Config + + +config = Config() + +def build_backbone(bb_name, pretrained=True, params_settings=''): + if bb_name == 'vgg16': + bb_net = list(vgg16(pretrained=VGG16_Weights.DEFAULT if pretrained else None).children())[0] + bb = nn.Sequential(OrderedDict({'conv1': bb_net[:4], 'conv2': bb_net[4:9], 'conv3': bb_net[9:16], 'conv4': bb_net[16:23]})) + elif bb_name == 'vgg16bn': + bb_net = list(vgg16_bn(pretrained=VGG16_BN_Weights.DEFAULT if pretrained else None).children())[0] + bb = nn.Sequential(OrderedDict({'conv1': bb_net[:6], 'conv2': bb_net[6:13], 'conv3': bb_net[13:23], 'conv4': bb_net[23:33]})) + elif bb_name == 'resnet50': + bb_net = list(resnet50(pretrained=ResNet50_Weights.DEFAULT if pretrained else None).children()) + bb = nn.Sequential(OrderedDict({'conv1': nn.Sequential(*bb_net[0:3]), 'conv2': bb_net[4], 'conv3': bb_net[5], 'conv4': bb_net[6]})) + else: + bb = eval('{}({})'.format(bb_name, params_settings)) + if pretrained: + bb = load_weights(bb, bb_name) + return bb + +def load_weights(model, model_name): + save_model = torch.load(config.weights[model_name], map_location='cpu') + model_dict = model.state_dict() + state_dict = {k: v if v.size() == model_dict[k].size() else model_dict[k] for k, v in save_model.items() if k in model_dict.keys()} + # to ignore the weights with mismatched size when I modify the backbone itself. + if not state_dict: + save_model_keys = list(save_model.keys()) + sub_item = save_model_keys[0] if len(save_model_keys) == 1 else None + state_dict = {k: v if v.size() == model_dict[k].size() else model_dict[k] for k, v in save_model[sub_item].items() if k in model_dict.keys()} + if not state_dict or not sub_item: + print('Weights are not successully loaded. Check the state dict of weights file.') + return None + else: + print('Found correct weights in the "{}" item of loaded state_dict.'.format(sub_item)) + model_dict.update(state_dict) + model.load_state_dict(model_dict) + return model + + + +### models/modules/decoder_blocks.py + +import torch +import torch.nn as nn +# from models.aspp import ASPP, ASPPDeformable +# from config import Config + + +# config = Config() + + +class BasicDecBlk(nn.Module): + def __init__(self, in_channels=64, out_channels=64, inter_channels=64): + super(BasicDecBlk, self).__init__() + inter_channels = in_channels // 4 if config.dec_channels_inter == 'adap' else 64 + self.conv_in = nn.Conv2d(in_channels, inter_channels, 3, 1, padding=1) + self.relu_in = nn.ReLU(inplace=True) + if config.dec_att == 'ASPP': + self.dec_att = ASPP(in_channels=inter_channels) + elif config.dec_att == 'ASPPDeformable': + self.dec_att = ASPPDeformable(in_channels=inter_channels) + self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, padding=1) + self.bn_in = nn.BatchNorm2d(inter_channels) if config.batch_size > 1 else nn.Identity() + self.bn_out = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity() + + def forward(self, x): + x = self.conv_in(x) + x = self.bn_in(x) + x = self.relu_in(x) + if hasattr(self, 'dec_att'): + x = self.dec_att(x) + x = self.conv_out(x) + x = self.bn_out(x) + return x + + +class ResBlk(nn.Module): + def __init__(self, in_channels=64, out_channels=None, inter_channels=64): + super(ResBlk, self).__init__() + if out_channels is None: + out_channels = in_channels + inter_channels = in_channels // 4 if config.dec_channels_inter == 'adap' else 64 + + self.conv_in = nn.Conv2d(in_channels, inter_channels, 3, 1, padding=1) + self.bn_in = nn.BatchNorm2d(inter_channels) if config.batch_size > 1 else nn.Identity() + self.relu_in = nn.ReLU(inplace=True) + + if config.dec_att == 'ASPP': + self.dec_att = ASPP(in_channels=inter_channels) + elif config.dec_att == 'ASPPDeformable': + self.dec_att = ASPPDeformable(in_channels=inter_channels) + + self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, padding=1) + self.bn_out = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity() + + self.conv_resi = nn.Conv2d(in_channels, out_channels, 1, 1, 0) + + def forward(self, x): + _x = self.conv_resi(x) + x = self.conv_in(x) + x = self.bn_in(x) + x = self.relu_in(x) + if hasattr(self, 'dec_att'): + x = self.dec_att(x) + x = self.conv_out(x) + x = self.bn_out(x) + return x + _x + + + +### models/modules/lateral_blocks.py + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from functools import partial + +# from config import Config + + +# config = Config() + + +class BasicLatBlk(nn.Module): + def __init__(self, in_channels=64, out_channels=64, inter_channels=64): + super(BasicLatBlk, self).__init__() + inter_channels = in_channels // 4 if config.dec_channels_inter == 'adap' else 64 + self.conv = nn.Conv2d(in_channels, out_channels, 1, 1, 0) + + def forward(self, x): + x = self.conv(x) + return x + + + +### models/modules/aspp.py + +import torch +import torch.nn as nn +import torch.nn.functional as F +# from models.deform_conv import DeformableConv2d +# from config import Config + + +# config = Config() + + +class _ASPPModule(nn.Module): + def __init__(self, in_channels, planes, kernel_size, padding, dilation): + super(_ASPPModule, self).__init__() + self.atrous_conv = nn.Conv2d(in_channels, planes, kernel_size=kernel_size, + stride=1, padding=padding, dilation=dilation, bias=False) + self.bn = nn.BatchNorm2d(planes) if config.batch_size > 1 else nn.Identity() + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + x = self.atrous_conv(x) + x = self.bn(x) + + return self.relu(x) + + +class ASPP(nn.Module): + def __init__(self, in_channels=64, out_channels=None, output_stride=16): + super(ASPP, self).__init__() + self.down_scale = 1 + if out_channels is None: + out_channels = in_channels + self.in_channelster = 256 // self.down_scale + if output_stride == 16: + dilations = [1, 6, 12, 18] + elif output_stride == 8: + dilations = [1, 12, 24, 36] + else: + raise NotImplementedError + + self.aspp1 = _ASPPModule(in_channels, self.in_channelster, 1, padding=0, dilation=dilations[0]) + self.aspp2 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[1], dilation=dilations[1]) + self.aspp3 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[2], dilation=dilations[2]) + self.aspp4 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[3], dilation=dilations[3]) + + self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)), + nn.Conv2d(in_channels, self.in_channelster, 1, stride=1, bias=False), + nn.BatchNorm2d(self.in_channelster) if config.batch_size > 1 else nn.Identity(), + nn.ReLU(inplace=True)) + self.conv1 = nn.Conv2d(self.in_channelster * 5, out_channels, 1, bias=False) + self.bn1 = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity() + self.relu = nn.ReLU(inplace=True) + self.dropout = nn.Dropout(0.5) + + def forward(self, x): + x1 = self.aspp1(x) + x2 = self.aspp2(x) + x3 = self.aspp3(x) + x4 = self.aspp4(x) + x5 = self.global_avg_pool(x) + x5 = F.interpolate(x5, size=x1.size()[2:], mode='bilinear', align_corners=True) + x = torch.cat((x1, x2, x3, x4, x5), dim=1) + + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + + return self.dropout(x) + + +##################### Deformable +class _ASPPModuleDeformable(nn.Module): + def __init__(self, in_channels, planes, kernel_size, padding): + super(_ASPPModuleDeformable, self).__init__() + self.atrous_conv = DeformableConv2d(in_channels, planes, kernel_size=kernel_size, + stride=1, padding=padding, bias=False) + self.bn = nn.BatchNorm2d(planes) if config.batch_size > 1 else nn.Identity() + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + x = self.atrous_conv(x) + x = self.bn(x) + + return self.relu(x) + + +class ASPPDeformable(nn.Module): + def __init__(self, in_channels, out_channels=None, parallel_block_sizes=[1, 3, 7]): + super(ASPPDeformable, self).__init__() + self.down_scale = 1 + if out_channels is None: + out_channels = in_channels + self.in_channelster = 256 // self.down_scale + + self.aspp1 = _ASPPModuleDeformable(in_channels, self.in_channelster, 1, padding=0) + self.aspp_deforms = nn.ModuleList([ + _ASPPModuleDeformable(in_channels, self.in_channelster, conv_size, padding=int(conv_size//2)) for conv_size in parallel_block_sizes + ]) + + self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)), + nn.Conv2d(in_channels, self.in_channelster, 1, stride=1, bias=False), + nn.BatchNorm2d(self.in_channelster) if config.batch_size > 1 else nn.Identity(), + nn.ReLU(inplace=True)) + self.conv1 = nn.Conv2d(self.in_channelster * (2 + len(self.aspp_deforms)), out_channels, 1, bias=False) + self.bn1 = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity() + self.relu = nn.ReLU(inplace=True) + self.dropout = nn.Dropout(0.5) + + def forward(self, x): + x1 = self.aspp1(x) + x_aspp_deforms = [aspp_deform(x) for aspp_deform in self.aspp_deforms] + x5 = self.global_avg_pool(x) + x5 = F.interpolate(x5, size=x1.size()[2:], mode='bilinear', align_corners=True) + x = torch.cat((x1, *x_aspp_deforms, x5), dim=1) + + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + + return self.dropout(x) + + + +### models/refinement/refiner.py + +import torch +import torch.nn as nn +from collections import OrderedDict +import torch +import torch.nn as nn +import torch.nn.functional as F +from torchvision.models import vgg16, vgg16_bn +from torchvision.models import resnet50 + +# from config import Config +# from dataset import class_labels_TR_sorted +# from models.build_backbone import build_backbone +# from models.decoder_blocks import BasicDecBlk +# from models.lateral_blocks import BasicLatBlk +# from models.ing import * +# from models.stem_layer import StemLayer + + +class RefinerPVTInChannels4(nn.Module): + def __init__(self, in_channels=3+1): + super(RefinerPVTInChannels4, self).__init__() + self.config = Config() + self.epoch = 1 + self.bb = build_backbone(self.config.bb, params_settings='in_channels=4') + + lateral_channels_in_collection = { + 'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64], + 'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64], + 'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192], + } + channels = lateral_channels_in_collection[self.config.bb] + self.squeeze_module = BasicDecBlk(channels[0], channels[0]) + + self.decoder = Decoder(channels) + + if 0: + for key, value in self.named_parameters(): + if 'bb.' in key: + value.requires_grad = False + + def forward(self, x): + if isinstance(x, list): + x = torch.cat(x, dim=1) + ########## Encoder ########## + if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']: + x1 = self.bb.conv1(x) + x2 = self.bb.conv2(x1) + x3 = self.bb.conv3(x2) + x4 = self.bb.conv4(x3) + else: + x1, x2, x3, x4 = self.bb(x) + + x4 = self.squeeze_module(x4) + + ########## Decoder ########## + + features = [x, x1, x2, x3, x4] + scaled_preds = self.decoder(features) + + return scaled_preds + + +class Refiner(nn.Module): + def __init__(self, in_channels=3+1): + super(Refiner, self).__init__() + self.config = Config() + self.epoch = 1 + self.stem_layer = StemLayer(in_channels=in_channels, inter_channels=48, out_channels=3, norm_layer='BN' if self.config.batch_size > 1 else 'LN') + self.bb = build_backbone(self.config.bb) + + lateral_channels_in_collection = { + 'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64], + 'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64], + 'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192], + } + channels = lateral_channels_in_collection[self.config.bb] + self.squeeze_module = BasicDecBlk(channels[0], channels[0]) + + self.decoder = Decoder(channels) + + if 0: + for key, value in self.named_parameters(): + if 'bb.' in key: + value.requires_grad = False + + def forward(self, x): + if isinstance(x, list): + x = torch.cat(x, dim=1) + x = self.stem_layer(x) + ########## Encoder ########## + if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']: + x1 = self.bb.conv1(x) + x2 = self.bb.conv2(x1) + x3 = self.bb.conv3(x2) + x4 = self.bb.conv4(x3) + else: + x1, x2, x3, x4 = self.bb(x) + + x4 = self.squeeze_module(x4) + + ########## Decoder ########## + + features = [x, x1, x2, x3, x4] + scaled_preds = self.decoder(features) + + return scaled_preds + + +class Decoder(nn.Module): + def __init__(self, channels): + super(Decoder, self).__init__() + self.config = Config() + DecoderBlock = eval('BasicDecBlk') + LateralBlock = eval('BasicLatBlk') + + self.decoder_block4 = DecoderBlock(channels[0], channels[1]) + self.decoder_block3 = DecoderBlock(channels[1], channels[2]) + self.decoder_block2 = DecoderBlock(channels[2], channels[3]) + self.decoder_block1 = DecoderBlock(channels[3], channels[3]//2) + + self.lateral_block4 = LateralBlock(channels[1], channels[1]) + self.lateral_block3 = LateralBlock(channels[2], channels[2]) + self.lateral_block2 = LateralBlock(channels[3], channels[3]) + + if self.config.ms_supervision: + self.conv_ms_spvn_4 = nn.Conv2d(channels[1], 1, 1, 1, 0) + self.conv_ms_spvn_3 = nn.Conv2d(channels[2], 1, 1, 1, 0) + self.conv_ms_spvn_2 = nn.Conv2d(channels[3], 1, 1, 1, 0) + self.conv_out1 = nn.Sequential(nn.Conv2d(channels[3]//2, 1, 1, 1, 0)) + + def forward(self, features): + x, x1, x2, x3, x4 = features + outs = [] + p4 = self.decoder_block4(x4) + _p4 = F.interpolate(p4, size=x3.shape[2:], mode='bilinear', align_corners=True) + _p3 = _p4 + self.lateral_block4(x3) + + p3 = self.decoder_block3(_p3) + _p3 = F.interpolate(p3, size=x2.shape[2:], mode='bilinear', align_corners=True) + _p2 = _p3 + self.lateral_block3(x2) + + p2 = self.decoder_block2(_p2) + _p2 = F.interpolate(p2, size=x1.shape[2:], mode='bilinear', align_corners=True) + _p1 = _p2 + self.lateral_block2(x1) + + _p1 = self.decoder_block1(_p1) + _p1 = F.interpolate(_p1, size=x.shape[2:], mode='bilinear', align_corners=True) + p1_out = self.conv_out1(_p1) + + if self.config.ms_supervision: + outs.append(self.conv_ms_spvn_4(p4)) + outs.append(self.conv_ms_spvn_3(p3)) + outs.append(self.conv_ms_spvn_2(p2)) + outs.append(p1_out) + return outs + + +class RefUNet(nn.Module): + # Refinement + def __init__(self, in_channels=3+1): + super(RefUNet, self).__init__() + self.encoder_1 = nn.Sequential( + nn.Conv2d(in_channels, 64, 3, 1, 1), + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.encoder_2 = nn.Sequential( + nn.MaxPool2d(2, 2, ceil_mode=True), + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.encoder_3 = nn.Sequential( + nn.MaxPool2d(2, 2, ceil_mode=True), + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.encoder_4 = nn.Sequential( + nn.MaxPool2d(2, 2, ceil_mode=True), + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.pool4 = nn.MaxPool2d(2, 2, ceil_mode=True) + ##### + self.decoder_5 = nn.Sequential( + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + ##### + self.decoder_4 = nn.Sequential( + nn.Conv2d(128, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.decoder_3 = nn.Sequential( + nn.Conv2d(128, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.decoder_2 = nn.Sequential( + nn.Conv2d(128, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.decoder_1 = nn.Sequential( + nn.Conv2d(128, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.conv_d0 = nn.Conv2d(64, 1, 3, 1, 1) + + self.upscore2 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) + + def forward(self, x): + outs = [] + if isinstance(x, list): + x = torch.cat(x, dim=1) + hx = x + + hx1 = self.encoder_1(hx) + hx2 = self.encoder_2(hx1) + hx3 = self.encoder_3(hx2) + hx4 = self.encoder_4(hx3) + + hx = self.decoder_5(self.pool4(hx4)) + hx = torch.cat((self.upscore2(hx), hx4), 1) + + d4 = self.decoder_4(hx) + hx = torch.cat((self.upscore2(d4), hx3), 1) + + d3 = self.decoder_3(hx) + hx = torch.cat((self.upscore2(d3), hx2), 1) + + d2 = self.decoder_2(hx) + hx = torch.cat((self.upscore2(d2), hx1), 1) + + d1 = self.decoder_1(hx) + + x = self.conv_d0(d1) + outs.append(x) + return outs + + + +### models/stem_layer.py + +import torch.nn as nn +# from utils import build_act_layer, build_norm_layer + + +class StemLayer(nn.Module): + r""" Stem layer of InternImage + Args: + in_channels (int): number of input channels + out_channels (int): number of output channels + act_layer (str): activation layer + norm_layer (str): normalization layer + """ + + def __init__(self, + in_channels=3+1, + inter_channels=48, + out_channels=96, + act_layer='GELU', + norm_layer='BN'): + super().__init__() + self.conv1 = nn.Conv2d(in_channels, + inter_channels, + kernel_size=3, + stride=1, + padding=1) + self.norm1 = build_norm_layer( + inter_channels, norm_layer, 'channels_first', 'channels_first' + ) + self.act = build_act_layer(act_layer) + self.conv2 = nn.Conv2d(inter_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1) + self.norm2 = build_norm_layer( + out_channels, norm_layer, 'channels_first', 'channels_first' + ) + + def forward(self, x): + x = self.conv1(x) + x = self.norm1(x) + x = self.act(x) + x = self.conv2(x) + x = self.norm2(x) + return x + + +### models/birefnet.py + +import torch +import torch.nn as nn +import torch.nn.functional as F +from kornia.filters import laplacian +from transformers import PreTrainedModel +from einops import rearrange + +# from config import Config +# from dataset import class_labels_TR_sorted +# from models.build_backbone import build_backbone +# from models.decoder_blocks import BasicDecBlk, ResBlk, HierarAttDecBlk +# from models.lateral_blocks import BasicLatBlk +# from models.aspp import ASPP, ASPPDeformable +# from models.ing import * +# from models.refiner import Refiner, RefinerPVTInChannels4, RefUNet +# from models.stem_layer import StemLayer +from BiRefNet_config import BiRefNetConfig + + +def image2patches(image, grid_h=2, grid_w=2, patch_ref=None, transformation='b c (hg h) (wg w) -> (b hg wg) c h w'): + if patch_ref is not None: + grid_h, grid_w = image.shape[-2] // patch_ref.shape[-2], image.shape[-1] // patch_ref.shape[-1] + patches = rearrange(image, transformation, hg=grid_h, wg=grid_w) + return patches + +def patches2image(patches, grid_h=2, grid_w=2, patch_ref=None, transformation='(b hg wg) c h w -> b c (hg h) (wg w)'): + if patch_ref is not None: + grid_h, grid_w = patch_ref.shape[-2] // patches[0].shape[-2], patch_ref.shape[-1] // patches[0].shape[-1] + image = rearrange(patches, transformation, hg=grid_h, wg=grid_w) + return image + +class BiRefNet( + PreTrainedModel +): + config_class = BiRefNetConfig + def __init__(self, bb_pretrained=True, config=BiRefNetConfig()): + super(BiRefNet, self).__init__(config) + bb_pretrained = config.bb_pretrained + self.config = Config() + self.epoch = 1 + self.bb = build_backbone(self.config.bb, pretrained=bb_pretrained) + + channels = self.config.lateral_channels_in_collection + + if self.config.auxiliary_classification: + self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) + self.cls_head = nn.Sequential( + nn.Linear(channels[0], len(class_labels_TR_sorted)) + ) + + if self.config.squeeze_block: + self.squeeze_module = nn.Sequential(*[ + eval(self.config.squeeze_block.split('_x')[0])(channels[0]+sum(self.config.cxt), channels[0]) + for _ in range(eval(self.config.squeeze_block.split('_x')[1])) + ]) + + self.decoder = Decoder(channels) + + if self.config.ender: + self.dec_end = nn.Sequential( + nn.Conv2d(1, 16, 3, 1, 1), + nn.Conv2d(16, 1, 3, 1, 1), + nn.ReLU(inplace=True), + ) + + # refine patch-level segmentation + if self.config.refine: + if self.config.refine == 'itself': + self.stem_layer = StemLayer(in_channels=3+1, inter_channels=48, out_channels=3, norm_layer='BN' if self.config.batch_size > 1 else 'LN') + else: + self.refiner = eval('{}({})'.format(self.config.refine, 'in_channels=3+1')) + + if self.config.freeze_bb: + # Freeze the backbone... + print(self.named_parameters()) + for key, value in self.named_parameters(): + if 'bb.' in key and 'refiner.' not in key: + value.requires_grad = False + + def forward_enc(self, x): + if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']: + x1 = self.bb.conv1(x); x2 = self.bb.conv2(x1); x3 = self.bb.conv3(x2); x4 = self.bb.conv4(x3) + else: + x1, x2, x3, x4 = self.bb(x) + if self.config.mul_scl_ipt == 'cat': + B, C, H, W = x.shape + x1_, x2_, x3_, x4_ = self.bb(F.interpolate(x, size=(H//2, W//2), mode='bilinear', align_corners=True)) + x1 = torch.cat([x1, F.interpolate(x1_, size=x1.shape[2:], mode='bilinear', align_corners=True)], dim=1) + x2 = torch.cat([x2, F.interpolate(x2_, size=x2.shape[2:], mode='bilinear', align_corners=True)], dim=1) + x3 = torch.cat([x3, F.interpolate(x3_, size=x3.shape[2:], mode='bilinear', align_corners=True)], dim=1) + x4 = torch.cat([x4, F.interpolate(x4_, size=x4.shape[2:], mode='bilinear', align_corners=True)], dim=1) + elif self.config.mul_scl_ipt == 'add': + B, C, H, W = x.shape + x1_, x2_, x3_, x4_ = self.bb(F.interpolate(x, size=(H//2, W//2), mode='bilinear', align_corners=True)) + x1 = x1 + F.interpolate(x1_, size=x1.shape[2:], mode='bilinear', align_corners=True) + x2 = x2 + F.interpolate(x2_, size=x2.shape[2:], mode='bilinear', align_corners=True) + x3 = x3 + F.interpolate(x3_, size=x3.shape[2:], mode='bilinear', align_corners=True) + x4 = x4 + F.interpolate(x4_, size=x4.shape[2:], mode='bilinear', align_corners=True) + class_preds = self.cls_head(self.avgpool(x4).view(x4.shape[0], -1)) if self.training and self.config.auxiliary_classification else None + if self.config.cxt: + x4 = torch.cat( + ( + *[ + F.interpolate(x1, size=x4.shape[2:], mode='bilinear', align_corners=True), + F.interpolate(x2, size=x4.shape[2:], mode='bilinear', align_corners=True), + F.interpolate(x3, size=x4.shape[2:], mode='bilinear', align_corners=True), + ][-len(self.config.cxt):], + x4 + ), + dim=1 + ) + return (x1, x2, x3, x4), class_preds + + def forward_ori(self, x): + ########## Encoder ########## + (x1, x2, x3, x4), class_preds = self.forward_enc(x) + if self.config.squeeze_block: + x4 = self.squeeze_module(x4) + ########## Decoder ########## + features = [x, x1, x2, x3, x4] + if self.training and self.config.out_ref: + features.append(laplacian(torch.mean(x, dim=1).unsqueeze(1), kernel_size=5)) + scaled_preds = self.decoder(features) + return scaled_preds, class_preds + + def forward(self, x): + scaled_preds, class_preds = self.forward_ori(x) + class_preds_lst = [class_preds] + return [scaled_preds, class_preds_lst] if self.training else scaled_preds + + +class Decoder(nn.Module): + def __init__(self, channels): + super(Decoder, self).__init__() + self.config = Config() + DecoderBlock = eval(self.config.dec_blk) + LateralBlock = eval(self.config.lat_blk) + + if self.config.dec_ipt: + self.split = self.config.dec_ipt_split + N_dec_ipt = 64 + DBlock = SimpleConvs + ic = 64 + ipt_cha_opt = 1 + self.ipt_blk5 = DBlock(2**10*3 if self.split else 3, [N_dec_ipt, channels[0]//8][ipt_cha_opt], inter_channels=ic) + self.ipt_blk4 = DBlock(2**8*3 if self.split else 3, [N_dec_ipt, channels[0]//8][ipt_cha_opt], inter_channels=ic) + self.ipt_blk3 = DBlock(2**6*3 if self.split else 3, [N_dec_ipt, channels[1]//8][ipt_cha_opt], inter_channels=ic) + self.ipt_blk2 = DBlock(2**4*3 if self.split else 3, [N_dec_ipt, channels[2]//8][ipt_cha_opt], inter_channels=ic) + self.ipt_blk1 = DBlock(2**0*3 if self.split else 3, [N_dec_ipt, channels[3]//8][ipt_cha_opt], inter_channels=ic) + else: + self.split = None + + self.decoder_block4 = DecoderBlock(channels[0]+([N_dec_ipt, channels[0]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[1]) + self.decoder_block3 = DecoderBlock(channels[1]+([N_dec_ipt, channels[0]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[2]) + self.decoder_block2 = DecoderBlock(channels[2]+([N_dec_ipt, channels[1]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[3]) + self.decoder_block1 = DecoderBlock(channels[3]+([N_dec_ipt, channels[2]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[3]//2) + self.conv_out1 = nn.Sequential(nn.Conv2d(channels[3]//2+([N_dec_ipt, channels[3]//8][ipt_cha_opt] if self.config.dec_ipt else 0), 1, 1, 1, 0)) + + self.lateral_block4 = LateralBlock(channels[1], channels[1]) + self.lateral_block3 = LateralBlock(channels[2], channels[2]) + self.lateral_block2 = LateralBlock(channels[3], channels[3]) + + if self.config.ms_supervision: + self.conv_ms_spvn_4 = nn.Conv2d(channels[1], 1, 1, 1, 0) + self.conv_ms_spvn_3 = nn.Conv2d(channels[2], 1, 1, 1, 0) + self.conv_ms_spvn_2 = nn.Conv2d(channels[3], 1, 1, 1, 0) + + if self.config.out_ref: + _N = 16 + self.gdt_convs_4 = nn.Sequential(nn.Conv2d(channels[1], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True)) + self.gdt_convs_3 = nn.Sequential(nn.Conv2d(channels[2], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True)) + self.gdt_convs_2 = nn.Sequential(nn.Conv2d(channels[3], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True)) + + self.gdt_convs_pred_4 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + self.gdt_convs_pred_3 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + self.gdt_convs_pred_2 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + + self.gdt_convs_attn_4 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + self.gdt_convs_attn_3 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + self.gdt_convs_attn_2 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + + def forward(self, features): + if self.training and self.config.out_ref: + outs_gdt_pred = [] + outs_gdt_label = [] + x, x1, x2, x3, x4, gdt_gt = features + else: + x, x1, x2, x3, x4 = features + outs = [] + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=x4, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + x4 = torch.cat((x4, self.ipt_blk5(F.interpolate(patches_batch, size=x4.shape[2:], mode='bilinear', align_corners=True))), 1) + p4 = self.decoder_block4(x4) + m4 = self.conv_ms_spvn_4(p4) if self.config.ms_supervision and self.training else None + if self.config.out_ref: + p4_gdt = self.gdt_convs_4(p4) + if self.training: + # >> GT: + m4_dia = m4 + gdt_label_main_4 = gdt_gt * F.interpolate(m4_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True) + outs_gdt_label.append(gdt_label_main_4) + # >> Pred: + gdt_pred_4 = self.gdt_convs_pred_4(p4_gdt) + outs_gdt_pred.append(gdt_pred_4) + gdt_attn_4 = self.gdt_convs_attn_4(p4_gdt).sigmoid() + # >> Finally: + p4 = p4 * gdt_attn_4 + _p4 = F.interpolate(p4, size=x3.shape[2:], mode='bilinear', align_corners=True) + _p3 = _p4 + self.lateral_block4(x3) + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=_p3, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + _p3 = torch.cat((_p3, self.ipt_blk4(F.interpolate(patches_batch, size=x3.shape[2:], mode='bilinear', align_corners=True))), 1) + p3 = self.decoder_block3(_p3) + m3 = self.conv_ms_spvn_3(p3) if self.config.ms_supervision and self.training else None + if self.config.out_ref: + p3_gdt = self.gdt_convs_3(p3) + if self.training: + # >> GT: + # m3 --dilation--> m3_dia + # G_3^gt * m3_dia --> G_3^m, which is the label of gradient + m3_dia = m3 + gdt_label_main_3 = gdt_gt * F.interpolate(m3_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True) + outs_gdt_label.append(gdt_label_main_3) + # >> Pred: + # p3 --conv--BN--> F_3^G, where F_3^G predicts the \hat{G_3} with xx + # F_3^G --sigmoid--> A_3^G + gdt_pred_3 = self.gdt_convs_pred_3(p3_gdt) + outs_gdt_pred.append(gdt_pred_3) + gdt_attn_3 = self.gdt_convs_attn_3(p3_gdt).sigmoid() + # >> Finally: + # p3 = p3 * A_3^G + p3 = p3 * gdt_attn_3 + _p3 = F.interpolate(p3, size=x2.shape[2:], mode='bilinear', align_corners=True) + _p2 = _p3 + self.lateral_block3(x2) + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=_p2, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + _p2 = torch.cat((_p2, self.ipt_blk3(F.interpolate(patches_batch, size=x2.shape[2:], mode='bilinear', align_corners=True))), 1) + p2 = self.decoder_block2(_p2) + m2 = self.conv_ms_spvn_2(p2) if self.config.ms_supervision and self.training else None + if self.config.out_ref: + p2_gdt = self.gdt_convs_2(p2) + if self.training: + # >> GT: + m2_dia = m2 + gdt_label_main_2 = gdt_gt * F.interpolate(m2_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True) + outs_gdt_label.append(gdt_label_main_2) + # >> Pred: + gdt_pred_2 = self.gdt_convs_pred_2(p2_gdt) + outs_gdt_pred.append(gdt_pred_2) + gdt_attn_2 = self.gdt_convs_attn_2(p2_gdt).sigmoid() + # >> Finally: + p2 = p2 * gdt_attn_2 + _p2 = F.interpolate(p2, size=x1.shape[2:], mode='bilinear', align_corners=True) + _p1 = _p2 + self.lateral_block2(x1) + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=_p1, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + _p1 = torch.cat((_p1, self.ipt_blk2(F.interpolate(patches_batch, size=x1.shape[2:], mode='bilinear', align_corners=True))), 1) + _p1 = self.decoder_block1(_p1) + _p1 = F.interpolate(_p1, size=x.shape[2:], mode='bilinear', align_corners=True) + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=_p1, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + _p1 = torch.cat((_p1, self.ipt_blk1(F.interpolate(patches_batch, size=x.shape[2:], mode='bilinear', align_corners=True))), 1) + p1_out = self.conv_out1(_p1) + + if self.config.ms_supervision and self.training: + outs.append(m4) + outs.append(m3) + outs.append(m2) + outs.append(p1_out) + return outs if not (self.config.out_ref and self.training) else ([outs_gdt_pred, outs_gdt_label], outs) + + +class SimpleConvs(nn.Module): + def __init__( + self, in_channels: int, out_channels: int, inter_channels=64 + ) -> None: + super().__init__() + self.conv1 = nn.Conv2d(in_channels, inter_channels, 3, 1, 1) + self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, 1) + + def forward(self, x): + return self.conv_out(self.conv1(x)) diff --git a/models/RMBG/BiRefNet/birefnet_lite.py b/models/RMBG/BiRefNet/birefnet_lite.py new file mode 100644 index 0000000000000000000000000000000000000000..8de5e88533b04010a03c819c1a5ddf5ffd85fe0a --- /dev/null +++ b/models/RMBG/BiRefNet/birefnet_lite.py @@ -0,0 +1,2248 @@ +### config.py + +import os +import math + + +class Config(): + def __init__(self) -> None: + # PATH settings + self.sys_home_dir = os.path.expanduser('~') # Make up your file system as: SYS_HOME_DIR/codes/dis/BiRefNet, SYS_HOME_DIR/datasets/dis/xx, SYS_HOME_DIR/weights/xx + + # TASK settings + self.task = ['DIS5K', 'COD', 'HRSOD', 'DIS5K+HRSOD+HRS10K', 'P3M-10k'][0] + self.training_set = { + 'DIS5K': ['DIS-TR', 'DIS-TR+DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4'][0], + 'COD': 'TR-COD10K+TR-CAMO', + 'HRSOD': ['TR-DUTS', 'TR-HRSOD', 'TR-UHRSD', 'TR-DUTS+TR-HRSOD', 'TR-DUTS+TR-UHRSD', 'TR-HRSOD+TR-UHRSD', 'TR-DUTS+TR-HRSOD+TR-UHRSD'][5], + 'DIS5K+HRSOD+HRS10K': 'DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4+DIS-TR+TE-HRS10K+TE-HRSOD+TE-UHRSD+TR-HRS10K+TR-HRSOD+TR-UHRSD', # leave DIS-VD for evaluation. + 'P3M-10k': 'TR-P3M-10k', + }[self.task] + self.prompt4loc = ['dense', 'sparse'][0] + + # Faster-Training settings + self.load_all = True + self.compile = True # 1. Trigger CPU memory leak in some extend, which is an inherent problem of PyTorch. + # Machines with > 70GB CPU memory can run the whole training on DIS5K with default setting. + # 2. Higher PyTorch version may fix it: https://github.com/pytorch/pytorch/issues/119607. + # 3. But compile in Pytorch > 2.0.1 seems to bring no acceleration for training. + self.precisionHigh = True + + # MODEL settings + self.ms_supervision = True + self.out_ref = self.ms_supervision and True + self.dec_ipt = True + self.dec_ipt_split = True + self.cxt_num = [0, 3][1] # multi-scale skip connections from encoder + self.mul_scl_ipt = ['', 'add', 'cat'][2] + self.dec_att = ['', 'ASPP', 'ASPPDeformable'][2] + self.squeeze_block = ['', 'BasicDecBlk_x1', 'ResBlk_x4', 'ASPP_x3', 'ASPPDeformable_x3'][1] + self.dec_blk = ['BasicDecBlk', 'ResBlk', 'HierarAttDecBlk'][0] + + # TRAINING settings + self.batch_size = 4 + self.IoU_finetune_last_epochs = [ + 0, + { + 'DIS5K': -50, + 'COD': -20, + 'HRSOD': -20, + 'DIS5K+HRSOD+HRS10K': -20, + 'P3M-10k': -20, + }[self.task] + ][1] # choose 0 to skip + self.lr = (1e-4 if 'DIS5K' in self.task else 1e-5) * math.sqrt(self.batch_size / 4) # DIS needs high lr to converge faster. Adapt the lr linearly + self.size = 1024 + self.num_workers = max(4, self.batch_size) # will be decrease to min(it, batch_size) at the initialization of the data_loader + + # Backbone settings + self.bb = [ + 'vgg16', 'vgg16bn', 'resnet50', # 0, 1, 2 + 'swin_v1_t', 'swin_v1_s', # 3, 4 + 'swin_v1_b', 'swin_v1_l', # 5-bs9, 6-bs4 + 'pvt_v2_b0', 'pvt_v2_b1', # 7, 8 + 'pvt_v2_b2', 'pvt_v2_b5', # 9-bs10, 10-bs5 + ][3] + self.lateral_channels_in_collection = { + 'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64], + 'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64], + 'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192], + 'swin_v1_t': [768, 384, 192, 96], 'swin_v1_s': [768, 384, 192, 96], + 'pvt_v2_b0': [256, 160, 64, 32], 'pvt_v2_b1': [512, 320, 128, 64], + }[self.bb] + if self.mul_scl_ipt == 'cat': + self.lateral_channels_in_collection = [channel * 2 for channel in self.lateral_channels_in_collection] + self.cxt = self.lateral_channels_in_collection[1:][::-1][-self.cxt_num:] if self.cxt_num else [] + + # MODEL settings - inactive + self.lat_blk = ['BasicLatBlk'][0] + self.dec_channels_inter = ['fixed', 'adap'][0] + self.refine = ['', 'itself', 'RefUNet', 'Refiner', 'RefinerPVTInChannels4'][0] + self.progressive_ref = self.refine and True + self.ender = self.progressive_ref and False + self.scale = self.progressive_ref and 2 + self.auxiliary_classification = False # Only for DIS5K, where class labels are saved in `dataset.py`. + self.refine_iteration = 1 + self.freeze_bb = False + self.model = [ + 'BiRefNet', + ][0] + if self.dec_blk == 'HierarAttDecBlk': + self.batch_size = 2 ** [0, 1, 2, 3, 4][2] + + # TRAINING settings - inactive + self.preproc_methods = ['flip', 'enhance', 'rotate', 'pepper', 'crop'][:4] + self.optimizer = ['Adam', 'AdamW'][1] + self.lr_decay_epochs = [1e5] # Set to negative N to decay the lr in the last N-th epoch. + self.lr_decay_rate = 0.5 + # Loss + self.lambdas_pix_last = { + # not 0 means opening this loss + # original rate -- 1 : 30 : 1.5 : 0.2, bce x 30 + 'bce': 30 * 1, # high performance + 'iou': 0.5 * 1, # 0 / 255 + 'iou_patch': 0.5 * 0, # 0 / 255, win_size = (64, 64) + 'mse': 150 * 0, # can smooth the saliency map + 'triplet': 3 * 0, + 'reg': 100 * 0, + 'ssim': 10 * 1, # help contours, + 'cnt': 5 * 0, # help contours + 'structure': 5 * 0, # structure loss from codes of MVANet. A little improvement on DIS-TE[1,2,3], a bit more decrease on DIS-TE4. + } + self.lambdas_cls = { + 'ce': 5.0 + } + # Adv + self.lambda_adv_g = 10. * 0 # turn to 0 to avoid adv training + self.lambda_adv_d = 3. * (self.lambda_adv_g > 0) + + # PATH settings - inactive + self.data_root_dir = os.path.join(self.sys_home_dir, 'datasets/dis') + self.weights_root_dir = os.path.join(self.sys_home_dir, 'weights') + self.weights = { + 'pvt_v2_b2': os.path.join(self.weights_root_dir, 'pvt_v2_b2.pth'), + 'pvt_v2_b5': os.path.join(self.weights_root_dir, ['pvt_v2_b5.pth', 'pvt_v2_b5_22k.pth'][0]), + 'swin_v1_b': os.path.join(self.weights_root_dir, ['swin_base_patch4_window12_384_22kto1k.pth', 'swin_base_patch4_window12_384_22k.pth'][0]), + 'swin_v1_l': os.path.join(self.weights_root_dir, ['swin_large_patch4_window12_384_22kto1k.pth', 'swin_large_patch4_window12_384_22k.pth'][0]), + 'swin_v1_t': os.path.join(self.weights_root_dir, ['swin_tiny_patch4_window7_224_22kto1k_finetune.pth'][0]), + 'swin_v1_s': os.path.join(self.weights_root_dir, ['swin_small_patch4_window7_224_22kto1k_finetune.pth'][0]), + 'pvt_v2_b0': os.path.join(self.weights_root_dir, ['pvt_v2_b0.pth'][0]), + 'pvt_v2_b1': os.path.join(self.weights_root_dir, ['pvt_v2_b1.pth'][0]), + } + + # Callbacks - inactive + self.verbose_eval = True + self.only_S_MAE = False + self.use_fp16 = False # Bugs. It may cause nan in training. + self.SDPA_enabled = False # Bugs. Slower and errors occur in multi-GPUs + + # others + self.device = [0, 'cpu'][0] # .to(0) == .to('cuda:0') + + self.batch_size_valid = 1 + self.rand_seed = 7 + # run_sh_file = [f for f in os.listdir('.') if 'train.sh' == f] + [os.path.join('..', f) for f in os.listdir('..') if 'train.sh' == f] + # with open(run_sh_file[0], 'r') as f: + # lines = f.readlines() + # self.save_last = int([l.strip() for l in lines if '"{}")'.format(self.task) in l and 'val_last=' in l][0].split('val_last=')[-1].split()[0]) + # self.save_step = int([l.strip() for l in lines if '"{}")'.format(self.task) in l and 'step=' in l][0].split('step=')[-1].split()[0]) + # self.val_step = [0, self.save_step][0] + + def print_task(self) -> None: + # Return task for choosing settings in shell scripts. + print(self.task) + + + +### models/backbones/pvt_v2.py + +import torch +import torch.nn as nn +from functools import partial + +from timm.models.layers import DropPath, to_2tuple, trunc_normal_ +from timm.models.registry import register_model + +import math + +# from config import Config + +# config = Config() + +class Mlp(nn.Module): + def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.dwconv = DWConv(hidden_features) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def forward(self, x, H, W): + x = self.fc1(x) + x = self.dwconv(x, H, W) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +class Attention(nn.Module): + def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., sr_ratio=1): + super().__init__() + assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}." + + self.dim = dim + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = qk_scale or head_dim ** -0.5 + + self.q = nn.Linear(dim, dim, bias=qkv_bias) + self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias) + self.attn_drop_prob = attn_drop + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + self.sr_ratio = sr_ratio + if sr_ratio > 1: + self.sr = nn.Conv2d(dim, dim, kernel_size=sr_ratio, stride=sr_ratio) + self.norm = nn.LayerNorm(dim) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def forward(self, x, H, W): + B, N, C = x.shape + q = self.q(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) + + if self.sr_ratio > 1: + x_ = x.permute(0, 2, 1).reshape(B, C, H, W) + x_ = self.sr(x_).reshape(B, C, -1).permute(0, 2, 1) + x_ = self.norm(x_) + kv = self.kv(x_).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + else: + kv = self.kv(x).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + k, v = kv[0], kv[1] + + if config.SDPA_enabled: + x = torch.nn.functional.scaled_dot_product_attention( + q, k, v, + attn_mask=None, dropout_p=self.attn_drop_prob, is_causal=False + ).transpose(1, 2).reshape(B, N, C) + else: + attn = (q @ k.transpose(-2, -1)) * self.scale + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B, N, C) + x = self.proj(x) + x = self.proj_drop(x) + + return x + + +class Block(nn.Module): + + def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., + drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, sr_ratio=1): + super().__init__() + self.norm1 = norm_layer(dim) + self.attn = Attention( + dim, + num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, + attn_drop=attn_drop, proj_drop=drop, sr_ratio=sr_ratio) + # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def forward(self, x, H, W): + x = x + self.drop_path(self.attn(self.norm1(x), H, W)) + x = x + self.drop_path(self.mlp(self.norm2(x), H, W)) + + return x + + +class OverlapPatchEmbed(nn.Module): + """ Image to Patch Embedding + """ + + def __init__(self, img_size=224, patch_size=7, stride=4, in_channels=3, embed_dim=768): + super().__init__() + img_size = to_2tuple(img_size) + patch_size = to_2tuple(patch_size) + + self.img_size = img_size + self.patch_size = patch_size + self.H, self.W = img_size[0] // patch_size[0], img_size[1] // patch_size[1] + self.num_patches = self.H * self.W + self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=stride, + padding=(patch_size[0] // 2, patch_size[1] // 2)) + self.norm = nn.LayerNorm(embed_dim) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def forward(self, x): + x = self.proj(x) + _, _, H, W = x.shape + x = x.flatten(2).transpose(1, 2) + x = self.norm(x) + + return x, H, W + + +class PyramidVisionTransformerImpr(nn.Module): + def __init__(self, img_size=224, patch_size=16, in_channels=3, num_classes=1000, embed_dims=[64, 128, 256, 512], + num_heads=[1, 2, 4, 8], mlp_ratios=[4, 4, 4, 4], qkv_bias=False, qk_scale=None, drop_rate=0., + attn_drop_rate=0., drop_path_rate=0., norm_layer=nn.LayerNorm, + depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1]): + super().__init__() + self.num_classes = num_classes + self.depths = depths + + # patch_embed + self.patch_embed1 = OverlapPatchEmbed(img_size=img_size, patch_size=7, stride=4, in_channels=in_channels, + embed_dim=embed_dims[0]) + self.patch_embed2 = OverlapPatchEmbed(img_size=img_size // 4, patch_size=3, stride=2, in_channels=embed_dims[0], + embed_dim=embed_dims[1]) + self.patch_embed3 = OverlapPatchEmbed(img_size=img_size // 8, patch_size=3, stride=2, in_channels=embed_dims[1], + embed_dim=embed_dims[2]) + self.patch_embed4 = OverlapPatchEmbed(img_size=img_size // 16, patch_size=3, stride=2, in_channels=embed_dims[2], + embed_dim=embed_dims[3]) + + # transformer encoder + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule + cur = 0 + self.block1 = nn.ModuleList([Block( + dim=embed_dims[0], num_heads=num_heads[0], mlp_ratio=mlp_ratios[0], qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, + sr_ratio=sr_ratios[0]) + for i in range(depths[0])]) + self.norm1 = norm_layer(embed_dims[0]) + + cur += depths[0] + self.block2 = nn.ModuleList([Block( + dim=embed_dims[1], num_heads=num_heads[1], mlp_ratio=mlp_ratios[1], qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, + sr_ratio=sr_ratios[1]) + for i in range(depths[1])]) + self.norm2 = norm_layer(embed_dims[1]) + + cur += depths[1] + self.block3 = nn.ModuleList([Block( + dim=embed_dims[2], num_heads=num_heads[2], mlp_ratio=mlp_ratios[2], qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, + sr_ratio=sr_ratios[2]) + for i in range(depths[2])]) + self.norm3 = norm_layer(embed_dims[2]) + + cur += depths[2] + self.block4 = nn.ModuleList([Block( + dim=embed_dims[3], num_heads=num_heads[3], mlp_ratio=mlp_ratios[3], qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, + sr_ratio=sr_ratios[3]) + for i in range(depths[3])]) + self.norm4 = norm_layer(embed_dims[3]) + + # classification head + # self.head = nn.Linear(embed_dims[3], num_classes) if num_classes > 0 else nn.Identity() + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def init_weights(self, pretrained=None): + if isinstance(pretrained, str): + logger = 1 + #load_checkpoint(self, pretrained, map_location='cpu', strict=False, logger=logger) + + def reset_drop_path(self, drop_path_rate): + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(self.depths))] + cur = 0 + for i in range(self.depths[0]): + self.block1[i].drop_path.drop_prob = dpr[cur + i] + + cur += self.depths[0] + for i in range(self.depths[1]): + self.block2[i].drop_path.drop_prob = dpr[cur + i] + + cur += self.depths[1] + for i in range(self.depths[2]): + self.block3[i].drop_path.drop_prob = dpr[cur + i] + + cur += self.depths[2] + for i in range(self.depths[3]): + self.block4[i].drop_path.drop_prob = dpr[cur + i] + + def freeze_patch_emb(self): + self.patch_embed1.requires_grad = False + + @torch.jit.ignore + def no_weight_decay(self): + return {'pos_embed1', 'pos_embed2', 'pos_embed3', 'pos_embed4', 'cls_token'} # has pos_embed may be better + + def get_classifier(self): + return self.head + + def reset_classifier(self, num_classes, global_pool=''): + self.num_classes = num_classes + self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity() + + def forward_features(self, x): + B = x.shape[0] + outs = [] + + # stage 1 + x, H, W = self.patch_embed1(x) + for i, blk in enumerate(self.block1): + x = blk(x, H, W) + x = self.norm1(x) + x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() + outs.append(x) + + # stage 2 + x, H, W = self.patch_embed2(x) + for i, blk in enumerate(self.block2): + x = blk(x, H, W) + x = self.norm2(x) + x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() + outs.append(x) + + # stage 3 + x, H, W = self.patch_embed3(x) + for i, blk in enumerate(self.block3): + x = blk(x, H, W) + x = self.norm3(x) + x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() + outs.append(x) + + # stage 4 + x, H, W = self.patch_embed4(x) + for i, blk in enumerate(self.block4): + x = blk(x, H, W) + x = self.norm4(x) + x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() + outs.append(x) + + return outs + + # return x.mean(dim=1) + + def forward(self, x): + x = self.forward_features(x) + # x = self.head(x) + + return x + + +class DWConv(nn.Module): + def __init__(self, dim=768): + super(DWConv, self).__init__() + self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim) + + def forward(self, x, H, W): + B, N, C = x.shape + x = x.transpose(1, 2).view(B, C, H, W).contiguous() + x = self.dwconv(x) + x = x.flatten(2).transpose(1, 2) + + return x + + +def _conv_filter(state_dict, patch_size=16): + """ convert patch embedding weight from manual patchify + linear proj to conv""" + out_dict = {} + for k, v in state_dict.items(): + if 'patch_embed.proj.weight' in k: + v = v.reshape((v.shape[0], 3, patch_size, patch_size)) + out_dict[k] = v + + return out_dict + + +## @register_model +class pvt_v2_b0(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b0, self).__init__( + patch_size=4, embed_dims=[32, 64, 160, 256], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[2, 2, 2, 2], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + + + +## @register_model +class pvt_v2_b1(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b1, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[2, 2, 2, 2], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + +## @register_model +class pvt_v2_b2(PyramidVisionTransformerImpr): + def __init__(self, in_channels=3, **kwargs): + super(pvt_v2_b2, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1, in_channels=in_channels) + +## @register_model +class pvt_v2_b3(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b3, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 18, 3], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + +## @register_model +class pvt_v2_b4(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b4, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 8, 27, 3], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + + +## @register_model +class pvt_v2_b5(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b5, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[4, 4, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 6, 40, 3], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + + + +### models/backbones/swin_v1.py + +# -------------------------------------------------------- +# Swin Transformer +# Copyright (c) 2021 Microsoft +# Licensed under The MIT License [see LICENSE for details] +# Written by Ze Liu, Yutong Lin, Yixuan Wei +# -------------------------------------------------------- + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as checkpoint +import numpy as np +from timm.models.layers import DropPath, to_2tuple, trunc_normal_ + +# from config import Config + + +# config = Config() + + +class Mlp(nn.Module): + """ Multilayer perceptron.""" + + def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +def window_partition(x, window_size): + """ + Args: + x: (B, H, W, C) + window_size (int): window size + + Returns: + windows: (num_windows*B, window_size, window_size, C) + """ + B, H, W, C = x.shape + x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) + windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) + return windows + + +def window_reverse(windows, window_size, H, W): + """ + Args: + windows: (num_windows*B, window_size, window_size, C) + window_size (int): Window size + H (int): Height of image + W (int): Width of image + + Returns: + x: (B, H, W, C) + """ + B = int(windows.shape[0] / (H * W / window_size / window_size)) + x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1) + x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) + return x + + +class WindowAttention(nn.Module): + """ Window based multi-head self attention (W-MSA) module with relative position bias. + It supports both of shifted and non-shifted window. + + Args: + dim (int): Number of input channels. + window_size (tuple[int]): The height and width of the window. + num_heads (int): Number of attention heads. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set + attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0 + proj_drop (float, optional): Dropout ratio of output. Default: 0.0 + """ + + def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.): + + super().__init__() + self.dim = dim + self.window_size = window_size # Wh, Ww + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = qk_scale or head_dim ** -0.5 + + # define a parameter table of relative position bias + self.relative_position_bias_table = nn.Parameter( + torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)) # 2*Wh-1 * 2*Ww-1, nH + + # get pair-wise relative position index for each token inside the window + coords_h = torch.arange(self.window_size[0]) + coords_w = torch.arange(self.window_size[1]) + coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing='ij')) # 2, Wh, Ww + coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww + relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww + relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2 + relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0 + relative_coords[:, :, 1] += self.window_size[1] - 1 + relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1 + relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww + self.register_buffer("relative_position_index", relative_position_index) + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop_prob = attn_drop + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + trunc_normal_(self.relative_position_bias_table, std=.02) + self.softmax = nn.Softmax(dim=-1) + + def forward(self, x, mask=None): + """ Forward function. + + Args: + x: input features with shape of (num_windows*B, N, C) + mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None + """ + B_, N, C = x.shape + qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple) + + q = q * self.scale + + if config.SDPA_enabled: + x = torch.nn.functional.scaled_dot_product_attention( + q, k, v, + attn_mask=None, dropout_p=self.attn_drop_prob, is_causal=False + ).transpose(1, 2).reshape(B_, N, C) + else: + attn = (q @ k.transpose(-2, -1)) + + relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view( + self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1 + ) # Wh*Ww, Wh*Ww, nH + relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww + attn = attn + relative_position_bias.unsqueeze(0) + + if mask is not None: + nW = mask.shape[0] + attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0) + attn = attn.view(-1, self.num_heads, N, N) + attn = self.softmax(attn) + else: + attn = self.softmax(attn) + + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B_, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class SwinTransformerBlock(nn.Module): + """ Swin Transformer Block. + + Args: + dim (int): Number of input channels. + num_heads (int): Number of attention heads. + window_size (int): Window size. + shift_size (int): Shift size for SW-MSA. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float, optional): Stochastic depth rate. Default: 0.0 + act_layer (nn.Module, optional): Activation layer. Default: nn.GELU + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + """ + + def __init__(self, dim, num_heads, window_size=7, shift_size=0, + mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0., + act_layer=nn.GELU, norm_layer=nn.LayerNorm): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.window_size = window_size + self.shift_size = shift_size + self.mlp_ratio = mlp_ratio + assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size" + + self.norm1 = norm_layer(dim) + self.attn = WindowAttention( + dim, window_size=to_2tuple(self.window_size), num_heads=num_heads, + qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop) + + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + + self.H = None + self.W = None + + def forward(self, x, mask_matrix): + """ Forward function. + + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + mask_matrix: Attention mask for cyclic shift. + """ + B, L, C = x.shape + H, W = self.H, self.W + assert L == H * W, "input feature has wrong size" + + shortcut = x + x = self.norm1(x) + x = x.view(B, H, W, C) + + # pad feature maps to multiples of window size + pad_l = pad_t = 0 + pad_r = (self.window_size - W % self.window_size) % self.window_size + pad_b = (self.window_size - H % self.window_size) % self.window_size + x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b)) + _, Hp, Wp, _ = x.shape + + # cyclic shift + if self.shift_size > 0: + shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)) + attn_mask = mask_matrix + else: + shifted_x = x + attn_mask = None + + # partition windows + x_windows = window_partition(shifted_x, self.window_size) # nW*B, window_size, window_size, C + x_windows = x_windows.view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C + + # W-MSA/SW-MSA + attn_windows = self.attn(x_windows, mask=attn_mask) # nW*B, window_size*window_size, C + + # merge windows + attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C) + shifted_x = window_reverse(attn_windows, self.window_size, Hp, Wp) # B H' W' C + + # reverse cyclic shift + if self.shift_size > 0: + x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2)) + else: + x = shifted_x + + if pad_r > 0 or pad_b > 0: + x = x[:, :H, :W, :].contiguous() + + x = x.view(B, H * W, C) + + # FFN + x = shortcut + self.drop_path(x) + x = x + self.drop_path(self.mlp(self.norm2(x))) + + return x + + +class PatchMerging(nn.Module): + """ Patch Merging Layer + + Args: + dim (int): Number of input channels. + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + """ + def __init__(self, dim, norm_layer=nn.LayerNorm): + super().__init__() + self.dim = dim + self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False) + self.norm = norm_layer(4 * dim) + + def forward(self, x, H, W): + """ Forward function. + + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + """ + B, L, C = x.shape + assert L == H * W, "input feature has wrong size" + + x = x.view(B, H, W, C) + + # padding + pad_input = (H % 2 == 1) or (W % 2 == 1) + if pad_input: + x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2)) + + x0 = x[:, 0::2, 0::2, :] # B H/2 W/2 C + x1 = x[:, 1::2, 0::2, :] # B H/2 W/2 C + x2 = x[:, 0::2, 1::2, :] # B H/2 W/2 C + x3 = x[:, 1::2, 1::2, :] # B H/2 W/2 C + x = torch.cat([x0, x1, x2, x3], -1) # B H/2 W/2 4*C + x = x.view(B, -1, 4 * C) # B H/2*W/2 4*C + + x = self.norm(x) + x = self.reduction(x) + + return x + + +class BasicLayer(nn.Module): + """ A basic Swin Transformer layer for one stage. + + Args: + dim (int): Number of feature channels + depth (int): Depths of this stage. + num_heads (int): Number of attention head. + window_size (int): Local window size. Default: 7. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. + """ + + def __init__(self, + dim, + depth, + num_heads, + window_size=7, + mlp_ratio=4., + qkv_bias=True, + qk_scale=None, + drop=0., + attn_drop=0., + drop_path=0., + norm_layer=nn.LayerNorm, + downsample=None, + use_checkpoint=False): + super().__init__() + self.window_size = window_size + self.shift_size = window_size // 2 + self.depth = depth + self.use_checkpoint = use_checkpoint + + # build blocks + self.blocks = nn.ModuleList([ + SwinTransformerBlock( + dim=dim, + num_heads=num_heads, + window_size=window_size, + shift_size=0 if (i % 2 == 0) else window_size // 2, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop, + attn_drop=attn_drop, + drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path, + norm_layer=norm_layer) + for i in range(depth)]) + + # patch merging layer + if downsample is not None: + self.downsample = downsample(dim=dim, norm_layer=norm_layer) + else: + self.downsample = None + + def forward(self, x, H, W): + """ Forward function. + + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + """ + + # calculate attention mask for SW-MSA + # Turn int to torch.tensor for the compatiability with torch.compile in PyTorch 2.5. + Hp = torch.ceil(torch.tensor(H) / self.window_size).to(torch.int64) * self.window_size + Wp = torch.ceil(torch.tensor(W) / self.window_size).to(torch.int64) * self.window_size + img_mask = torch.zeros((1, Hp, Wp, 1), device=x.device) # 1 Hp Wp 1 + h_slices = (slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None)) + w_slices = (slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None)) + cnt = 0 + for h in h_slices: + for w in w_slices: + img_mask[:, h, w, :] = cnt + cnt += 1 + + mask_windows = window_partition(img_mask, self.window_size) # nW, window_size, window_size, 1 + mask_windows = mask_windows.view(-1, self.window_size * self.window_size) + attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) + attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0)).to(x.dtype) + + for blk in self.blocks: + blk.H, blk.W = H, W + if self.use_checkpoint: + x = checkpoint.checkpoint(blk, x, attn_mask) + else: + x = blk(x, attn_mask) + if self.downsample is not None: + x_down = self.downsample(x, H, W) + Wh, Ww = (H + 1) // 2, (W + 1) // 2 + return x, H, W, x_down, Wh, Ww + else: + return x, H, W, x, H, W + + +class PatchEmbed(nn.Module): + """ Image to Patch Embedding + + Args: + patch_size (int): Patch token size. Default: 4. + in_channels (int): Number of input image channels. Default: 3. + embed_dim (int): Number of linear projection output channels. Default: 96. + norm_layer (nn.Module, optional): Normalization layer. Default: None + """ + + def __init__(self, patch_size=4, in_channels=3, embed_dim=96, norm_layer=None): + super().__init__() + patch_size = to_2tuple(patch_size) + self.patch_size = patch_size + + self.in_channels = in_channels + self.embed_dim = embed_dim + + self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size) + if norm_layer is not None: + self.norm = norm_layer(embed_dim) + else: + self.norm = None + + def forward(self, x): + """Forward function.""" + # padding + _, _, H, W = x.size() + if W % self.patch_size[1] != 0: + x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1])) + if H % self.patch_size[0] != 0: + x = F.pad(x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0])) + + x = self.proj(x) # B C Wh Ww + if self.norm is not None: + Wh, Ww = x.size(2), x.size(3) + x = x.flatten(2).transpose(1, 2) + x = self.norm(x) + x = x.transpose(1, 2).view(-1, self.embed_dim, Wh, Ww) + + return x + + +class SwinTransformer(nn.Module): + """ Swin Transformer backbone. + A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows` - + https://arxiv.org/pdf/2103.14030 + + Args: + pretrain_img_size (int): Input image size for training the pretrained model, + used in absolute postion embedding. Default 224. + patch_size (int | tuple(int)): Patch size. Default: 4. + in_channels (int): Number of input image channels. Default: 3. + embed_dim (int): Number of linear projection output channels. Default: 96. + depths (tuple[int]): Depths of each Swin Transformer stage. + num_heads (tuple[int]): Number of attention head of each stage. + window_size (int): Window size. Default: 7. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4. + qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. + drop_rate (float): Dropout rate. + attn_drop_rate (float): Attention dropout rate. Default: 0. + drop_path_rate (float): Stochastic depth rate. Default: 0.2. + norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm. + ape (bool): If True, add absolute position embedding to the patch embedding. Default: False. + patch_norm (bool): If True, add normalization after patch embedding. Default: True. + out_indices (Sequence[int]): Output from which stages. + frozen_stages (int): Stages to be frozen (stop grad and set eval mode). + -1 means not freezing any parameters. + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. + """ + + def __init__(self, + pretrain_img_size=224, + patch_size=4, + in_channels=3, + embed_dim=96, + depths=[2, 2, 6, 2], + num_heads=[3, 6, 12, 24], + window_size=7, + mlp_ratio=4., + qkv_bias=True, + qk_scale=None, + drop_rate=0., + attn_drop_rate=0., + drop_path_rate=0.2, + norm_layer=nn.LayerNorm, + ape=False, + patch_norm=True, + out_indices=(0, 1, 2, 3), + frozen_stages=-1, + use_checkpoint=False): + super().__init__() + + self.pretrain_img_size = pretrain_img_size + self.num_layers = len(depths) + self.embed_dim = embed_dim + self.ape = ape + self.patch_norm = patch_norm + self.out_indices = out_indices + self.frozen_stages = frozen_stages + + # split image into non-overlapping patches + self.patch_embed = PatchEmbed( + patch_size=patch_size, in_channels=in_channels, embed_dim=embed_dim, + norm_layer=norm_layer if self.patch_norm else None) + + # absolute position embedding + if self.ape: + pretrain_img_size = to_2tuple(pretrain_img_size) + patch_size = to_2tuple(patch_size) + patches_resolution = [pretrain_img_size[0] // patch_size[0], pretrain_img_size[1] // patch_size[1]] + + self.absolute_pos_embed = nn.Parameter(torch.zeros(1, embed_dim, patches_resolution[0], patches_resolution[1])) + trunc_normal_(self.absolute_pos_embed, std=.02) + + self.pos_drop = nn.Dropout(p=drop_rate) + + # stochastic depth + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule + + # build layers + self.layers = nn.ModuleList() + for i_layer in range(self.num_layers): + layer = BasicLayer( + dim=int(embed_dim * 2 ** i_layer), + depth=depths[i_layer], + num_heads=num_heads[i_layer], + window_size=window_size, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])], + norm_layer=norm_layer, + downsample=PatchMerging if (i_layer < self.num_layers - 1) else None, + use_checkpoint=use_checkpoint) + self.layers.append(layer) + + num_features = [int(embed_dim * 2 ** i) for i in range(self.num_layers)] + self.num_features = num_features + + # add a norm layer for each output + for i_layer in out_indices: + layer = norm_layer(num_features[i_layer]) + layer_name = f'norm{i_layer}' + self.add_module(layer_name, layer) + + self._freeze_stages() + + def _freeze_stages(self): + if self.frozen_stages >= 0: + self.patch_embed.eval() + for param in self.patch_embed.parameters(): + param.requires_grad = False + + if self.frozen_stages >= 1 and self.ape: + self.absolute_pos_embed.requires_grad = False + + if self.frozen_stages >= 2: + self.pos_drop.eval() + for i in range(0, self.frozen_stages - 1): + m = self.layers[i] + m.eval() + for param in m.parameters(): + param.requires_grad = False + + + def forward(self, x): + """Forward function.""" + x = self.patch_embed(x) + + Wh, Ww = x.size(2), x.size(3) + if self.ape: + # interpolate the position embedding to the corresponding size + absolute_pos_embed = F.interpolate(self.absolute_pos_embed, size=(Wh, Ww), mode='bicubic') + x = (x + absolute_pos_embed) # B Wh*Ww C + + outs = []#x.contiguous()] + x = x.flatten(2).transpose(1, 2) + x = self.pos_drop(x) + for i in range(self.num_layers): + layer = self.layers[i] + x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww) + + if i in self.out_indices: + norm_layer = getattr(self, f'norm{i}') + x_out = norm_layer(x_out) + + out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous() + outs.append(out) + + return tuple(outs) + + def train(self, mode=True): + """Convert the model into training mode while keep layers freezed.""" + super(SwinTransformer, self).train(mode) + self._freeze_stages() + +def swin_v1_t(): + model = SwinTransformer(embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24], window_size=7) + return model + +def swin_v1_s(): + model = SwinTransformer(embed_dim=96, depths=[2, 2, 18, 2], num_heads=[3, 6, 12, 24], window_size=7) + return model + +def swin_v1_b(): + model = SwinTransformer(embed_dim=128, depths=[2, 2, 18, 2], num_heads=[4, 8, 16, 32], window_size=12) + return model + +def swin_v1_l(): + model = SwinTransformer(embed_dim=192, depths=[2, 2, 18, 2], num_heads=[6, 12, 24, 48], window_size=12) + return model + + + +### models/modules/deform_conv.py + +import torch +import torch.nn as nn +from torchvision.ops import deform_conv2d + + +class DeformableConv2d(nn.Module): + def __init__(self, + in_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1, + bias=False): + + super(DeformableConv2d, self).__init__() + + assert type(kernel_size) == tuple or type(kernel_size) == int + + kernel_size = kernel_size if type(kernel_size) == tuple else (kernel_size, kernel_size) + self.stride = stride if type(stride) == tuple else (stride, stride) + self.padding = padding + + self.offset_conv = nn.Conv2d(in_channels, + 2 * kernel_size[0] * kernel_size[1], + kernel_size=kernel_size, + stride=stride, + padding=self.padding, + bias=True) + + nn.init.constant_(self.offset_conv.weight, 0.) + nn.init.constant_(self.offset_conv.bias, 0.) + + self.modulator_conv = nn.Conv2d(in_channels, + 1 * kernel_size[0] * kernel_size[1], + kernel_size=kernel_size, + stride=stride, + padding=self.padding, + bias=True) + + nn.init.constant_(self.modulator_conv.weight, 0.) + nn.init.constant_(self.modulator_conv.bias, 0.) + + self.regular_conv = nn.Conv2d(in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=self.padding, + bias=bias) + + def forward(self, x): + #h, w = x.shape[2:] + #max_offset = max(h, w)/4. + + offset = self.offset_conv(x)#.clamp(-max_offset, max_offset) + modulator = 2. * torch.sigmoid(self.modulator_conv(x)) + + x = deform_conv2d( + input=x, + offset=offset, + weight=self.regular_conv.weight, + bias=self.regular_conv.bias, + padding=self.padding, + mask=modulator, + stride=self.stride, + ) + return x + + + + +### utils.py + +import torch.nn as nn + + +def build_act_layer(act_layer): + if act_layer == 'ReLU': + return nn.ReLU(inplace=True) + elif act_layer == 'SiLU': + return nn.SiLU(inplace=True) + elif act_layer == 'GELU': + return nn.GELU() + + raise NotImplementedError(f'build_act_layer does not support {act_layer}') + + +def build_norm_layer(dim, + norm_layer, + in_format='channels_last', + out_format='channels_last', + eps=1e-6): + layers = [] + if norm_layer == 'BN': + if in_format == 'channels_last': + layers.append(to_channels_first()) + layers.append(nn.BatchNorm2d(dim)) + if out_format == 'channels_last': + layers.append(to_channels_last()) + elif norm_layer == 'LN': + if in_format == 'channels_first': + layers.append(to_channels_last()) + layers.append(nn.LayerNorm(dim, eps=eps)) + if out_format == 'channels_first': + layers.append(to_channels_first()) + else: + raise NotImplementedError( + f'build_norm_layer does not support {norm_layer}') + return nn.Sequential(*layers) + + +class to_channels_first(nn.Module): + + def __init__(self): + super().__init__() + + def forward(self, x): + return x.permute(0, 3, 1, 2) + + +class to_channels_last(nn.Module): + + def __init__(self): + super().__init__() + + def forward(self, x): + return x.permute(0, 2, 3, 1) + + + +### dataset.py + +_class_labels_TR_sorted = ( + 'Airplane, Ant, Antenna, Archery, Axe, BabyCarriage, Bag, BalanceBeam, Balcony, Balloon, Basket, BasketballHoop, Beatle, Bed, Bee, Bench, Bicycle, ' + 'BicycleFrame, BicycleStand, Boat, Bonsai, BoomLift, Bridge, BunkBed, Butterfly, Button, Cable, CableLift, Cage, Camcorder, Cannon, Canoe, Car, ' + 'CarParkDropArm, Carriage, Cart, Caterpillar, CeilingLamp, Centipede, Chair, Clip, Clock, Clothes, CoatHanger, Comb, ConcretePumpTruck, Crack, Crane, ' + 'Cup, DentalChair, Desk, DeskChair, Diagram, DishRack, DoorHandle, Dragonfish, Dragonfly, Drum, Earphone, Easel, ElectricIron, Excavator, Eyeglasses, ' + 'Fan, Fence, Fencing, FerrisWheel, FireExtinguisher, Fishing, Flag, FloorLamp, Forklift, GasStation, Gate, Gear, Goal, Golf, GymEquipment, Hammock, ' + 'Handcart, Handcraft, Handrail, HangGlider, Harp, Harvester, Headset, Helicopter, Helmet, Hook, HorizontalBar, Hydrovalve, IroningTable, Jewelry, Key, ' + 'KidsPlayground, Kitchenware, Kite, Knife, Ladder, LaundryRack, Lightning, Lobster, Locust, Machine, MachineGun, MagazineRack, Mantis, Medal, MemorialArchway, ' + 'Microphone, Missile, MobileHolder, Monitor, Mosquito, Motorcycle, MovingTrolley, Mower, MusicPlayer, MusicStand, ObservationTower, Octopus, OilWell, ' + 'OlympicLogo, OperatingTable, OutdoorFitnessEquipment, Parachute, Pavilion, Piano, Pipe, PlowHarrow, PoleVault, Punchbag, Rack, Racket, Rifle, Ring, Robot, ' + 'RockClimbing, Rope, Sailboat, Satellite, Scaffold, Scale, Scissor, Scooter, Sculpture, Seadragon, Seahorse, Seal, SewingMachine, Ship, Shoe, ShoppingCart, ' + 'ShoppingTrolley, Shower, Shrimp, Signboard, Skateboarding, Skeleton, Skiing, Spade, SpeedBoat, Spider, Spoon, Stair, Stand, Stationary, SteeringWheel, ' + 'Stethoscope, Stool, Stove, StreetLamp, SweetStand, Swing, Sword, TV, Table, TableChair, TableLamp, TableTennis, Tank, Tapeline, Teapot, Telescope, Tent, ' + 'TobaccoPipe, Toy, Tractor, TrafficLight, TrafficSign, Trampoline, TransmissionTower, Tree, Tricycle, TrimmerCover, Tripod, Trombone, Truck, Trumpet, Tuba, ' + 'UAV, Umbrella, UnevenBars, UtilityPole, VacuumCleaner, Violin, Wakesurfing, Watch, WaterTower, WateringPot, Well, WellLid, Wheel, Wheelchair, WindTurbine, Windmill, WineGlass, WireWhisk, Yacht' +) +class_labels_TR_sorted = _class_labels_TR_sorted.split(', ') + + +### models/backbones/build_backbones.py + +import torch +import torch.nn as nn +from collections import OrderedDict +from torchvision.models import vgg16, vgg16_bn, VGG16_Weights, VGG16_BN_Weights, resnet50, ResNet50_Weights +# from models.pvt_v2 import pvt_v2_b0, pvt_v2_b1, pvt_v2_b2, pvt_v2_b5 +# from models.swin_v1 import swin_v1_t, swin_v1_s, swin_v1_b, swin_v1_l +# from config import Config + + +config = Config() + +def build_backbone(bb_name, pretrained=True, params_settings=''): + if bb_name == 'vgg16': + bb_net = list(vgg16(pretrained=VGG16_Weights.DEFAULT if pretrained else None).children())[0] + bb = nn.Sequential(OrderedDict({'conv1': bb_net[:4], 'conv2': bb_net[4:9], 'conv3': bb_net[9:16], 'conv4': bb_net[16:23]})) + elif bb_name == 'vgg16bn': + bb_net = list(vgg16_bn(pretrained=VGG16_BN_Weights.DEFAULT if pretrained else None).children())[0] + bb = nn.Sequential(OrderedDict({'conv1': bb_net[:6], 'conv2': bb_net[6:13], 'conv3': bb_net[13:23], 'conv4': bb_net[23:33]})) + elif bb_name == 'resnet50': + bb_net = list(resnet50(pretrained=ResNet50_Weights.DEFAULT if pretrained else None).children()) + bb = nn.Sequential(OrderedDict({'conv1': nn.Sequential(*bb_net[0:3]), 'conv2': bb_net[4], 'conv3': bb_net[5], 'conv4': bb_net[6]})) + else: + bb = eval('{}({})'.format(bb_name, params_settings)) + if pretrained: + bb = load_weights(bb, bb_name) + return bb + +def load_weights(model, model_name): + save_model = torch.load(config.weights[model_name], map_location='cpu') + model_dict = model.state_dict() + state_dict = {k: v if v.size() == model_dict[k].size() else model_dict[k] for k, v in save_model.items() if k in model_dict.keys()} + # to ignore the weights with mismatched size when I modify the backbone itself. + if not state_dict: + save_model_keys = list(save_model.keys()) + sub_item = save_model_keys[0] if len(save_model_keys) == 1 else None + state_dict = {k: v if v.size() == model_dict[k].size() else model_dict[k] for k, v in save_model[sub_item].items() if k in model_dict.keys()} + if not state_dict or not sub_item: + print('Weights are not successully loaded. Check the state dict of weights file.') + return None + else: + print('Found correct weights in the "{}" item of loaded state_dict.'.format(sub_item)) + model_dict.update(state_dict) + model.load_state_dict(model_dict) + return model + + + +### models/modules/decoder_blocks.py + +import torch +import torch.nn as nn +# from models.aspp import ASPP, ASPPDeformable +# from config import Config + + +# config = Config() + + +class BasicDecBlk(nn.Module): + def __init__(self, in_channels=64, out_channels=64, inter_channels=64): + super(BasicDecBlk, self).__init__() + inter_channels = in_channels // 4 if config.dec_channels_inter == 'adap' else 64 + self.conv_in = nn.Conv2d(in_channels, inter_channels, 3, 1, padding=1) + self.relu_in = nn.ReLU(inplace=True) + if config.dec_att == 'ASPP': + self.dec_att = ASPP(in_channels=inter_channels) + elif config.dec_att == 'ASPPDeformable': + self.dec_att = ASPPDeformable(in_channels=inter_channels) + self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, padding=1) + self.bn_in = nn.BatchNorm2d(inter_channels) if config.batch_size > 1 else nn.Identity() + self.bn_out = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity() + + def forward(self, x): + x = self.conv_in(x) + x = self.bn_in(x) + x = self.relu_in(x) + if hasattr(self, 'dec_att'): + x = self.dec_att(x) + x = self.conv_out(x) + x = self.bn_out(x) + return x + + +class ResBlk(nn.Module): + def __init__(self, in_channels=64, out_channels=None, inter_channels=64): + super(ResBlk, self).__init__() + if out_channels is None: + out_channels = in_channels + inter_channels = in_channels // 4 if config.dec_channels_inter == 'adap' else 64 + + self.conv_in = nn.Conv2d(in_channels, inter_channels, 3, 1, padding=1) + self.bn_in = nn.BatchNorm2d(inter_channels) if config.batch_size > 1 else nn.Identity() + self.relu_in = nn.ReLU(inplace=True) + + if config.dec_att == 'ASPP': + self.dec_att = ASPP(in_channels=inter_channels) + elif config.dec_att == 'ASPPDeformable': + self.dec_att = ASPPDeformable(in_channels=inter_channels) + + self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, padding=1) + self.bn_out = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity() + + self.conv_resi = nn.Conv2d(in_channels, out_channels, 1, 1, 0) + + def forward(self, x): + _x = self.conv_resi(x) + x = self.conv_in(x) + x = self.bn_in(x) + x = self.relu_in(x) + if hasattr(self, 'dec_att'): + x = self.dec_att(x) + x = self.conv_out(x) + x = self.bn_out(x) + return x + _x + + + +### models/modules/lateral_blocks.py + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from functools import partial + +# from config import Config + + +# config = Config() + + +class BasicLatBlk(nn.Module): + def __init__(self, in_channels=64, out_channels=64, inter_channels=64): + super(BasicLatBlk, self).__init__() + inter_channels = in_channels // 4 if config.dec_channels_inter == 'adap' else 64 + self.conv = nn.Conv2d(in_channels, out_channels, 1, 1, 0) + + def forward(self, x): + x = self.conv(x) + return x + + + +### models/modules/aspp.py + +import torch +import torch.nn as nn +import torch.nn.functional as F +# from models.deform_conv import DeformableConv2d +# from config import Config + + +# config = Config() + + +class _ASPPModule(nn.Module): + def __init__(self, in_channels, planes, kernel_size, padding, dilation): + super(_ASPPModule, self).__init__() + self.atrous_conv = nn.Conv2d(in_channels, planes, kernel_size=kernel_size, + stride=1, padding=padding, dilation=dilation, bias=False) + self.bn = nn.BatchNorm2d(planes) if config.batch_size > 1 else nn.Identity() + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + x = self.atrous_conv(x) + x = self.bn(x) + + return self.relu(x) + + +class ASPP(nn.Module): + def __init__(self, in_channels=64, out_channels=None, output_stride=16): + super(ASPP, self).__init__() + self.down_scale = 1 + if out_channels is None: + out_channels = in_channels + self.in_channelster = 256 // self.down_scale + if output_stride == 16: + dilations = [1, 6, 12, 18] + elif output_stride == 8: + dilations = [1, 12, 24, 36] + else: + raise NotImplementedError + + self.aspp1 = _ASPPModule(in_channels, self.in_channelster, 1, padding=0, dilation=dilations[0]) + self.aspp2 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[1], dilation=dilations[1]) + self.aspp3 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[2], dilation=dilations[2]) + self.aspp4 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[3], dilation=dilations[3]) + + self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)), + nn.Conv2d(in_channels, self.in_channelster, 1, stride=1, bias=False), + nn.BatchNorm2d(self.in_channelster) if config.batch_size > 1 else nn.Identity(), + nn.ReLU(inplace=True)) + self.conv1 = nn.Conv2d(self.in_channelster * 5, out_channels, 1, bias=False) + self.bn1 = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity() + self.relu = nn.ReLU(inplace=True) + self.dropout = nn.Dropout(0.5) + + def forward(self, x): + x1 = self.aspp1(x) + x2 = self.aspp2(x) + x3 = self.aspp3(x) + x4 = self.aspp4(x) + x5 = self.global_avg_pool(x) + x5 = F.interpolate(x5, size=x1.size()[2:], mode='bilinear', align_corners=True) + x = torch.cat((x1, x2, x3, x4, x5), dim=1) + + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + + return self.dropout(x) + + +##################### Deformable +class _ASPPModuleDeformable(nn.Module): + def __init__(self, in_channels, planes, kernel_size, padding): + super(_ASPPModuleDeformable, self).__init__() + self.atrous_conv = DeformableConv2d(in_channels, planes, kernel_size=kernel_size, + stride=1, padding=padding, bias=False) + self.bn = nn.BatchNorm2d(planes) if config.batch_size > 1 else nn.Identity() + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + x = self.atrous_conv(x) + x = self.bn(x) + + return self.relu(x) + + +class ASPPDeformable(nn.Module): + def __init__(self, in_channels, out_channels=None, parallel_block_sizes=[1, 3, 7]): + super(ASPPDeformable, self).__init__() + self.down_scale = 1 + if out_channels is None: + out_channels = in_channels + self.in_channelster = 256 // self.down_scale + + self.aspp1 = _ASPPModuleDeformable(in_channels, self.in_channelster, 1, padding=0) + self.aspp_deforms = nn.ModuleList([ + _ASPPModuleDeformable(in_channels, self.in_channelster, conv_size, padding=int(conv_size//2)) for conv_size in parallel_block_sizes + ]) + + self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)), + nn.Conv2d(in_channels, self.in_channelster, 1, stride=1, bias=False), + nn.BatchNorm2d(self.in_channelster) if config.batch_size > 1 else nn.Identity(), + nn.ReLU(inplace=True)) + self.conv1 = nn.Conv2d(self.in_channelster * (2 + len(self.aspp_deforms)), out_channels, 1, bias=False) + self.bn1 = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity() + self.relu = nn.ReLU(inplace=True) + self.dropout = nn.Dropout(0.5) + + def forward(self, x): + x1 = self.aspp1(x) + x_aspp_deforms = [aspp_deform(x) for aspp_deform in self.aspp_deforms] + x5 = self.global_avg_pool(x) + x5 = F.interpolate(x5, size=x1.size()[2:], mode='bilinear', align_corners=True) + x = torch.cat((x1, *x_aspp_deforms, x5), dim=1) + + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + + return self.dropout(x) + + + +### models/refinement/refiner.py + +import torch +import torch.nn as nn +from collections import OrderedDict +import torch +import torch.nn as nn +import torch.nn.functional as F +from torchvision.models import vgg16, vgg16_bn +from torchvision.models import resnet50 + +# from config import Config +# from dataset import class_labels_TR_sorted +# from models.build_backbone import build_backbone +# from models.decoder_blocks import BasicDecBlk +# from models.lateral_blocks import BasicLatBlk +# from models.ing import * +# from models.stem_layer import StemLayer + + +class RefinerPVTInChannels4(nn.Module): + def __init__(self, in_channels=3+1): + super(RefinerPVTInChannels4, self).__init__() + self.config = Config() + self.epoch = 1 + self.bb = build_backbone(self.config.bb, params_settings='in_channels=4') + + lateral_channels_in_collection = { + 'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64], + 'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64], + 'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192], + } + channels = lateral_channels_in_collection[self.config.bb] + self.squeeze_module = BasicDecBlk(channels[0], channels[0]) + + self.decoder = Decoder(channels) + + if 0: + for key, value in self.named_parameters(): + if 'bb.' in key: + value.requires_grad = False + + def forward(self, x): + if isinstance(x, list): + x = torch.cat(x, dim=1) + ########## Encoder ########## + if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']: + x1 = self.bb.conv1(x) + x2 = self.bb.conv2(x1) + x3 = self.bb.conv3(x2) + x4 = self.bb.conv4(x3) + else: + x1, x2, x3, x4 = self.bb(x) + + x4 = self.squeeze_module(x4) + + ########## Decoder ########## + + features = [x, x1, x2, x3, x4] + scaled_preds = self.decoder(features) + + return scaled_preds + + +class Refiner(nn.Module): + def __init__(self, in_channels=3+1): + super(Refiner, self).__init__() + self.config = Config() + self.epoch = 1 + self.stem_layer = StemLayer(in_channels=in_channels, inter_channels=48, out_channels=3, norm_layer='BN' if self.config.batch_size > 1 else 'LN') + self.bb = build_backbone(self.config.bb) + + lateral_channels_in_collection = { + 'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64], + 'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64], + 'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192], + } + channels = lateral_channels_in_collection[self.config.bb] + self.squeeze_module = BasicDecBlk(channels[0], channels[0]) + + self.decoder = Decoder(channels) + + if 0: + for key, value in self.named_parameters(): + if 'bb.' in key: + value.requires_grad = False + + def forward(self, x): + if isinstance(x, list): + x = torch.cat(x, dim=1) + x = self.stem_layer(x) + ########## Encoder ########## + if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']: + x1 = self.bb.conv1(x) + x2 = self.bb.conv2(x1) + x3 = self.bb.conv3(x2) + x4 = self.bb.conv4(x3) + else: + x1, x2, x3, x4 = self.bb(x) + + x4 = self.squeeze_module(x4) + + ########## Decoder ########## + + features = [x, x1, x2, x3, x4] + scaled_preds = self.decoder(features) + + return scaled_preds + + +class Decoder(nn.Module): + def __init__(self, channels): + super(Decoder, self).__init__() + self.config = Config() + DecoderBlock = eval('BasicDecBlk') + LateralBlock = eval('BasicLatBlk') + + self.decoder_block4 = DecoderBlock(channels[0], channels[1]) + self.decoder_block3 = DecoderBlock(channels[1], channels[2]) + self.decoder_block2 = DecoderBlock(channels[2], channels[3]) + self.decoder_block1 = DecoderBlock(channels[3], channels[3]//2) + + self.lateral_block4 = LateralBlock(channels[1], channels[1]) + self.lateral_block3 = LateralBlock(channels[2], channels[2]) + self.lateral_block2 = LateralBlock(channels[3], channels[3]) + + if self.config.ms_supervision: + self.conv_ms_spvn_4 = nn.Conv2d(channels[1], 1, 1, 1, 0) + self.conv_ms_spvn_3 = nn.Conv2d(channels[2], 1, 1, 1, 0) + self.conv_ms_spvn_2 = nn.Conv2d(channels[3], 1, 1, 1, 0) + self.conv_out1 = nn.Sequential(nn.Conv2d(channels[3]//2, 1, 1, 1, 0)) + + def forward(self, features): + x, x1, x2, x3, x4 = features + outs = [] + p4 = self.decoder_block4(x4) + _p4 = F.interpolate(p4, size=x3.shape[2:], mode='bilinear', align_corners=True) + _p3 = _p4 + self.lateral_block4(x3) + + p3 = self.decoder_block3(_p3) + _p3 = F.interpolate(p3, size=x2.shape[2:], mode='bilinear', align_corners=True) + _p2 = _p3 + self.lateral_block3(x2) + + p2 = self.decoder_block2(_p2) + _p2 = F.interpolate(p2, size=x1.shape[2:], mode='bilinear', align_corners=True) + _p1 = _p2 + self.lateral_block2(x1) + + _p1 = self.decoder_block1(_p1) + _p1 = F.interpolate(_p1, size=x.shape[2:], mode='bilinear', align_corners=True) + p1_out = self.conv_out1(_p1) + + if self.config.ms_supervision: + outs.append(self.conv_ms_spvn_4(p4)) + outs.append(self.conv_ms_spvn_3(p3)) + outs.append(self.conv_ms_spvn_2(p2)) + outs.append(p1_out) + return outs + + +class RefUNet(nn.Module): + # Refinement + def __init__(self, in_channels=3+1): + super(RefUNet, self).__init__() + self.encoder_1 = nn.Sequential( + nn.Conv2d(in_channels, 64, 3, 1, 1), + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.encoder_2 = nn.Sequential( + nn.MaxPool2d(2, 2, ceil_mode=True), + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.encoder_3 = nn.Sequential( + nn.MaxPool2d(2, 2, ceil_mode=True), + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.encoder_4 = nn.Sequential( + nn.MaxPool2d(2, 2, ceil_mode=True), + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.pool4 = nn.MaxPool2d(2, 2, ceil_mode=True) + ##### + self.decoder_5 = nn.Sequential( + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + ##### + self.decoder_4 = nn.Sequential( + nn.Conv2d(128, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.decoder_3 = nn.Sequential( + nn.Conv2d(128, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.decoder_2 = nn.Sequential( + nn.Conv2d(128, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.decoder_1 = nn.Sequential( + nn.Conv2d(128, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.conv_d0 = nn.Conv2d(64, 1, 3, 1, 1) + + self.upscore2 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) + + def forward(self, x): + outs = [] + if isinstance(x, list): + x = torch.cat(x, dim=1) + hx = x + + hx1 = self.encoder_1(hx) + hx2 = self.encoder_2(hx1) + hx3 = self.encoder_3(hx2) + hx4 = self.encoder_4(hx3) + + hx = self.decoder_5(self.pool4(hx4)) + hx = torch.cat((self.upscore2(hx), hx4), 1) + + d4 = self.decoder_4(hx) + hx = torch.cat((self.upscore2(d4), hx3), 1) + + d3 = self.decoder_3(hx) + hx = torch.cat((self.upscore2(d3), hx2), 1) + + d2 = self.decoder_2(hx) + hx = torch.cat((self.upscore2(d2), hx1), 1) + + d1 = self.decoder_1(hx) + + x = self.conv_d0(d1) + outs.append(x) + return outs + + + +### models/stem_layer.py + +import torch.nn as nn +# from utils import build_act_layer, build_norm_layer + + +class StemLayer(nn.Module): + r""" Stem layer of InternImage + Args: + in_channels (int): number of input channels + out_channels (int): number of output channels + act_layer (str): activation layer + norm_layer (str): normalization layer + """ + + def __init__(self, + in_channels=3+1, + inter_channels=48, + out_channels=96, + act_layer='GELU', + norm_layer='BN'): + super().__init__() + self.conv1 = nn.Conv2d(in_channels, + inter_channels, + kernel_size=3, + stride=1, + padding=1) + self.norm1 = build_norm_layer( + inter_channels, norm_layer, 'channels_first', 'channels_first' + ) + self.act = build_act_layer(act_layer) + self.conv2 = nn.Conv2d(inter_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1) + self.norm2 = build_norm_layer( + out_channels, norm_layer, 'channels_first', 'channels_first' + ) + + def forward(self, x): + x = self.conv1(x) + x = self.norm1(x) + x = self.act(x) + x = self.conv2(x) + x = self.norm2(x) + return x + + +### models/birefnet.py + +import torch +import torch.nn as nn +import torch.nn.functional as F +from kornia.filters import laplacian +from transformers import PreTrainedModel +from einops import rearrange + +# from config import Config +# from dataset import class_labels_TR_sorted +# from models.build_backbone import build_backbone +# from models.decoder_blocks import BasicDecBlk, ResBlk, HierarAttDecBlk +# from models.lateral_blocks import BasicLatBlk +# from models.aspp import ASPP, ASPPDeformable +# from models.ing import * +# from models.refiner import Refiner, RefinerPVTInChannels4, RefUNet +# from models.stem_layer import StemLayer +from BiRefNet_config import BiRefNetConfig + + +def image2patches(image, grid_h=2, grid_w=2, patch_ref=None, transformation='b c (hg h) (wg w) -> (b hg wg) c h w'): + if patch_ref is not None: + grid_h, grid_w = image.shape[-2] // patch_ref.shape[-2], image.shape[-1] // patch_ref.shape[-1] + patches = rearrange(image, transformation, hg=grid_h, wg=grid_w) + return patches + +def patches2image(patches, grid_h=2, grid_w=2, patch_ref=None, transformation='(b hg wg) c h w -> b c (hg h) (wg w)'): + if patch_ref is not None: + grid_h, grid_w = patch_ref.shape[-2] // patches[0].shape[-2], patch_ref.shape[-1] // patches[0].shape[-1] + image = rearrange(patches, transformation, hg=grid_h, wg=grid_w) + return image + +class BiRefNet( + PreTrainedModel +): + config_class = BiRefNetConfig + def __init__(self, bb_pretrained=True, config=BiRefNetConfig()): + super(BiRefNet, self).__init__(config) + bb_pretrained = config.bb_pretrained + self.config = Config() + self.epoch = 1 + self.bb = build_backbone(self.config.bb, pretrained=bb_pretrained) + + channels = self.config.lateral_channels_in_collection + + if self.config.auxiliary_classification: + self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) + self.cls_head = nn.Sequential( + nn.Linear(channels[0], len(class_labels_TR_sorted)) + ) + + if self.config.squeeze_block: + self.squeeze_module = nn.Sequential(*[ + eval(self.config.squeeze_block.split('_x')[0])(channels[0]+sum(self.config.cxt), channels[0]) + for _ in range(eval(self.config.squeeze_block.split('_x')[1])) + ]) + + self.decoder = Decoder(channels) + + if self.config.ender: + self.dec_end = nn.Sequential( + nn.Conv2d(1, 16, 3, 1, 1), + nn.Conv2d(16, 1, 3, 1, 1), + nn.ReLU(inplace=True), + ) + + # refine patch-level segmentation + if self.config.refine: + if self.config.refine == 'itself': + self.stem_layer = StemLayer(in_channels=3+1, inter_channels=48, out_channels=3, norm_layer='BN' if self.config.batch_size > 1 else 'LN') + else: + self.refiner = eval('{}({})'.format(self.config.refine, 'in_channels=3+1')) + + if self.config.freeze_bb: + # Freeze the backbone... + print(self.named_parameters()) + for key, value in self.named_parameters(): + if 'bb.' in key and 'refiner.' not in key: + value.requires_grad = False + + def forward_enc(self, x): + if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']: + x1 = self.bb.conv1(x); x2 = self.bb.conv2(x1); x3 = self.bb.conv3(x2); x4 = self.bb.conv4(x3) + else: + x1, x2, x3, x4 = self.bb(x) + if self.config.mul_scl_ipt == 'cat': + B, C, H, W = x.shape + x1_, x2_, x3_, x4_ = self.bb(F.interpolate(x, size=(H//2, W//2), mode='bilinear', align_corners=True)) + x1 = torch.cat([x1, F.interpolate(x1_, size=x1.shape[2:], mode='bilinear', align_corners=True)], dim=1) + x2 = torch.cat([x2, F.interpolate(x2_, size=x2.shape[2:], mode='bilinear', align_corners=True)], dim=1) + x3 = torch.cat([x3, F.interpolate(x3_, size=x3.shape[2:], mode='bilinear', align_corners=True)], dim=1) + x4 = torch.cat([x4, F.interpolate(x4_, size=x4.shape[2:], mode='bilinear', align_corners=True)], dim=1) + elif self.config.mul_scl_ipt == 'add': + B, C, H, W = x.shape + x1_, x2_, x3_, x4_ = self.bb(F.interpolate(x, size=(H//2, W//2), mode='bilinear', align_corners=True)) + x1 = x1 + F.interpolate(x1_, size=x1.shape[2:], mode='bilinear', align_corners=True) + x2 = x2 + F.interpolate(x2_, size=x2.shape[2:], mode='bilinear', align_corners=True) + x3 = x3 + F.interpolate(x3_, size=x3.shape[2:], mode='bilinear', align_corners=True) + x4 = x4 + F.interpolate(x4_, size=x4.shape[2:], mode='bilinear', align_corners=True) + class_preds = self.cls_head(self.avgpool(x4).view(x4.shape[0], -1)) if self.training and self.config.auxiliary_classification else None + if self.config.cxt: + x4 = torch.cat( + ( + *[ + F.interpolate(x1, size=x4.shape[2:], mode='bilinear', align_corners=True), + F.interpolate(x2, size=x4.shape[2:], mode='bilinear', align_corners=True), + F.interpolate(x3, size=x4.shape[2:], mode='bilinear', align_corners=True), + ][-len(self.config.cxt):], + x4 + ), + dim=1 + ) + return (x1, x2, x3, x4), class_preds + + def forward_ori(self, x): + ########## Encoder ########## + (x1, x2, x3, x4), class_preds = self.forward_enc(x) + if self.config.squeeze_block: + x4 = self.squeeze_module(x4) + ########## Decoder ########## + features = [x, x1, x2, x3, x4] + if self.training and self.config.out_ref: + features.append(laplacian(torch.mean(x, dim=1).unsqueeze(1), kernel_size=5)) + scaled_preds = self.decoder(features) + return scaled_preds, class_preds + + def forward(self, x): + scaled_preds, class_preds = self.forward_ori(x) + class_preds_lst = [class_preds] + return [scaled_preds, class_preds_lst] if self.training else scaled_preds + + +class Decoder(nn.Module): + def __init__(self, channels): + super(Decoder, self).__init__() + self.config = Config() + DecoderBlock = eval(self.config.dec_blk) + LateralBlock = eval(self.config.lat_blk) + + if self.config.dec_ipt: + self.split = self.config.dec_ipt_split + N_dec_ipt = 64 + DBlock = SimpleConvs + ic = 64 + ipt_cha_opt = 1 + self.ipt_blk5 = DBlock(2**10*3 if self.split else 3, [N_dec_ipt, channels[0]//8][ipt_cha_opt], inter_channels=ic) + self.ipt_blk4 = DBlock(2**8*3 if self.split else 3, [N_dec_ipt, channels[0]//8][ipt_cha_opt], inter_channels=ic) + self.ipt_blk3 = DBlock(2**6*3 if self.split else 3, [N_dec_ipt, channels[1]//8][ipt_cha_opt], inter_channels=ic) + self.ipt_blk2 = DBlock(2**4*3 if self.split else 3, [N_dec_ipt, channels[2]//8][ipt_cha_opt], inter_channels=ic) + self.ipt_blk1 = DBlock(2**0*3 if self.split else 3, [N_dec_ipt, channels[3]//8][ipt_cha_opt], inter_channels=ic) + else: + self.split = None + + self.decoder_block4 = DecoderBlock(channels[0]+([N_dec_ipt, channels[0]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[1]) + self.decoder_block3 = DecoderBlock(channels[1]+([N_dec_ipt, channels[0]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[2]) + self.decoder_block2 = DecoderBlock(channels[2]+([N_dec_ipt, channels[1]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[3]) + self.decoder_block1 = DecoderBlock(channels[3]+([N_dec_ipt, channels[2]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[3]//2) + self.conv_out1 = nn.Sequential(nn.Conv2d(channels[3]//2+([N_dec_ipt, channels[3]//8][ipt_cha_opt] if self.config.dec_ipt else 0), 1, 1, 1, 0)) + + self.lateral_block4 = LateralBlock(channels[1], channels[1]) + self.lateral_block3 = LateralBlock(channels[2], channels[2]) + self.lateral_block2 = LateralBlock(channels[3], channels[3]) + + if self.config.ms_supervision: + self.conv_ms_spvn_4 = nn.Conv2d(channels[1], 1, 1, 1, 0) + self.conv_ms_spvn_3 = nn.Conv2d(channels[2], 1, 1, 1, 0) + self.conv_ms_spvn_2 = nn.Conv2d(channels[3], 1, 1, 1, 0) + + if self.config.out_ref: + _N = 16 + self.gdt_convs_4 = nn.Sequential(nn.Conv2d(channels[1], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True)) + self.gdt_convs_3 = nn.Sequential(nn.Conv2d(channels[2], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True)) + self.gdt_convs_2 = nn.Sequential(nn.Conv2d(channels[3], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True)) + + self.gdt_convs_pred_4 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + self.gdt_convs_pred_3 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + self.gdt_convs_pred_2 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + + self.gdt_convs_attn_4 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + self.gdt_convs_attn_3 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + self.gdt_convs_attn_2 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + + def forward(self, features): + if self.training and self.config.out_ref: + outs_gdt_pred = [] + outs_gdt_label = [] + x, x1, x2, x3, x4, gdt_gt = features + else: + x, x1, x2, x3, x4 = features + outs = [] + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=x4, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + x4 = torch.cat((x4, self.ipt_blk5(F.interpolate(patches_batch, size=x4.shape[2:], mode='bilinear', align_corners=True))), 1) + p4 = self.decoder_block4(x4) + m4 = self.conv_ms_spvn_4(p4) if self.config.ms_supervision and self.training else None + if self.config.out_ref: + p4_gdt = self.gdt_convs_4(p4) + if self.training: + # >> GT: + m4_dia = m4 + gdt_label_main_4 = gdt_gt * F.interpolate(m4_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True) + outs_gdt_label.append(gdt_label_main_4) + # >> Pred: + gdt_pred_4 = self.gdt_convs_pred_4(p4_gdt) + outs_gdt_pred.append(gdt_pred_4) + gdt_attn_4 = self.gdt_convs_attn_4(p4_gdt).sigmoid() + # >> Finally: + p4 = p4 * gdt_attn_4 + _p4 = F.interpolate(p4, size=x3.shape[2:], mode='bilinear', align_corners=True) + _p3 = _p4 + self.lateral_block4(x3) + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=_p3, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + _p3 = torch.cat((_p3, self.ipt_blk4(F.interpolate(patches_batch, size=x3.shape[2:], mode='bilinear', align_corners=True))), 1) + p3 = self.decoder_block3(_p3) + m3 = self.conv_ms_spvn_3(p3) if self.config.ms_supervision and self.training else None + if self.config.out_ref: + p3_gdt = self.gdt_convs_3(p3) + if self.training: + # >> GT: + # m3 --dilation--> m3_dia + # G_3^gt * m3_dia --> G_3^m, which is the label of gradient + m3_dia = m3 + gdt_label_main_3 = gdt_gt * F.interpolate(m3_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True) + outs_gdt_label.append(gdt_label_main_3) + # >> Pred: + # p3 --conv--BN--> F_3^G, where F_3^G predicts the \hat{G_3} with xx + # F_3^G --sigmoid--> A_3^G + gdt_pred_3 = self.gdt_convs_pred_3(p3_gdt) + outs_gdt_pred.append(gdt_pred_3) + gdt_attn_3 = self.gdt_convs_attn_3(p3_gdt).sigmoid() + # >> Finally: + # p3 = p3 * A_3^G + p3 = p3 * gdt_attn_3 + _p3 = F.interpolate(p3, size=x2.shape[2:], mode='bilinear', align_corners=True) + _p2 = _p3 + self.lateral_block3(x2) + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=_p2, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + _p2 = torch.cat((_p2, self.ipt_blk3(F.interpolate(patches_batch, size=x2.shape[2:], mode='bilinear', align_corners=True))), 1) + p2 = self.decoder_block2(_p2) + m2 = self.conv_ms_spvn_2(p2) if self.config.ms_supervision and self.training else None + if self.config.out_ref: + p2_gdt = self.gdt_convs_2(p2) + if self.training: + # >> GT: + m2_dia = m2 + gdt_label_main_2 = gdt_gt * F.interpolate(m2_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True) + outs_gdt_label.append(gdt_label_main_2) + # >> Pred: + gdt_pred_2 = self.gdt_convs_pred_2(p2_gdt) + outs_gdt_pred.append(gdt_pred_2) + gdt_attn_2 = self.gdt_convs_attn_2(p2_gdt).sigmoid() + # >> Finally: + p2 = p2 * gdt_attn_2 + _p2 = F.interpolate(p2, size=x1.shape[2:], mode='bilinear', align_corners=True) + _p1 = _p2 + self.lateral_block2(x1) + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=_p1, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + _p1 = torch.cat((_p1, self.ipt_blk2(F.interpolate(patches_batch, size=x1.shape[2:], mode='bilinear', align_corners=True))), 1) + _p1 = self.decoder_block1(_p1) + _p1 = F.interpolate(_p1, size=x.shape[2:], mode='bilinear', align_corners=True) + + if self.config.dec_ipt: + patches_batch = image2patches(x, patch_ref=_p1, transformation='b c (hg h) (wg w) -> b (c hg wg) h w') if self.split else x + _p1 = torch.cat((_p1, self.ipt_blk1(F.interpolate(patches_batch, size=x.shape[2:], mode='bilinear', align_corners=True))), 1) + p1_out = self.conv_out1(_p1) + + if self.config.ms_supervision and self.training: + outs.append(m4) + outs.append(m3) + outs.append(m2) + outs.append(p1_out) + return outs if not (self.config.out_ref and self.training) else ([outs_gdt_pred, outs_gdt_label], outs) + + +class SimpleConvs(nn.Module): + def __init__( + self, in_channels: int, out_channels: int, inter_channels=64 + ) -> None: + super().__init__() + self.conv1 = nn.Conv2d(in_channels, inter_channels, 3, 1, 1) + self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, 1) + + def forward(self, x): + return self.conv_out(self.conv1(x)) diff --git a/models/RMBG/BiRefNet/config.json b/models/RMBG/BiRefNet/config.json new file mode 100644 index 0000000000000000000000000000000000000000..25977f83be51d6c9715a2919221f307973c2d1b1 --- /dev/null +++ b/models/RMBG/BiRefNet/config.json @@ -0,0 +1,20 @@ +{ + "_name_or_path": "1038lab/BiRefNet", + "architectures": [ + "BiRefNet" + ], + "auto_map": { + "AutoConfig": "BiRefNet_config.BiRefNetConfig", + "AutoModelForImageSegmentation": "birefnet.BiRefNet" + }, + "custom_pipelines": { + "image-segmentation": { + "pt": [ + "AutoModelForImageSegmentation" + ], + "tf": [], + "type": "image" + } + }, + "bb_pretrained": false +} \ No newline at end of file diff --git a/models/RMBG/BiRefNet/gitattributes b/models/RMBG/BiRefNet/gitattributes new file mode 100644 index 0000000000000000000000000000000000000000..a6344aac8c09253b3b630fb776ae94478aa0275b --- /dev/null +++ b/models/RMBG/BiRefNet/gitattributes @@ -0,0 +1,35 @@ +*.7z filter=lfs diff=lfs merge=lfs -text +*.arrow filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +*.bz2 filter=lfs diff=lfs merge=lfs -text +*.ckpt filter=lfs diff=lfs merge=lfs -text +*.ftz filter=lfs diff=lfs merge=lfs -text +*.gz filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text +*.joblib filter=lfs diff=lfs merge=lfs -text +*.lfs.* filter=lfs diff=lfs merge=lfs -text +*.mlmodel filter=lfs diff=lfs merge=lfs -text +*.model filter=lfs diff=lfs merge=lfs -text +*.msgpack filter=lfs diff=lfs merge=lfs -text +*.npy filter=lfs diff=lfs merge=lfs -text +*.npz filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.ot filter=lfs diff=lfs merge=lfs -text +*.parquet filter=lfs diff=lfs merge=lfs -text +*.pb filter=lfs diff=lfs merge=lfs -text +*.pickle filter=lfs diff=lfs merge=lfs -text +*.pkl filter=lfs diff=lfs merge=lfs -text +*.pt filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.rar filter=lfs diff=lfs merge=lfs -text +*.safetensors filter=lfs diff=lfs merge=lfs -text +saved_model/**/* filter=lfs diff=lfs merge=lfs -text +*.tar.* filter=lfs diff=lfs merge=lfs -text +*.tar filter=lfs diff=lfs merge=lfs -text +*.tflite filter=lfs diff=lfs merge=lfs -text +*.tgz filter=lfs diff=lfs merge=lfs -text +*.wasm filter=lfs diff=lfs merge=lfs -text +*.xz filter=lfs diff=lfs merge=lfs -text +*.zip filter=lfs diff=lfs merge=lfs -text +*.zst filter=lfs diff=lfs merge=lfs -text +*tfevents* filter=lfs diff=lfs merge=lfs -text diff --git a/models/RMBG/INSPYRENET/gitattributes b/models/RMBG/INSPYRENET/gitattributes new file mode 100644 index 0000000000000000000000000000000000000000..a6344aac8c09253b3b630fb776ae94478aa0275b --- /dev/null +++ b/models/RMBG/INSPYRENET/gitattributes @@ -0,0 +1,35 @@ +*.7z filter=lfs diff=lfs merge=lfs -text +*.arrow filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +*.bz2 filter=lfs diff=lfs merge=lfs -text +*.ckpt filter=lfs diff=lfs merge=lfs -text +*.ftz filter=lfs diff=lfs merge=lfs -text +*.gz filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text +*.joblib filter=lfs diff=lfs merge=lfs -text +*.lfs.* filter=lfs diff=lfs merge=lfs -text +*.mlmodel filter=lfs diff=lfs merge=lfs -text +*.model filter=lfs diff=lfs merge=lfs -text +*.msgpack filter=lfs diff=lfs merge=lfs -text +*.npy filter=lfs diff=lfs merge=lfs -text +*.npz filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.ot filter=lfs diff=lfs merge=lfs -text +*.parquet filter=lfs diff=lfs merge=lfs -text +*.pb filter=lfs diff=lfs merge=lfs -text +*.pickle filter=lfs diff=lfs merge=lfs -text +*.pkl filter=lfs diff=lfs merge=lfs -text +*.pt filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.rar filter=lfs diff=lfs merge=lfs -text +*.safetensors filter=lfs diff=lfs merge=lfs -text +saved_model/**/* filter=lfs diff=lfs merge=lfs -text +*.tar.* filter=lfs diff=lfs merge=lfs -text +*.tar filter=lfs diff=lfs merge=lfs -text +*.tflite filter=lfs diff=lfs merge=lfs -text +*.tgz filter=lfs diff=lfs merge=lfs -text +*.wasm filter=lfs diff=lfs merge=lfs -text +*.xz filter=lfs diff=lfs merge=lfs -text +*.zip filter=lfs diff=lfs merge=lfs -text +*.zst filter=lfs diff=lfs merge=lfs -text +*tfevents* filter=lfs diff=lfs merge=lfs -text diff --git a/models/RMBG/INSPYRENET/inspyrenet.pth b/models/RMBG/INSPYRENET/inspyrenet.pth new file mode 100644 index 0000000000000000000000000000000000000000..8bc76f6aa859d02cc6df5df4a5510c889405adb6 --- /dev/null +++ b/models/RMBG/INSPYRENET/inspyrenet.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0a6fe2a73ab0532d6d0b8d82849a9760a226df719e3063d09b4149ece6f80fcd +size 367520613 diff --git a/models/RMBG/INSPYRENET/inspyrenet.safetensors b/models/RMBG/INSPYRENET/inspyrenet.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..9bdd2bc3bd890bf60f848c2c23fa3b318ba1e879 --- /dev/null +++ b/models/RMBG/INSPYRENET/inspyrenet.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5c8ab4541f9a960b376a80362b07f1d96262e689c501365e6323faa958de47ea +size 367120444 diff --git a/models/RMBG/RMBG-2.0/BiRefNet_config.py b/models/RMBG/RMBG-2.0/BiRefNet_config.py new file mode 100644 index 0000000000000000000000000000000000000000..37c8ac58bec2f52dac34204978a7b61b69e3da76 --- /dev/null +++ b/models/RMBG/RMBG-2.0/BiRefNet_config.py @@ -0,0 +1,11 @@ +from transformers import PretrainedConfig + +class BiRefNetConfig(PretrainedConfig): + model_type = "SegformerForSemanticSegmentation" + def __init__( + self, + bb_pretrained=False, + **kwargs + ): + self.bb_pretrained = bb_pretrained + super().__init__(**kwargs) diff --git a/models/RMBG/RMBG-2.0/__pycache__/BiRefNet_config.cpython-310.pyc b/models/RMBG/RMBG-2.0/__pycache__/BiRefNet_config.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a74404adda606b026cc2a51616156f376999a135 Binary files /dev/null and b/models/RMBG/RMBG-2.0/__pycache__/BiRefNet_config.cpython-310.pyc differ diff --git a/models/RMBG/RMBG-2.0/birefnet.py b/models/RMBG/RMBG-2.0/birefnet.py new file mode 100644 index 0000000000000000000000000000000000000000..1ed28de11ce84e5f6b2dbf2a0dbb66412bc28513 --- /dev/null +++ b/models/RMBG/RMBG-2.0/birefnet.py @@ -0,0 +1,2244 @@ +### config.py + +import os +import math + + +class Config(): + def __init__(self) -> None: + # PATH settings + self.sys_home_dir = os.path.expanduser('~') # Make up your file system as: SYS_HOME_DIR/codes/dis/BiRefNet, SYS_HOME_DIR/datasets/dis/xx, SYS_HOME_DIR/weights/xx + + # TASK settings + self.task = ['DIS5K', 'COD', 'HRSOD', 'DIS5K+HRSOD+HRS10K', 'P3M-10k'][0] + self.training_set = { + 'DIS5K': ['DIS-TR', 'DIS-TR+DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4'][0], + 'COD': 'TR-COD10K+TR-CAMO', + 'HRSOD': ['TR-DUTS', 'TR-HRSOD', 'TR-UHRSD', 'TR-DUTS+TR-HRSOD', 'TR-DUTS+TR-UHRSD', 'TR-HRSOD+TR-UHRSD', 'TR-DUTS+TR-HRSOD+TR-UHRSD'][5], + 'DIS5K+HRSOD+HRS10K': 'DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4+DIS-TR+TE-HRS10K+TE-HRSOD+TE-UHRSD+TR-HRS10K+TR-HRSOD+TR-UHRSD', # leave DIS-VD for evaluation. + 'P3M-10k': 'TR-P3M-10k', + }[self.task] + self.prompt4loc = ['dense', 'sparse'][0] + + # Faster-Training settings + self.load_all = True + self.compile = True # 1. Trigger CPU memory leak in some extend, which is an inherent problem of PyTorch. + # Machines with > 70GB CPU memory can run the whole training on DIS5K with default setting. + # 2. Higher PyTorch version may fix it: https://github.com/pytorch/pytorch/issues/119607. + # 3. But compile in Pytorch > 2.0.1 seems to bring no acceleration for training. + self.precisionHigh = True + + # MODEL settings + self.ms_supervision = True + self.out_ref = self.ms_supervision and True + self.dec_ipt = True + self.dec_ipt_split = True + self.cxt_num = [0, 3][1] # multi-scale skip connections from encoder + self.mul_scl_ipt = ['', 'add', 'cat'][2] + self.dec_att = ['', 'ASPP', 'ASPPDeformable'][2] + self.squeeze_block = ['', 'BasicDecBlk_x1', 'ResBlk_x4', 'ASPP_x3', 'ASPPDeformable_x3'][1] + self.dec_blk = ['BasicDecBlk', 'ResBlk', 'HierarAttDecBlk'][0] + + # TRAINING settings + self.batch_size = 4 + self.IoU_finetune_last_epochs = [ + 0, + { + 'DIS5K': -50, + 'COD': -20, + 'HRSOD': -20, + 'DIS5K+HRSOD+HRS10K': -20, + 'P3M-10k': -20, + }[self.task] + ][1] # choose 0 to skip + self.lr = (1e-4 if 'DIS5K' in self.task else 1e-5) * math.sqrt(self.batch_size / 4) # DIS needs high lr to converge faster. Adapt the lr linearly + self.size = 1024 + self.num_workers = max(4, self.batch_size) # will be decrease to min(it, batch_size) at the initialization of the data_loader + + # Backbone settings + self.bb = [ + 'vgg16', 'vgg16bn', 'resnet50', # 0, 1, 2 + 'swin_v1_t', 'swin_v1_s', # 3, 4 + 'swin_v1_b', 'swin_v1_l', # 5-bs9, 6-bs4 + 'pvt_v2_b0', 'pvt_v2_b1', # 7, 8 + 'pvt_v2_b2', 'pvt_v2_b5', # 9-bs10, 10-bs5 + ][6] + self.lateral_channels_in_collection = { + 'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64], + 'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64], + 'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192], + 'swin_v1_t': [768, 384, 192, 96], 'swin_v1_s': [768, 384, 192, 96], + 'pvt_v2_b0': [256, 160, 64, 32], 'pvt_v2_b1': [512, 320, 128, 64], + }[self.bb] + if self.mul_scl_ipt == 'cat': + self.lateral_channels_in_collection = [channel * 2 for channel in self.lateral_channels_in_collection] + self.cxt = self.lateral_channels_in_collection[1:][::-1][-self.cxt_num:] if self.cxt_num else [] + + # MODEL settings - inactive + self.lat_blk = ['BasicLatBlk'][0] + self.dec_channels_inter = ['fixed', 'adap'][0] + self.refine = ['', 'itself', 'RefUNet', 'Refiner', 'RefinerPVTInChannels4'][0] + self.progressive_ref = self.refine and True + self.ender = self.progressive_ref and False + self.scale = self.progressive_ref and 2 + self.auxiliary_classification = False # Only for DIS5K, where class labels are saved in `dataset.py`. + self.refine_iteration = 1 + self.freeze_bb = False + self.model = [ + 'BiRefNet', + ][0] + if self.dec_blk == 'HierarAttDecBlk': + self.batch_size = 2 ** [0, 1, 2, 3, 4][2] + + # TRAINING settings - inactive + self.preproc_methods = ['flip', 'enhance', 'rotate', 'pepper', 'crop'][:4] + self.optimizer = ['Adam', 'AdamW'][1] + self.lr_decay_epochs = [1e5] # Set to negative N to decay the lr in the last N-th epoch. + self.lr_decay_rate = 0.5 + # Loss + self.lambdas_pix_last = { + # not 0 means opening this loss + # original rate -- 1 : 30 : 1.5 : 0.2, bce x 30 + 'bce': 30 * 1, # high performance + 'iou': 0.5 * 1, # 0 / 255 + 'iou_patch': 0.5 * 0, # 0 / 255, win_size = (64, 64) + 'mse': 150 * 0, # can smooth the saliency map + 'triplet': 3 * 0, + 'reg': 100 * 0, + 'ssim': 10 * 1, # help contours, + 'cnt': 5 * 0, # help contours + 'structure': 5 * 0, # structure loss from codes of MVANet. A little improvement on DIS-TE[1,2,3], a bit more decrease on DIS-TE4. + } + self.lambdas_cls = { + 'ce': 5.0 + } + # Adv + self.lambda_adv_g = 10. * 0 # turn to 0 to avoid adv training + self.lambda_adv_d = 3. * (self.lambda_adv_g > 0) + + # PATH settings - inactive + self.data_root_dir = os.path.join(self.sys_home_dir, 'datasets/dis') + self.weights_root_dir = os.path.join(self.sys_home_dir, 'weights') + self.weights = { + 'pvt_v2_b2': os.path.join(self.weights_root_dir, 'pvt_v2_b2.pth'), + 'pvt_v2_b5': os.path.join(self.weights_root_dir, ['pvt_v2_b5.pth', 'pvt_v2_b5_22k.pth'][0]), + 'swin_v1_b': os.path.join(self.weights_root_dir, ['swin_base_patch4_window12_384_22kto1k.pth', 'swin_base_patch4_window12_384_22k.pth'][0]), + 'swin_v1_l': os.path.join(self.weights_root_dir, ['swin_large_patch4_window12_384_22kto1k.pth', 'swin_large_patch4_window12_384_22k.pth'][0]), + 'swin_v1_t': os.path.join(self.weights_root_dir, ['swin_tiny_patch4_window7_224_22kto1k_finetune.pth'][0]), + 'swin_v1_s': os.path.join(self.weights_root_dir, ['swin_small_patch4_window7_224_22kto1k_finetune.pth'][0]), + 'pvt_v2_b0': os.path.join(self.weights_root_dir, ['pvt_v2_b0.pth'][0]), + 'pvt_v2_b1': os.path.join(self.weights_root_dir, ['pvt_v2_b1.pth'][0]), + } + + # Callbacks - inactive + self.verbose_eval = True + self.only_S_MAE = False + self.use_fp16 = False # Bugs. It may cause nan in training. + self.SDPA_enabled = False # Bugs. Slower and errors occur in multi-GPUs + + # others + self.device = [0, 'cpu'][0] # .to(0) == .to('cuda:0') + + self.batch_size_valid = 1 + self.rand_seed = 7 + # run_sh_file = [f for f in os.listdir('.') if 'train.sh' == f] + [os.path.join('..', f) for f in os.listdir('..') if 'train.sh' == f] + # with open(run_sh_file[0], 'r') as f: + # lines = f.readlines() + # self.save_last = int([l.strip() for l in lines if '"{}")'.format(self.task) in l and 'val_last=' in l][0].split('val_last=')[-1].split()[0]) + # self.save_step = int([l.strip() for l in lines if '"{}")'.format(self.task) in l and 'step=' in l][0].split('step=')[-1].split()[0]) + # self.val_step = [0, self.save_step][0] + + def print_task(self) -> None: + # Return task for choosing settings in shell scripts. + print(self.task) + + + +### models/backbones/pvt_v2.py + +import torch +import torch.nn as nn +from functools import partial + +from timm.models.layers import DropPath, to_2tuple, trunc_normal_ +from timm.models.registry import register_model + +import math + +# from config import Config + +# config = Config() + +class Mlp(nn.Module): + def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.dwconv = DWConv(hidden_features) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def forward(self, x, H, W): + x = self.fc1(x) + x = self.dwconv(x, H, W) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +class Attention(nn.Module): + def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., sr_ratio=1): + super().__init__() + assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}." + + self.dim = dim + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = qk_scale or head_dim ** -0.5 + + self.q = nn.Linear(dim, dim, bias=qkv_bias) + self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias) + self.attn_drop_prob = attn_drop + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + self.sr_ratio = sr_ratio + if sr_ratio > 1: + self.sr = nn.Conv2d(dim, dim, kernel_size=sr_ratio, stride=sr_ratio) + self.norm = nn.LayerNorm(dim) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def forward(self, x, H, W): + B, N, C = x.shape + q = self.q(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) + + if self.sr_ratio > 1: + x_ = x.permute(0, 2, 1).reshape(B, C, H, W) + x_ = self.sr(x_).reshape(B, C, -1).permute(0, 2, 1) + x_ = self.norm(x_) + kv = self.kv(x_).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + else: + kv = self.kv(x).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + k, v = kv[0], kv[1] + + if config.SDPA_enabled: + x = torch.nn.functional.scaled_dot_product_attention( + q, k, v, + attn_mask=None, dropout_p=self.attn_drop_prob, is_causal=False + ).transpose(1, 2).reshape(B, N, C) + else: + attn = (q @ k.transpose(-2, -1)) * self.scale + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B, N, C) + x = self.proj(x) + x = self.proj_drop(x) + + return x + + +class Block(nn.Module): + + def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., + drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, sr_ratio=1): + super().__init__() + self.norm1 = norm_layer(dim) + self.attn = Attention( + dim, + num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, + attn_drop=attn_drop, proj_drop=drop, sr_ratio=sr_ratio) + # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def forward(self, x, H, W): + x = x + self.drop_path(self.attn(self.norm1(x), H, W)) + x = x + self.drop_path(self.mlp(self.norm2(x), H, W)) + + return x + + +class OverlapPatchEmbed(nn.Module): + """ Image to Patch Embedding + """ + + def __init__(self, img_size=224, patch_size=7, stride=4, in_channels=3, embed_dim=768): + super().__init__() + img_size = to_2tuple(img_size) + patch_size = to_2tuple(patch_size) + + self.img_size = img_size + self.patch_size = patch_size + self.H, self.W = img_size[0] // patch_size[0], img_size[1] // patch_size[1] + self.num_patches = self.H * self.W + self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=stride, + padding=(patch_size[0] // 2, patch_size[1] // 2)) + self.norm = nn.LayerNorm(embed_dim) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def forward(self, x): + x = self.proj(x) + _, _, H, W = x.shape + x = x.flatten(2).transpose(1, 2) + x = self.norm(x) + + return x, H, W + + +class PyramidVisionTransformerImpr(nn.Module): + def __init__(self, img_size=224, patch_size=16, in_channels=3, num_classes=1000, embed_dims=[64, 128, 256, 512], + num_heads=[1, 2, 4, 8], mlp_ratios=[4, 4, 4, 4], qkv_bias=False, qk_scale=None, drop_rate=0., + attn_drop_rate=0., drop_path_rate=0., norm_layer=nn.LayerNorm, + depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1]): + super().__init__() + self.num_classes = num_classes + self.depths = depths + + # patch_embed + self.patch_embed1 = OverlapPatchEmbed(img_size=img_size, patch_size=7, stride=4, in_channels=in_channels, + embed_dim=embed_dims[0]) + self.patch_embed2 = OverlapPatchEmbed(img_size=img_size // 4, patch_size=3, stride=2, in_channels=embed_dims[0], + embed_dim=embed_dims[1]) + self.patch_embed3 = OverlapPatchEmbed(img_size=img_size // 8, patch_size=3, stride=2, in_channels=embed_dims[1], + embed_dim=embed_dims[2]) + self.patch_embed4 = OverlapPatchEmbed(img_size=img_size // 16, patch_size=3, stride=2, in_channels=embed_dims[2], + embed_dim=embed_dims[3]) + + # transformer encoder + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule + cur = 0 + self.block1 = nn.ModuleList([Block( + dim=embed_dims[0], num_heads=num_heads[0], mlp_ratio=mlp_ratios[0], qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, + sr_ratio=sr_ratios[0]) + for i in range(depths[0])]) + self.norm1 = norm_layer(embed_dims[0]) + + cur += depths[0] + self.block2 = nn.ModuleList([Block( + dim=embed_dims[1], num_heads=num_heads[1], mlp_ratio=mlp_ratios[1], qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, + sr_ratio=sr_ratios[1]) + for i in range(depths[1])]) + self.norm2 = norm_layer(embed_dims[1]) + + cur += depths[1] + self.block3 = nn.ModuleList([Block( + dim=embed_dims[2], num_heads=num_heads[2], mlp_ratio=mlp_ratios[2], qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, + sr_ratio=sr_ratios[2]) + for i in range(depths[2])]) + self.norm3 = norm_layer(embed_dims[2]) + + cur += depths[2] + self.block4 = nn.ModuleList([Block( + dim=embed_dims[3], num_heads=num_heads[3], mlp_ratio=mlp_ratios[3], qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, + sr_ratio=sr_ratios[3]) + for i in range(depths[3])]) + self.norm4 = norm_layer(embed_dims[3]) + + # classification head + # self.head = nn.Linear(embed_dims[3], num_classes) if num_classes > 0 else nn.Identity() + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + elif isinstance(m, nn.Conv2d): + fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + fan_out //= m.groups + m.weight.data.normal_(0, math.sqrt(2.0 / fan_out)) + if m.bias is not None: + m.bias.data.zero_() + + def init_weights(self, pretrained=None): + if isinstance(pretrained, str): + logger = 1 + #load_checkpoint(self, pretrained, map_location='cpu', strict=False, logger=logger) + + def reset_drop_path(self, drop_path_rate): + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(self.depths))] + cur = 0 + for i in range(self.depths[0]): + self.block1[i].drop_path.drop_prob = dpr[cur + i] + + cur += self.depths[0] + for i in range(self.depths[1]): + self.block2[i].drop_path.drop_prob = dpr[cur + i] + + cur += self.depths[1] + for i in range(self.depths[2]): + self.block3[i].drop_path.drop_prob = dpr[cur + i] + + cur += self.depths[2] + for i in range(self.depths[3]): + self.block4[i].drop_path.drop_prob = dpr[cur + i] + + def freeze_patch_emb(self): + self.patch_embed1.requires_grad = False + + @torch.jit.ignore + def no_weight_decay(self): + return {'pos_embed1', 'pos_embed2', 'pos_embed3', 'pos_embed4', 'cls_token'} # has pos_embed may be better + + def get_classifier(self): + return self.head + + def reset_classifier(self, num_classes, global_pool=''): + self.num_classes = num_classes + self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity() + + def forward_features(self, x): + B = x.shape[0] + outs = [] + + # stage 1 + x, H, W = self.patch_embed1(x) + for i, blk in enumerate(self.block1): + x = blk(x, H, W) + x = self.norm1(x) + x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() + outs.append(x) + + # stage 2 + x, H, W = self.patch_embed2(x) + for i, blk in enumerate(self.block2): + x = blk(x, H, W) + x = self.norm2(x) + x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() + outs.append(x) + + # stage 3 + x, H, W = self.patch_embed3(x) + for i, blk in enumerate(self.block3): + x = blk(x, H, W) + x = self.norm3(x) + x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() + outs.append(x) + + # stage 4 + x, H, W = self.patch_embed4(x) + for i, blk in enumerate(self.block4): + x = blk(x, H, W) + x = self.norm4(x) + x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() + outs.append(x) + + return outs + + # return x.mean(dim=1) + + def forward(self, x): + x = self.forward_features(x) + # x = self.head(x) + + return x + + +class DWConv(nn.Module): + def __init__(self, dim=768): + super(DWConv, self).__init__() + self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim) + + def forward(self, x, H, W): + B, N, C = x.shape + x = x.transpose(1, 2).view(B, C, H, W).contiguous() + x = self.dwconv(x) + x = x.flatten(2).transpose(1, 2) + + return x + + +def _conv_filter(state_dict, patch_size=16): + """ convert patch embedding weight from manual patchify + linear proj to conv""" + out_dict = {} + for k, v in state_dict.items(): + if 'patch_embed.proj.weight' in k: + v = v.reshape((v.shape[0], 3, patch_size, patch_size)) + out_dict[k] = v + + return out_dict + + +## @register_model +class pvt_v2_b0(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b0, self).__init__( + patch_size=4, embed_dims=[32, 64, 160, 256], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[2, 2, 2, 2], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + + + +## @register_model +class pvt_v2_b1(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b1, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[2, 2, 2, 2], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + +## @register_model +class pvt_v2_b2(PyramidVisionTransformerImpr): + def __init__(self, in_channels=3, **kwargs): + super(pvt_v2_b2, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1, in_channels=in_channels) + +## @register_model +class pvt_v2_b3(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b3, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 18, 3], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + +## @register_model +class pvt_v2_b4(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b4, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 8, 27, 3], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + + +## @register_model +class pvt_v2_b5(PyramidVisionTransformerImpr): + def __init__(self, **kwargs): + super(pvt_v2_b5, self).__init__( + patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[4, 4, 4, 4], + qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 6, 40, 3], sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, drop_path_rate=0.1) + + + +### models/backbones/swin_v1.py + +# -------------------------------------------------------- +# Swin Transformer +# Copyright (c) 2021 Microsoft +# Licensed under The MIT License [see LICENSE for details] +# Written by Ze Liu, Yutong Lin, Yixuan Wei +# -------------------------------------------------------- + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as checkpoint +import numpy as np +from timm.models.layers import DropPath, to_2tuple, trunc_normal_ + +# from config import Config + + +# config = Config() + +class Mlp(nn.Module): + """ Multilayer perceptron.""" + + def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +def window_partition(x, window_size): + """ + Args: + x: (B, H, W, C) + window_size (int): window size + + Returns: + windows: (num_windows*B, window_size, window_size, C) + """ + B, H, W, C = x.shape + x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) + windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) + return windows + + +def window_reverse(windows, window_size, H, W): + """ + Args: + windows: (num_windows*B, window_size, window_size, C) + window_size (int): Window size + H (int): Height of image + W (int): Width of image + + Returns: + x: (B, H, W, C) + """ + B = int(windows.shape[0] / (H * W / window_size / window_size)) + x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1) + x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) + return x + + +class WindowAttention(nn.Module): + """ Window based multi-head self attention (W-MSA) module with relative position bias. + It supports both of shifted and non-shifted window. + + Args: + dim (int): Number of input channels. + window_size (tuple[int]): The height and width of the window. + num_heads (int): Number of attention heads. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set + attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0 + proj_drop (float, optional): Dropout ratio of output. Default: 0.0 + """ + + def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.): + + super().__init__() + self.dim = dim + self.window_size = window_size # Wh, Ww + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = qk_scale or head_dim ** -0.5 + + # define a parameter table of relative position bias + self.relative_position_bias_table = nn.Parameter( + torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)) # 2*Wh-1 * 2*Ww-1, nH + + # get pair-wise relative position index for each token inside the window + coords_h = torch.arange(self.window_size[0]) + coords_w = torch.arange(self.window_size[1]) + coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing='ij')) # 2, Wh, Ww + coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww + relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww + relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2 + relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0 + relative_coords[:, :, 1] += self.window_size[1] - 1 + relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1 + relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww + self.register_buffer("relative_position_index", relative_position_index) + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop_prob = attn_drop + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + trunc_normal_(self.relative_position_bias_table, std=.02) + self.softmax = nn.Softmax(dim=-1) + + def forward(self, x, mask=None): + """ Forward function. + + Args: + x: input features with shape of (num_windows*B, N, C) + mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None + """ + B_, N, C = x.shape + qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple) + + q = q * self.scale + + if config.SDPA_enabled: + x = torch.nn.functional.scaled_dot_product_attention( + q, k, v, + attn_mask=None, dropout_p=self.attn_drop_prob, is_causal=False + ).transpose(1, 2).reshape(B_, N, C) + else: + attn = (q @ k.transpose(-2, -1)) + + relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view( + self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1) # Wh*Ww,Wh*Ww,nH + relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww + attn = attn + relative_position_bias.unsqueeze(0) + + if mask is not None: + nW = mask.shape[0] + attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0) + attn = attn.view(-1, self.num_heads, N, N) + attn = self.softmax(attn) + else: + attn = self.softmax(attn) + + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B_, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class SwinTransformerBlock(nn.Module): + """ Swin Transformer Block. + + Args: + dim (int): Number of input channels. + num_heads (int): Number of attention heads. + window_size (int): Window size. + shift_size (int): Shift size for SW-MSA. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float, optional): Stochastic depth rate. Default: 0.0 + act_layer (nn.Module, optional): Activation layer. Default: nn.GELU + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + """ + + def __init__(self, dim, num_heads, window_size=7, shift_size=0, + mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0., + act_layer=nn.GELU, norm_layer=nn.LayerNorm): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.window_size = window_size + self.shift_size = shift_size + self.mlp_ratio = mlp_ratio + assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size" + + self.norm1 = norm_layer(dim) + self.attn = WindowAttention( + dim, window_size=to_2tuple(self.window_size), num_heads=num_heads, + qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop) + + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + + self.H = None + self.W = None + + def forward(self, x, mask_matrix): + """ Forward function. + + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + mask_matrix: Attention mask for cyclic shift. + """ + B, L, C = x.shape + H, W = self.H, self.W + assert L == H * W, "input feature has wrong size" + + shortcut = x + x = self.norm1(x) + x = x.view(B, H, W, C) + + # pad feature maps to multiples of window size + pad_l = pad_t = 0 + pad_r = (self.window_size - W % self.window_size) % self.window_size + pad_b = (self.window_size - H % self.window_size) % self.window_size + x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b)) + _, Hp, Wp, _ = x.shape + + # cyclic shift + if self.shift_size > 0: + shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)) + attn_mask = mask_matrix + else: + shifted_x = x + attn_mask = None + + # partition windows + x_windows = window_partition(shifted_x, self.window_size) # nW*B, window_size, window_size, C + x_windows = x_windows.view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C + + # W-MSA/SW-MSA + attn_windows = self.attn(x_windows, mask=attn_mask) # nW*B, window_size*window_size, C + + # merge windows + attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C) + shifted_x = window_reverse(attn_windows, self.window_size, Hp, Wp) # B H' W' C + + # reverse cyclic shift + if self.shift_size > 0: + x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2)) + else: + x = shifted_x + + if pad_r > 0 or pad_b > 0: + x = x[:, :H, :W, :].contiguous() + + x = x.view(B, H * W, C) + + # FFN + x = shortcut + self.drop_path(x) + x = x + self.drop_path(self.mlp(self.norm2(x))) + + return x + + +class PatchMerging(nn.Module): + """ Patch Merging Layer + + Args: + dim (int): Number of input channels. + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + """ + def __init__(self, dim, norm_layer=nn.LayerNorm): + super().__init__() + self.dim = dim + self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False) + self.norm = norm_layer(4 * dim) + + def forward(self, x, H, W): + """ Forward function. + + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + """ + B, L, C = x.shape + assert L == H * W, "input feature has wrong size" + + x = x.view(B, H, W, C) + + # padding + pad_input = (H % 2 == 1) or (W % 2 == 1) + if pad_input: + x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2)) + + x0 = x[:, 0::2, 0::2, :] # B H/2 W/2 C + x1 = x[:, 1::2, 0::2, :] # B H/2 W/2 C + x2 = x[:, 0::2, 1::2, :] # B H/2 W/2 C + x3 = x[:, 1::2, 1::2, :] # B H/2 W/2 C + x = torch.cat([x0, x1, x2, x3], -1) # B H/2 W/2 4*C + x = x.view(B, -1, 4 * C) # B H/2*W/2 4*C + + x = self.norm(x) + x = self.reduction(x) + + return x + + +class BasicLayer(nn.Module): + """ A basic Swin Transformer layer for one stage. + + Args: + dim (int): Number of feature channels + depth (int): Depths of this stage. + num_heads (int): Number of attention head. + window_size (int): Local window size. Default: 7. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. + """ + + def __init__(self, + dim, + depth, + num_heads, + window_size=7, + mlp_ratio=4., + qkv_bias=True, + qk_scale=None, + drop=0., + attn_drop=0., + drop_path=0., + norm_layer=nn.LayerNorm, + downsample=None, + use_checkpoint=False): + super().__init__() + self.window_size = window_size + self.shift_size = window_size // 2 + self.depth = depth + self.use_checkpoint = use_checkpoint + + # build blocks + self.blocks = nn.ModuleList([ + SwinTransformerBlock( + dim=dim, + num_heads=num_heads, + window_size=window_size, + shift_size=0 if (i % 2 == 0) else window_size // 2, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop, + attn_drop=attn_drop, + drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path, + norm_layer=norm_layer) + for i in range(depth)]) + + # patch merging layer + if downsample is not None: + self.downsample = downsample(dim=dim, norm_layer=norm_layer) + else: + self.downsample = None + + def forward(self, x, H, W): + """ Forward function. + + Args: + x: Input feature, tensor size (B, H*W, C). + H, W: Spatial resolution of the input feature. + """ + + # calculate attention mask for SW-MSA + Hp = int(np.ceil(H / self.window_size)) * self.window_size + Wp = int(np.ceil(W / self.window_size)) * self.window_size + img_mask = torch.zeros((1, Hp, Wp, 1), device=x.device) # 1 Hp Wp 1 + h_slices = (slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None)) + w_slices = (slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None)) + cnt = 0 + for h in h_slices: + for w in w_slices: + img_mask[:, h, w, :] = cnt + cnt += 1 + + mask_windows = window_partition(img_mask, self.window_size) # nW, window_size, window_size, 1 + mask_windows = mask_windows.view(-1, self.window_size * self.window_size) + attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) + attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0)) + + for blk in self.blocks: + blk.H, blk.W = H, W + if self.use_checkpoint: + x = checkpoint.checkpoint(blk, x, attn_mask) + else: + x = blk(x, attn_mask) + if self.downsample is not None: + x_down = self.downsample(x, H, W) + Wh, Ww = (H + 1) // 2, (W + 1) // 2 + return x, H, W, x_down, Wh, Ww + else: + return x, H, W, x, H, W + + +class PatchEmbed(nn.Module): + """ Image to Patch Embedding + + Args: + patch_size (int): Patch token size. Default: 4. + in_channels (int): Number of input image channels. Default: 3. + embed_dim (int): Number of linear projection output channels. Default: 96. + norm_layer (nn.Module, optional): Normalization layer. Default: None + """ + + def __init__(self, patch_size=4, in_channels=3, embed_dim=96, norm_layer=None): + super().__init__() + patch_size = to_2tuple(patch_size) + self.patch_size = patch_size + + self.in_channels = in_channels + self.embed_dim = embed_dim + + self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size) + if norm_layer is not None: + self.norm = norm_layer(embed_dim) + else: + self.norm = None + + def forward(self, x): + """Forward function.""" + # padding + _, _, H, W = x.size() + if W % self.patch_size[1] != 0: + x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1])) + if H % self.patch_size[0] != 0: + x = F.pad(x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0])) + + x = self.proj(x) # B C Wh Ww + if self.norm is not None: + Wh, Ww = x.size(2), x.size(3) + x = x.flatten(2).transpose(1, 2) + x = self.norm(x) + x = x.transpose(1, 2).view(-1, self.embed_dim, Wh, Ww) + + return x + + +class SwinTransformer(nn.Module): + """ Swin Transformer backbone. + A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows` - + https://arxiv.org/pdf/2103.14030 + + Args: + pretrain_img_size (int): Input image size for training the pretrained model, + used in absolute postion embedding. Default 224. + patch_size (int | tuple(int)): Patch size. Default: 4. + in_channels (int): Number of input image channels. Default: 3. + embed_dim (int): Number of linear projection output channels. Default: 96. + depths (tuple[int]): Depths of each Swin Transformer stage. + num_heads (tuple[int]): Number of attention head of each stage. + window_size (int): Window size. Default: 7. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4. + qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. + drop_rate (float): Dropout rate. + attn_drop_rate (float): Attention dropout rate. Default: 0. + drop_path_rate (float): Stochastic depth rate. Default: 0.2. + norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm. + ape (bool): If True, add absolute position embedding to the patch embedding. Default: False. + patch_norm (bool): If True, add normalization after patch embedding. Default: True. + out_indices (Sequence[int]): Output from which stages. + frozen_stages (int): Stages to be frozen (stop grad and set eval mode). + -1 means not freezing any parameters. + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. + """ + + def __init__(self, + pretrain_img_size=224, + patch_size=4, + in_channels=3, + embed_dim=96, + depths=[2, 2, 6, 2], + num_heads=[3, 6, 12, 24], + window_size=7, + mlp_ratio=4., + qkv_bias=True, + qk_scale=None, + drop_rate=0., + attn_drop_rate=0., + drop_path_rate=0.2, + norm_layer=nn.LayerNorm, + ape=False, + patch_norm=True, + out_indices=(0, 1, 2, 3), + frozen_stages=-1, + use_checkpoint=False): + super().__init__() + + self.pretrain_img_size = pretrain_img_size + self.num_layers = len(depths) + self.embed_dim = embed_dim + self.ape = ape + self.patch_norm = patch_norm + self.out_indices = out_indices + self.frozen_stages = frozen_stages + + # split image into non-overlapping patches + self.patch_embed = PatchEmbed( + patch_size=patch_size, in_channels=in_channels, embed_dim=embed_dim, + norm_layer=norm_layer if self.patch_norm else None) + + # absolute position embedding + if self.ape: + pretrain_img_size = to_2tuple(pretrain_img_size) + patch_size = to_2tuple(patch_size) + patches_resolution = [pretrain_img_size[0] // patch_size[0], pretrain_img_size[1] // patch_size[1]] + + self.absolute_pos_embed = nn.Parameter(torch.zeros(1, embed_dim, patches_resolution[0], patches_resolution[1])) + trunc_normal_(self.absolute_pos_embed, std=.02) + + self.pos_drop = nn.Dropout(p=drop_rate) + + # stochastic depth + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule + + # build layers + self.layers = nn.ModuleList() + for i_layer in range(self.num_layers): + layer = BasicLayer( + dim=int(embed_dim * 2 ** i_layer), + depth=depths[i_layer], + num_heads=num_heads[i_layer], + window_size=window_size, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])], + norm_layer=norm_layer, + downsample=PatchMerging if (i_layer < self.num_layers - 1) else None, + use_checkpoint=use_checkpoint) + self.layers.append(layer) + + num_features = [int(embed_dim * 2 ** i) for i in range(self.num_layers)] + self.num_features = num_features + + # add a norm layer for each output + for i_layer in out_indices: + layer = norm_layer(num_features[i_layer]) + layer_name = f'norm{i_layer}' + self.add_module(layer_name, layer) + + self._freeze_stages() + + def _freeze_stages(self): + if self.frozen_stages >= 0: + self.patch_embed.eval() + for param in self.patch_embed.parameters(): + param.requires_grad = False + + if self.frozen_stages >= 1 and self.ape: + self.absolute_pos_embed.requires_grad = False + + if self.frozen_stages >= 2: + self.pos_drop.eval() + for i in range(0, self.frozen_stages - 1): + m = self.layers[i] + m.eval() + for param in m.parameters(): + param.requires_grad = False + + + def forward(self, x): + """Forward function.""" + x = self.patch_embed(x) + + Wh, Ww = x.size(2), x.size(3) + if self.ape: + # interpolate the position embedding to the corresponding size + absolute_pos_embed = F.interpolate(self.absolute_pos_embed, size=(Wh, Ww), mode='bicubic') + x = (x + absolute_pos_embed) # B Wh*Ww C + + outs = []#x.contiguous()] + x = x.flatten(2).transpose(1, 2) + x = self.pos_drop(x) + for i in range(self.num_layers): + layer = self.layers[i] + x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww) + + if i in self.out_indices: + norm_layer = getattr(self, f'norm{i}') + x_out = norm_layer(x_out) + + out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous() + outs.append(out) + + return tuple(outs) + + def train(self, mode=True): + """Convert the model into training mode while keep layers freezed.""" + super(SwinTransformer, self).train(mode) + self._freeze_stages() + +def swin_v1_t(): + model = SwinTransformer(embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24], window_size=7) + return model + +def swin_v1_s(): + model = SwinTransformer(embed_dim=96, depths=[2, 2, 18, 2], num_heads=[3, 6, 12, 24], window_size=7) + return model + +def swin_v1_b(): + model = SwinTransformer(embed_dim=128, depths=[2, 2, 18, 2], num_heads=[4, 8, 16, 32], window_size=12) + return model + +def swin_v1_l(): + model = SwinTransformer(embed_dim=192, depths=[2, 2, 18, 2], num_heads=[6, 12, 24, 48], window_size=12) + return model + + + +### models/modules/deform_conv.py + +import torch +import torch.nn as nn +from torchvision.ops import deform_conv2d + + +class DeformableConv2d(nn.Module): + def __init__(self, + in_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1, + bias=False): + + super(DeformableConv2d, self).__init__() + + assert type(kernel_size) == tuple or type(kernel_size) == int + + kernel_size = kernel_size if type(kernel_size) == tuple else (kernel_size, kernel_size) + self.stride = stride if type(stride) == tuple else (stride, stride) + self.padding = padding + + self.offset_conv = nn.Conv2d(in_channels, + 2 * kernel_size[0] * kernel_size[1], + kernel_size=kernel_size, + stride=stride, + padding=self.padding, + bias=True) + + nn.init.constant_(self.offset_conv.weight, 0.) + nn.init.constant_(self.offset_conv.bias, 0.) + + self.modulator_conv = nn.Conv2d(in_channels, + 1 * kernel_size[0] * kernel_size[1], + kernel_size=kernel_size, + stride=stride, + padding=self.padding, + bias=True) + + nn.init.constant_(self.modulator_conv.weight, 0.) + nn.init.constant_(self.modulator_conv.bias, 0.) + + self.regular_conv = nn.Conv2d(in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=self.padding, + bias=bias) + + def forward(self, x): + #h, w = x.shape[2:] + #max_offset = max(h, w)/4. + + offset = self.offset_conv(x)#.clamp(-max_offset, max_offset) + modulator = 2. * torch.sigmoid(self.modulator_conv(x)) + + x = deform_conv2d( + input=x, + offset=offset, + weight=self.regular_conv.weight, + bias=self.regular_conv.bias, + padding=self.padding, + mask=modulator, + stride=self.stride, + ) + return x + + + + +### utils.py + +import torch.nn as nn + + +def build_act_layer(act_layer): + if act_layer == 'ReLU': + return nn.ReLU(inplace=True) + elif act_layer == 'SiLU': + return nn.SiLU(inplace=True) + elif act_layer == 'GELU': + return nn.GELU() + + raise NotImplementedError(f'build_act_layer does not support {act_layer}') + + +def build_norm_layer(dim, + norm_layer, + in_format='channels_last', + out_format='channels_last', + eps=1e-6): + layers = [] + if norm_layer == 'BN': + if in_format == 'channels_last': + layers.append(to_channels_first()) + layers.append(nn.BatchNorm2d(dim)) + if out_format == 'channels_last': + layers.append(to_channels_last()) + elif norm_layer == 'LN': + if in_format == 'channels_first': + layers.append(to_channels_last()) + layers.append(nn.LayerNorm(dim, eps=eps)) + if out_format == 'channels_first': + layers.append(to_channels_first()) + else: + raise NotImplementedError( + f'build_norm_layer does not support {norm_layer}') + return nn.Sequential(*layers) + + +class to_channels_first(nn.Module): + + def __init__(self): + super().__init__() + + def forward(self, x): + return x.permute(0, 3, 1, 2) + + +class to_channels_last(nn.Module): + + def __init__(self): + super().__init__() + + def forward(self, x): + return x.permute(0, 2, 3, 1) + + + +### dataset.py + +_class_labels_TR_sorted = ( + 'Airplane, Ant, Antenna, Archery, Axe, BabyCarriage, Bag, BalanceBeam, Balcony, Balloon, Basket, BasketballHoop, Beatle, Bed, Bee, Bench, Bicycle, ' + 'BicycleFrame, BicycleStand, Boat, Bonsai, BoomLift, Bridge, BunkBed, Butterfly, Button, Cable, CableLift, Cage, Camcorder, Cannon, Canoe, Car, ' + 'CarParkDropArm, Carriage, Cart, Caterpillar, CeilingLamp, Centipede, Chair, Clip, Clock, Clothes, CoatHanger, Comb, ConcretePumpTruck, Crack, Crane, ' + 'Cup, DentalChair, Desk, DeskChair, Diagram, DishRack, DoorHandle, Dragonfish, Dragonfly, Drum, Earphone, Easel, ElectricIron, Excavator, Eyeglasses, ' + 'Fan, Fence, Fencing, FerrisWheel, FireExtinguisher, Fishing, Flag, FloorLamp, Forklift, GasStation, Gate, Gear, Goal, Golf, GymEquipment, Hammock, ' + 'Handcart, Handcraft, Handrail, HangGlider, Harp, Harvester, Headset, Helicopter, Helmet, Hook, HorizontalBar, Hydrovalve, IroningTable, Jewelry, Key, ' + 'KidsPlayground, Kitchenware, Kite, Knife, Ladder, LaundryRack, Lightning, Lobster, Locust, Machine, MachineGun, MagazineRack, Mantis, Medal, MemorialArchway, ' + 'Microphone, Missile, MobileHolder, Monitor, Mosquito, Motorcycle, MovingTrolley, Mower, MusicPlayer, MusicStand, ObservationTower, Octopus, OilWell, ' + 'OlympicLogo, OperatingTable, OutdoorFitnessEquipment, Parachute, Pavilion, Piano, Pipe, PlowHarrow, PoleVault, Punchbag, Rack, Racket, Rifle, Ring, Robot, ' + 'RockClimbing, Rope, Sailboat, Satellite, Scaffold, Scale, Scissor, Scooter, Sculpture, Seadragon, Seahorse, Seal, SewingMachine, Ship, Shoe, ShoppingCart, ' + 'ShoppingTrolley, Shower, Shrimp, Signboard, Skateboarding, Skeleton, Skiing, Spade, SpeedBoat, Spider, Spoon, Stair, Stand, Stationary, SteeringWheel, ' + 'Stethoscope, Stool, Stove, StreetLamp, SweetStand, Swing, Sword, TV, Table, TableChair, TableLamp, TableTennis, Tank, Tapeline, Teapot, Telescope, Tent, ' + 'TobaccoPipe, Toy, Tractor, TrafficLight, TrafficSign, Trampoline, TransmissionTower, Tree, Tricycle, TrimmerCover, Tripod, Trombone, Truck, Trumpet, Tuba, ' + 'UAV, Umbrella, UnevenBars, UtilityPole, VacuumCleaner, Violin, Wakesurfing, Watch, WaterTower, WateringPot, Well, WellLid, Wheel, Wheelchair, WindTurbine, Windmill, WineGlass, WireWhisk, Yacht' +) +class_labels_TR_sorted = _class_labels_TR_sorted.split(', ') + + +### models/backbones/build_backbones.py + +import torch +import torch.nn as nn +from collections import OrderedDict +from torchvision.models import vgg16, vgg16_bn, VGG16_Weights, VGG16_BN_Weights, resnet50, ResNet50_Weights +# from models.pvt_v2 import pvt_v2_b0, pvt_v2_b1, pvt_v2_b2, pvt_v2_b5 +# from models.swin_v1 import swin_v1_t, swin_v1_s, swin_v1_b, swin_v1_l +# from config import Config + + +config = Config() + +def build_backbone(bb_name, pretrained=True, params_settings=''): + if bb_name == 'vgg16': + bb_net = list(vgg16(pretrained=VGG16_Weights.DEFAULT if pretrained else None).children())[0] + bb = nn.Sequential(OrderedDict({'conv1': bb_net[:4], 'conv2': bb_net[4:9], 'conv3': bb_net[9:16], 'conv4': bb_net[16:23]})) + elif bb_name == 'vgg16bn': + bb_net = list(vgg16_bn(pretrained=VGG16_BN_Weights.DEFAULT if pretrained else None).children())[0] + bb = nn.Sequential(OrderedDict({'conv1': bb_net[:6], 'conv2': bb_net[6:13], 'conv3': bb_net[13:23], 'conv4': bb_net[23:33]})) + elif bb_name == 'resnet50': + bb_net = list(resnet50(pretrained=ResNet50_Weights.DEFAULT if pretrained else None).children()) + bb = nn.Sequential(OrderedDict({'conv1': nn.Sequential(*bb_net[0:3]), 'conv2': bb_net[4], 'conv3': bb_net[5], 'conv4': bb_net[6]})) + else: + bb = eval('{}({})'.format(bb_name, params_settings)) + if pretrained: + bb = load_weights(bb, bb_name) + return bb + +def load_weights(model, model_name): + save_model = torch.load(config.weights[model_name], map_location='cpu') + model_dict = model.state_dict() + state_dict = {k: v if v.size() == model_dict[k].size() else model_dict[k] for k, v in save_model.items() if k in model_dict.keys()} + # to ignore the weights with mismatched size when I modify the backbone itself. + if not state_dict: + save_model_keys = list(save_model.keys()) + sub_item = save_model_keys[0] if len(save_model_keys) == 1 else None + state_dict = {k: v if v.size() == model_dict[k].size() else model_dict[k] for k, v in save_model[sub_item].items() if k in model_dict.keys()} + if not state_dict or not sub_item: + print('Weights are not successully loaded. Check the state dict of weights file.') + return None + else: + print('Found correct weights in the "{}" item of loaded state_dict.'.format(sub_item)) + model_dict.update(state_dict) + model.load_state_dict(model_dict) + return model + + + +### models/modules/decoder_blocks.py + +import torch +import torch.nn as nn +# from models.aspp import ASPP, ASPPDeformable +# from config import Config + + +# config = Config() + + +class BasicDecBlk(nn.Module): + def __init__(self, in_channels=64, out_channels=64, inter_channels=64): + super(BasicDecBlk, self).__init__() + inter_channels = in_channels // 4 if config.dec_channels_inter == 'adap' else 64 + self.conv_in = nn.Conv2d(in_channels, inter_channels, 3, 1, padding=1) + self.relu_in = nn.ReLU(inplace=True) + if config.dec_att == 'ASPP': + self.dec_att = ASPP(in_channels=inter_channels) + elif config.dec_att == 'ASPPDeformable': + self.dec_att = ASPPDeformable(in_channels=inter_channels) + self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, padding=1) + self.bn_in = nn.BatchNorm2d(inter_channels) if config.batch_size > 1 else nn.Identity() + self.bn_out = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity() + + def forward(self, x): + x = self.conv_in(x) + x = self.bn_in(x) + x = self.relu_in(x) + if hasattr(self, 'dec_att'): + x = self.dec_att(x) + x = self.conv_out(x) + x = self.bn_out(x) + return x + + +class ResBlk(nn.Module): + def __init__(self, in_channels=64, out_channels=None, inter_channels=64): + super(ResBlk, self).__init__() + if out_channels is None: + out_channels = in_channels + inter_channels = in_channels // 4 if config.dec_channels_inter == 'adap' else 64 + + self.conv_in = nn.Conv2d(in_channels, inter_channels, 3, 1, padding=1) + self.bn_in = nn.BatchNorm2d(inter_channels) if config.batch_size > 1 else nn.Identity() + self.relu_in = nn.ReLU(inplace=True) + + if config.dec_att == 'ASPP': + self.dec_att = ASPP(in_channels=inter_channels) + elif config.dec_att == 'ASPPDeformable': + self.dec_att = ASPPDeformable(in_channels=inter_channels) + + self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, padding=1) + self.bn_out = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity() + + self.conv_resi = nn.Conv2d(in_channels, out_channels, 1, 1, 0) + + def forward(self, x): + _x = self.conv_resi(x) + x = self.conv_in(x) + x = self.bn_in(x) + x = self.relu_in(x) + if hasattr(self, 'dec_att'): + x = self.dec_att(x) + x = self.conv_out(x) + x = self.bn_out(x) + return x + _x + + + +### models/modules/lateral_blocks.py + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from functools import partial + +# from config import Config + + +# config = Config() + + +class BasicLatBlk(nn.Module): + def __init__(self, in_channels=64, out_channels=64, inter_channels=64): + super(BasicLatBlk, self).__init__() + inter_channels = in_channels // 4 if config.dec_channels_inter == 'adap' else 64 + self.conv = nn.Conv2d(in_channels, out_channels, 1, 1, 0) + + def forward(self, x): + x = self.conv(x) + return x + + + +### models/modules/aspp.py + +import torch +import torch.nn as nn +import torch.nn.functional as F +# from models.deform_conv import DeformableConv2d +# from config import Config + + +# config = Config() + + +class _ASPPModule(nn.Module): + def __init__(self, in_channels, planes, kernel_size, padding, dilation): + super(_ASPPModule, self).__init__() + self.atrous_conv = nn.Conv2d(in_channels, planes, kernel_size=kernel_size, + stride=1, padding=padding, dilation=dilation, bias=False) + self.bn = nn.BatchNorm2d(planes) if config.batch_size > 1 else nn.Identity() + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + x = self.atrous_conv(x) + x = self.bn(x) + + return self.relu(x) + + +class ASPP(nn.Module): + def __init__(self, in_channels=64, out_channels=None, output_stride=16): + super(ASPP, self).__init__() + self.down_scale = 1 + if out_channels is None: + out_channels = in_channels + self.in_channelster = 256 // self.down_scale + if output_stride == 16: + dilations = [1, 6, 12, 18] + elif output_stride == 8: + dilations = [1, 12, 24, 36] + else: + raise NotImplementedError + + self.aspp1 = _ASPPModule(in_channels, self.in_channelster, 1, padding=0, dilation=dilations[0]) + self.aspp2 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[1], dilation=dilations[1]) + self.aspp3 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[2], dilation=dilations[2]) + self.aspp4 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[3], dilation=dilations[3]) + + self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)), + nn.Conv2d(in_channels, self.in_channelster, 1, stride=1, bias=False), + nn.BatchNorm2d(self.in_channelster) if config.batch_size > 1 else nn.Identity(), + nn.ReLU(inplace=True)) + self.conv1 = nn.Conv2d(self.in_channelster * 5, out_channels, 1, bias=False) + self.bn1 = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity() + self.relu = nn.ReLU(inplace=True) + self.dropout = nn.Dropout(0.5) + + def forward(self, x): + x1 = self.aspp1(x) + x2 = self.aspp2(x) + x3 = self.aspp3(x) + x4 = self.aspp4(x) + x5 = self.global_avg_pool(x) + x5 = F.interpolate(x5, size=x1.size()[2:], mode='bilinear', align_corners=True) + x = torch.cat((x1, x2, x3, x4, x5), dim=1) + + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + + return self.dropout(x) + + +##################### Deformable +class _ASPPModuleDeformable(nn.Module): + def __init__(self, in_channels, planes, kernel_size, padding): + super(_ASPPModuleDeformable, self).__init__() + self.atrous_conv = DeformableConv2d(in_channels, planes, kernel_size=kernel_size, + stride=1, padding=padding, bias=False) + self.bn = nn.BatchNorm2d(planes) if config.batch_size > 1 else nn.Identity() + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + x = self.atrous_conv(x) + x = self.bn(x) + + return self.relu(x) + + +class ASPPDeformable(nn.Module): + def __init__(self, in_channels, out_channels=None, parallel_block_sizes=[1, 3, 7]): + super(ASPPDeformable, self).__init__() + self.down_scale = 1 + if out_channels is None: + out_channels = in_channels + self.in_channelster = 256 // self.down_scale + + self.aspp1 = _ASPPModuleDeformable(in_channels, self.in_channelster, 1, padding=0) + self.aspp_deforms = nn.ModuleList([ + _ASPPModuleDeformable(in_channels, self.in_channelster, conv_size, padding=int(conv_size//2)) for conv_size in parallel_block_sizes + ]) + + self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)), + nn.Conv2d(in_channels, self.in_channelster, 1, stride=1, bias=False), + nn.BatchNorm2d(self.in_channelster) if config.batch_size > 1 else nn.Identity(), + nn.ReLU(inplace=True)) + self.conv1 = nn.Conv2d(self.in_channelster * (2 + len(self.aspp_deforms)), out_channels, 1, bias=False) + self.bn1 = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity() + self.relu = nn.ReLU(inplace=True) + self.dropout = nn.Dropout(0.5) + + def forward(self, x): + x1 = self.aspp1(x) + x_aspp_deforms = [aspp_deform(x) for aspp_deform in self.aspp_deforms] + x5 = self.global_avg_pool(x) + x5 = F.interpolate(x5, size=x1.size()[2:], mode='bilinear', align_corners=True) + x = torch.cat((x1, *x_aspp_deforms, x5), dim=1) + + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + + return self.dropout(x) + + + +### models/refinement/refiner.py + +import torch +import torch.nn as nn +from collections import OrderedDict +import torch +import torch.nn as nn +import torch.nn.functional as F +from torchvision.models import vgg16, vgg16_bn +from torchvision.models import resnet50 + +# from config import Config +# from dataset import class_labels_TR_sorted +# from models.build_backbone import build_backbone +# from models.decoder_blocks import BasicDecBlk +# from models.lateral_blocks import BasicLatBlk +# from models.ing import * +# from models.stem_layer import StemLayer + + +class RefinerPVTInChannels4(nn.Module): + def __init__(self, in_channels=3+1): + super(RefinerPVTInChannels4, self).__init__() + self.config = Config() + self.epoch = 1 + self.bb = build_backbone(self.config.bb, params_settings='in_channels=4') + + lateral_channels_in_collection = { + 'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64], + 'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64], + 'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192], + } + channels = lateral_channels_in_collection[self.config.bb] + self.squeeze_module = BasicDecBlk(channels[0], channels[0]) + + self.decoder = Decoder(channels) + + if 0: + for key, value in self.named_parameters(): + if 'bb.' in key: + value.requires_grad = False + + def forward(self, x): + if isinstance(x, list): + x = torch.cat(x, dim=1) + ########## Encoder ########## + if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']: + x1 = self.bb.conv1(x) + x2 = self.bb.conv2(x1) + x3 = self.bb.conv3(x2) + x4 = self.bb.conv4(x3) + else: + x1, x2, x3, x4 = self.bb(x) + + x4 = self.squeeze_module(x4) + + ########## Decoder ########## + + features = [x, x1, x2, x3, x4] + scaled_preds = self.decoder(features) + + return scaled_preds + + +class Refiner(nn.Module): + def __init__(self, in_channels=3+1): + super(Refiner, self).__init__() + self.config = Config() + self.epoch = 1 + self.stem_layer = StemLayer(in_channels=in_channels, inter_channels=48, out_channels=3, norm_layer='BN' if self.config.batch_size > 1 else 'LN') + self.bb = build_backbone(self.config.bb) + + lateral_channels_in_collection = { + 'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64], + 'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64], + 'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192], + } + channels = lateral_channels_in_collection[self.config.bb] + self.squeeze_module = BasicDecBlk(channels[0], channels[0]) + + self.decoder = Decoder(channels) + + if 0: + for key, value in self.named_parameters(): + if 'bb.' in key: + value.requires_grad = False + + def forward(self, x): + if isinstance(x, list): + x = torch.cat(x, dim=1) + x = self.stem_layer(x) + ########## Encoder ########## + if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']: + x1 = self.bb.conv1(x) + x2 = self.bb.conv2(x1) + x3 = self.bb.conv3(x2) + x4 = self.bb.conv4(x3) + else: + x1, x2, x3, x4 = self.bb(x) + + x4 = self.squeeze_module(x4) + + ########## Decoder ########## + + features = [x, x1, x2, x3, x4] + scaled_preds = self.decoder(features) + + return scaled_preds + + +class Decoder(nn.Module): + def __init__(self, channels): + super(Decoder, self).__init__() + self.config = Config() + DecoderBlock = eval('BasicDecBlk') + LateralBlock = eval('BasicLatBlk') + + self.decoder_block4 = DecoderBlock(channels[0], channels[1]) + self.decoder_block3 = DecoderBlock(channels[1], channels[2]) + self.decoder_block2 = DecoderBlock(channels[2], channels[3]) + self.decoder_block1 = DecoderBlock(channels[3], channels[3]//2) + + self.lateral_block4 = LateralBlock(channels[1], channels[1]) + self.lateral_block3 = LateralBlock(channels[2], channels[2]) + self.lateral_block2 = LateralBlock(channels[3], channels[3]) + + if self.config.ms_supervision: + self.conv_ms_spvn_4 = nn.Conv2d(channels[1], 1, 1, 1, 0) + self.conv_ms_spvn_3 = nn.Conv2d(channels[2], 1, 1, 1, 0) + self.conv_ms_spvn_2 = nn.Conv2d(channels[3], 1, 1, 1, 0) + self.conv_out1 = nn.Sequential(nn.Conv2d(channels[3]//2, 1, 1, 1, 0)) + + def forward(self, features): + x, x1, x2, x3, x4 = features + outs = [] + p4 = self.decoder_block4(x4) + _p4 = F.interpolate(p4, size=x3.shape[2:], mode='bilinear', align_corners=True) + _p3 = _p4 + self.lateral_block4(x3) + + p3 = self.decoder_block3(_p3) + _p3 = F.interpolate(p3, size=x2.shape[2:], mode='bilinear', align_corners=True) + _p2 = _p3 + self.lateral_block3(x2) + + p2 = self.decoder_block2(_p2) + _p2 = F.interpolate(p2, size=x1.shape[2:], mode='bilinear', align_corners=True) + _p1 = _p2 + self.lateral_block2(x1) + + _p1 = self.decoder_block1(_p1) + _p1 = F.interpolate(_p1, size=x.shape[2:], mode='bilinear', align_corners=True) + p1_out = self.conv_out1(_p1) + + if self.config.ms_supervision: + outs.append(self.conv_ms_spvn_4(p4)) + outs.append(self.conv_ms_spvn_3(p3)) + outs.append(self.conv_ms_spvn_2(p2)) + outs.append(p1_out) + return outs + + +class RefUNet(nn.Module): + # Refinement + def __init__(self, in_channels=3+1): + super(RefUNet, self).__init__() + self.encoder_1 = nn.Sequential( + nn.Conv2d(in_channels, 64, 3, 1, 1), + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.encoder_2 = nn.Sequential( + nn.MaxPool2d(2, 2, ceil_mode=True), + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.encoder_3 = nn.Sequential( + nn.MaxPool2d(2, 2, ceil_mode=True), + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.encoder_4 = nn.Sequential( + nn.MaxPool2d(2, 2, ceil_mode=True), + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.pool4 = nn.MaxPool2d(2, 2, ceil_mode=True) + ##### + self.decoder_5 = nn.Sequential( + nn.Conv2d(64, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + ##### + self.decoder_4 = nn.Sequential( + nn.Conv2d(128, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.decoder_3 = nn.Sequential( + nn.Conv2d(128, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.decoder_2 = nn.Sequential( + nn.Conv2d(128, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.decoder_1 = nn.Sequential( + nn.Conv2d(128, 64, 3, 1, 1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True) + ) + + self.conv_d0 = nn.Conv2d(64, 1, 3, 1, 1) + + self.upscore2 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) + + def forward(self, x): + outs = [] + if isinstance(x, list): + x = torch.cat(x, dim=1) + hx = x + + hx1 = self.encoder_1(hx) + hx2 = self.encoder_2(hx1) + hx3 = self.encoder_3(hx2) + hx4 = self.encoder_4(hx3) + + hx = self.decoder_5(self.pool4(hx4)) + hx = torch.cat((self.upscore2(hx), hx4), 1) + + d4 = self.decoder_4(hx) + hx = torch.cat((self.upscore2(d4), hx3), 1) + + d3 = self.decoder_3(hx) + hx = torch.cat((self.upscore2(d3), hx2), 1) + + d2 = self.decoder_2(hx) + hx = torch.cat((self.upscore2(d2), hx1), 1) + + d1 = self.decoder_1(hx) + + x = self.conv_d0(d1) + outs.append(x) + return outs + + + +### models/stem_layer.py + +import torch.nn as nn +# from utils import build_act_layer, build_norm_layer + + +class StemLayer(nn.Module): + r""" Stem layer of InternImage + Args: + in_channels (int): number of input channels + out_channels (int): number of output channels + act_layer (str): activation layer + norm_layer (str): normalization layer + """ + + def __init__(self, + in_channels=3+1, + inter_channels=48, + out_channels=96, + act_layer='GELU', + norm_layer='BN'): + super().__init__() + self.conv1 = nn.Conv2d(in_channels, + inter_channels, + kernel_size=3, + stride=1, + padding=1) + self.norm1 = build_norm_layer( + inter_channels, norm_layer, 'channels_first', 'channels_first' + ) + self.act = build_act_layer(act_layer) + self.conv2 = nn.Conv2d(inter_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1) + self.norm2 = build_norm_layer( + out_channels, norm_layer, 'channels_first', 'channels_first' + ) + + def forward(self, x): + x = self.conv1(x) + x = self.norm1(x) + x = self.act(x) + x = self.conv2(x) + x = self.norm2(x) + return x + + +### models/birefnet.py + +import torch +import torch.nn as nn +import torch.nn.functional as F +from kornia.filters import laplacian +from transformers import PreTrainedModel + +# from config import Config +# from dataset import class_labels_TR_sorted +# from models.build_backbone import build_backbone +# from models.decoder_blocks import BasicDecBlk, ResBlk, HierarAttDecBlk +# from models.lateral_blocks import BasicLatBlk +# from models.aspp import ASPP, ASPPDeformable +# from models.ing import * +# from models.refiner import Refiner, RefinerPVTInChannels4, RefUNet +# from models.stem_layer import StemLayer +from .BiRefNet_config import BiRefNetConfig + + +class BiRefNet( + PreTrainedModel +): + config_class = BiRefNetConfig + def __init__(self, bb_pretrained=True, config=BiRefNetConfig()): + super(BiRefNet, self).__init__(config) + bb_pretrained = config.bb_pretrained + self.config = Config() + self.epoch = 1 + self.bb = build_backbone(self.config.bb, pretrained=bb_pretrained) + + channels = self.config.lateral_channels_in_collection + + if self.config.auxiliary_classification: + self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) + self.cls_head = nn.Sequential( + nn.Linear(channels[0], len(class_labels_TR_sorted)) + ) + + if self.config.squeeze_block: + self.squeeze_module = nn.Sequential(*[ + eval(self.config.squeeze_block.split('_x')[0])(channels[0]+sum(self.config.cxt), channels[0]) + for _ in range(eval(self.config.squeeze_block.split('_x')[1])) + ]) + + self.decoder = Decoder(channels) + + if self.config.ender: + self.dec_end = nn.Sequential( + nn.Conv2d(1, 16, 3, 1, 1), + nn.Conv2d(16, 1, 3, 1, 1), + nn.ReLU(inplace=True), + ) + + # refine patch-level segmentation + if self.config.refine: + if self.config.refine == 'itself': + self.stem_layer = StemLayer(in_channels=3+1, inter_channels=48, out_channels=3, norm_layer='BN' if self.config.batch_size > 1 else 'LN') + else: + self.refiner = eval('{}({})'.format(self.config.refine, 'in_channels=3+1')) + + if self.config.freeze_bb: + # Freeze the backbone... + print(self.named_parameters()) + for key, value in self.named_parameters(): + if 'bb.' in key and 'refiner.' not in key: + value.requires_grad = False + + def forward_enc(self, x): + if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']: + x1 = self.bb.conv1(x); x2 = self.bb.conv2(x1); x3 = self.bb.conv3(x2); x4 = self.bb.conv4(x3) + else: + x1, x2, x3, x4 = self.bb(x) + if self.config.mul_scl_ipt == 'cat': + B, C, H, W = x.shape + x1_, x2_, x3_, x4_ = self.bb(F.interpolate(x, size=(H//2, W//2), mode='bilinear', align_corners=True)) + x1 = torch.cat([x1, F.interpolate(x1_, size=x1.shape[2:], mode='bilinear', align_corners=True)], dim=1) + x2 = torch.cat([x2, F.interpolate(x2_, size=x2.shape[2:], mode='bilinear', align_corners=True)], dim=1) + x3 = torch.cat([x3, F.interpolate(x3_, size=x3.shape[2:], mode='bilinear', align_corners=True)], dim=1) + x4 = torch.cat([x4, F.interpolate(x4_, size=x4.shape[2:], mode='bilinear', align_corners=True)], dim=1) + elif self.config.mul_scl_ipt == 'add': + B, C, H, W = x.shape + x1_, x2_, x3_, x4_ = self.bb(F.interpolate(x, size=(H//2, W//2), mode='bilinear', align_corners=True)) + x1 = x1 + F.interpolate(x1_, size=x1.shape[2:], mode='bilinear', align_corners=True) + x2 = x2 + F.interpolate(x2_, size=x2.shape[2:], mode='bilinear', align_corners=True) + x3 = x3 + F.interpolate(x3_, size=x3.shape[2:], mode='bilinear', align_corners=True) + x4 = x4 + F.interpolate(x4_, size=x4.shape[2:], mode='bilinear', align_corners=True) + class_preds = self.cls_head(self.avgpool(x4).view(x4.shape[0], -1)) if self.training and self.config.auxiliary_classification else None + if self.config.cxt: + x4 = torch.cat( + ( + *[ + F.interpolate(x1, size=x4.shape[2:], mode='bilinear', align_corners=True), + F.interpolate(x2, size=x4.shape[2:], mode='bilinear', align_corners=True), + F.interpolate(x3, size=x4.shape[2:], mode='bilinear', align_corners=True), + ][-len(self.config.cxt):], + x4 + ), + dim=1 + ) + return (x1, x2, x3, x4), class_preds + + def forward_ori(self, x): + ########## Encoder ########## + (x1, x2, x3, x4), class_preds = self.forward_enc(x) + if self.config.squeeze_block: + x4 = self.squeeze_module(x4) + ########## Decoder ########## + features = [x, x1, x2, x3, x4] + if self.training and self.config.out_ref: + features.append(laplacian(torch.mean(x, dim=1).unsqueeze(1), kernel_size=5)) + scaled_preds = self.decoder(features) + return scaled_preds, class_preds + + def forward(self, x): + scaled_preds, class_preds = self.forward_ori(x) + class_preds_lst = [class_preds] + return [scaled_preds, class_preds_lst] if self.training else scaled_preds + + +class Decoder(nn.Module): + def __init__(self, channels): + super(Decoder, self).__init__() + self.config = Config() + DecoderBlock = eval(self.config.dec_blk) + LateralBlock = eval(self.config.lat_blk) + + if self.config.dec_ipt: + self.split = self.config.dec_ipt_split + N_dec_ipt = 64 + DBlock = SimpleConvs + ic = 64 + ipt_cha_opt = 1 + self.ipt_blk5 = DBlock(2**10*3 if self.split else 3, [N_dec_ipt, channels[0]//8][ipt_cha_opt], inter_channels=ic) + self.ipt_blk4 = DBlock(2**8*3 if self.split else 3, [N_dec_ipt, channels[0]//8][ipt_cha_opt], inter_channels=ic) + self.ipt_blk3 = DBlock(2**6*3 if self.split else 3, [N_dec_ipt, channels[1]//8][ipt_cha_opt], inter_channels=ic) + self.ipt_blk2 = DBlock(2**4*3 if self.split else 3, [N_dec_ipt, channels[2]//8][ipt_cha_opt], inter_channels=ic) + self.ipt_blk1 = DBlock(2**0*3 if self.split else 3, [N_dec_ipt, channels[3]//8][ipt_cha_opt], inter_channels=ic) + else: + self.split = None + + self.decoder_block4 = DecoderBlock(channels[0]+([N_dec_ipt, channels[0]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[1]) + self.decoder_block3 = DecoderBlock(channels[1]+([N_dec_ipt, channels[0]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[2]) + self.decoder_block2 = DecoderBlock(channels[2]+([N_dec_ipt, channels[1]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[3]) + self.decoder_block1 = DecoderBlock(channels[3]+([N_dec_ipt, channels[2]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[3]//2) + self.conv_out1 = nn.Sequential(nn.Conv2d(channels[3]//2+([N_dec_ipt, channels[3]//8][ipt_cha_opt] if self.config.dec_ipt else 0), 1, 1, 1, 0)) + + self.lateral_block4 = LateralBlock(channels[1], channels[1]) + self.lateral_block3 = LateralBlock(channels[2], channels[2]) + self.lateral_block2 = LateralBlock(channels[3], channels[3]) + + if self.config.ms_supervision: + self.conv_ms_spvn_4 = nn.Conv2d(channels[1], 1, 1, 1, 0) + self.conv_ms_spvn_3 = nn.Conv2d(channels[2], 1, 1, 1, 0) + self.conv_ms_spvn_2 = nn.Conv2d(channels[3], 1, 1, 1, 0) + + if self.config.out_ref: + _N = 16 + self.gdt_convs_4 = nn.Sequential(nn.Conv2d(channels[1], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True)) + self.gdt_convs_3 = nn.Sequential(nn.Conv2d(channels[2], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True)) + self.gdt_convs_2 = nn.Sequential(nn.Conv2d(channels[3], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True)) + + self.gdt_convs_pred_4 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + self.gdt_convs_pred_3 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + self.gdt_convs_pred_2 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + + self.gdt_convs_attn_4 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + self.gdt_convs_attn_3 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + self.gdt_convs_attn_2 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0)) + + def get_patches_batch(self, x, p): + _size_h, _size_w = p.shape[2:] + patches_batch = [] + for idx in range(x.shape[0]): + columns_x = torch.split(x[idx], split_size_or_sections=_size_w, dim=-1) + patches_x = [] + for column_x in columns_x: + patches_x += [p.unsqueeze(0) for p in torch.split(column_x, split_size_or_sections=_size_h, dim=-2)] + patch_sample = torch.cat(patches_x, dim=1) + patches_batch.append(patch_sample) + return torch.cat(patches_batch, dim=0) + + def forward(self, features): + if self.training and self.config.out_ref: + outs_gdt_pred = [] + outs_gdt_label = [] + x, x1, x2, x3, x4, gdt_gt = features + else: + x, x1, x2, x3, x4 = features + outs = [] + + if self.config.dec_ipt: + patches_batch = self.get_patches_batch(x, x4) if self.split else x + x4 = torch.cat((x4, self.ipt_blk5(F.interpolate(patches_batch, size=x4.shape[2:], mode='bilinear', align_corners=True))), 1) + p4 = self.decoder_block4(x4) + m4 = self.conv_ms_spvn_4(p4) if self.config.ms_supervision else None + if self.config.out_ref: + p4_gdt = self.gdt_convs_4(p4) + if self.training: + # >> GT: + m4_dia = m4 + gdt_label_main_4 = gdt_gt * F.interpolate(m4_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True) + outs_gdt_label.append(gdt_label_main_4) + # >> Pred: + gdt_pred_4 = self.gdt_convs_pred_4(p4_gdt) + outs_gdt_pred.append(gdt_pred_4) + gdt_attn_4 = self.gdt_convs_attn_4(p4_gdt).sigmoid() + # >> Finally: + p4 = p4 * gdt_attn_4 + _p4 = F.interpolate(p4, size=x3.shape[2:], mode='bilinear', align_corners=True) + _p3 = _p4 + self.lateral_block4(x3) + + if self.config.dec_ipt: + patches_batch = self.get_patches_batch(x, _p3) if self.split else x + _p3 = torch.cat((_p3, self.ipt_blk4(F.interpolate(patches_batch, size=x3.shape[2:], mode='bilinear', align_corners=True))), 1) + p3 = self.decoder_block3(_p3) + m3 = self.conv_ms_spvn_3(p3) if self.config.ms_supervision else None + if self.config.out_ref: + p3_gdt = self.gdt_convs_3(p3) + if self.training: + # >> GT: + # m3 --dilation--> m3_dia + # G_3^gt * m3_dia --> G_3^m, which is the label of gradient + m3_dia = m3 + gdt_label_main_3 = gdt_gt * F.interpolate(m3_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True) + outs_gdt_label.append(gdt_label_main_3) + # >> Pred: + # p3 --conv--BN--> F_3^G, where F_3^G predicts the \hat{G_3} with xx + # F_3^G --sigmoid--> A_3^G + gdt_pred_3 = self.gdt_convs_pred_3(p3_gdt) + outs_gdt_pred.append(gdt_pred_3) + gdt_attn_3 = self.gdt_convs_attn_3(p3_gdt).sigmoid() + # >> Finally: + # p3 = p3 * A_3^G + p3 = p3 * gdt_attn_3 + _p3 = F.interpolate(p3, size=x2.shape[2:], mode='bilinear', align_corners=True) + _p2 = _p3 + self.lateral_block3(x2) + + if self.config.dec_ipt: + patches_batch = self.get_patches_batch(x, _p2) if self.split else x + _p2 = torch.cat((_p2, self.ipt_blk3(F.interpolate(patches_batch, size=x2.shape[2:], mode='bilinear', align_corners=True))), 1) + p2 = self.decoder_block2(_p2) + m2 = self.conv_ms_spvn_2(p2) if self.config.ms_supervision else None + if self.config.out_ref: + p2_gdt = self.gdt_convs_2(p2) + if self.training: + # >> GT: + m2_dia = m2 + gdt_label_main_2 = gdt_gt * F.interpolate(m2_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True) + outs_gdt_label.append(gdt_label_main_2) + # >> Pred: + gdt_pred_2 = self.gdt_convs_pred_2(p2_gdt) + outs_gdt_pred.append(gdt_pred_2) + gdt_attn_2 = self.gdt_convs_attn_2(p2_gdt).sigmoid() + # >> Finally: + p2 = p2 * gdt_attn_2 + _p2 = F.interpolate(p2, size=x1.shape[2:], mode='bilinear', align_corners=True) + _p1 = _p2 + self.lateral_block2(x1) + + if self.config.dec_ipt: + patches_batch = self.get_patches_batch(x, _p1) if self.split else x + _p1 = torch.cat((_p1, self.ipt_blk2(F.interpolate(patches_batch, size=x1.shape[2:], mode='bilinear', align_corners=True))), 1) + _p1 = self.decoder_block1(_p1) + _p1 = F.interpolate(_p1, size=x.shape[2:], mode='bilinear', align_corners=True) + + if self.config.dec_ipt: + patches_batch = self.get_patches_batch(x, _p1) if self.split else x + _p1 = torch.cat((_p1, self.ipt_blk1(F.interpolate(patches_batch, size=x.shape[2:], mode='bilinear', align_corners=True))), 1) + p1_out = self.conv_out1(_p1) + + if self.config.ms_supervision: + outs.append(m4) + outs.append(m3) + outs.append(m2) + outs.append(p1_out) + return outs if not (self.config.out_ref and self.training) else ([outs_gdt_pred, outs_gdt_label], outs) + + +class SimpleConvs(nn.Module): + def __init__( + self, in_channels: int, out_channels: int, inter_channels=64 + ) -> None: + super().__init__() + self.conv1 = nn.Conv2d(in_channels, inter_channels, 3, 1, 1) + self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, 1) + + def forward(self, x): + return self.conv_out(self.conv1(x)) diff --git a/models/RMBG/RMBG-2.0/config.json b/models/RMBG/RMBG-2.0/config.json new file mode 100644 index 0000000000000000000000000000000000000000..06d8fa9d7f2f4c6f1cf0dc6e7bfd194153176a42 --- /dev/null +++ b/models/RMBG/RMBG-2.0/config.json @@ -0,0 +1,20 @@ +{ + "_name_or_path": "ZhengPeng7/BiRefNet", + "architectures": [ + "BiRefNet" + ], + "auto_map": { + "AutoConfig": "BiRefNet_config.BiRefNetConfig", + "AutoModelForImageSegmentation": "birefnet.BiRefNet" + }, + "custom_pipelines": { + "image-segmentation": { + "pt": [ + "AutoModelForImageSegmentation" + ], + "tf": [], + "type": "image" + } + }, + "bb_pretrained": false +} \ No newline at end of file diff --git a/models/RMBG/RMBG-2.0/gitattributes b/models/RMBG/RMBG-2.0/gitattributes new file mode 100644 index 0000000000000000000000000000000000000000..a6344aac8c09253b3b630fb776ae94478aa0275b --- /dev/null +++ b/models/RMBG/RMBG-2.0/gitattributes @@ -0,0 +1,35 @@ +*.7z filter=lfs diff=lfs merge=lfs -text +*.arrow filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +*.bz2 filter=lfs diff=lfs merge=lfs -text +*.ckpt filter=lfs diff=lfs merge=lfs -text +*.ftz filter=lfs diff=lfs merge=lfs -text +*.gz filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text +*.joblib filter=lfs diff=lfs merge=lfs -text +*.lfs.* filter=lfs diff=lfs merge=lfs -text +*.mlmodel filter=lfs diff=lfs merge=lfs -text +*.model filter=lfs diff=lfs merge=lfs -text +*.msgpack filter=lfs diff=lfs merge=lfs -text +*.npy filter=lfs diff=lfs merge=lfs -text +*.npz filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.ot filter=lfs diff=lfs merge=lfs -text +*.parquet filter=lfs diff=lfs merge=lfs -text +*.pb filter=lfs diff=lfs merge=lfs -text +*.pickle filter=lfs diff=lfs merge=lfs -text +*.pkl filter=lfs diff=lfs merge=lfs -text +*.pt filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.rar filter=lfs diff=lfs merge=lfs -text +*.safetensors filter=lfs diff=lfs merge=lfs -text +saved_model/**/* filter=lfs diff=lfs merge=lfs -text +*.tar.* filter=lfs diff=lfs merge=lfs -text +*.tar filter=lfs diff=lfs merge=lfs -text +*.tflite filter=lfs diff=lfs merge=lfs -text +*.tgz filter=lfs diff=lfs merge=lfs -text +*.wasm filter=lfs diff=lfs merge=lfs -text +*.xz filter=lfs diff=lfs merge=lfs -text +*.zip filter=lfs diff=lfs merge=lfs -text +*.zst filter=lfs diff=lfs merge=lfs -text +*tfevents* filter=lfs diff=lfs merge=lfs -text diff --git a/models/RMBG/RMBG-2.0/model.safetensors b/models/RMBG/RMBG-2.0/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..121c871665f9d734c36f5be12556984df33667b1 --- /dev/null +++ b/models/RMBG/RMBG-2.0/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:566ed80c3d95f87ada6864d4cbe2290a1c5eb1c7bb0b123e984f60f76b02c3a7 +size 884878856 diff --git a/models/RMBG/SAM/Moonlit Serenade.mp3 b/models/RMBG/SAM/Moonlit Serenade.mp3 new file mode 100644 index 0000000000000000000000000000000000000000..6bee15f8da2207be7ed24c47f6b008fce0f750f6 --- /dev/null +++ b/models/RMBG/SAM/Moonlit Serenade.mp3 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1a219a66c5d8590a19a4dc7f8bc661170fd6afeab9db8e642c474bb3571ede4a +size 4812946 diff --git a/models/RMBG/SAM/README.md b/models/RMBG/SAM/README.md new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/models/RMBG/SAM/gitattributes b/models/RMBG/SAM/gitattributes new file mode 100644 index 0000000000000000000000000000000000000000..a6344aac8c09253b3b630fb776ae94478aa0275b --- /dev/null +++ b/models/RMBG/SAM/gitattributes @@ -0,0 +1,35 @@ +*.7z filter=lfs diff=lfs merge=lfs -text +*.arrow filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +*.bz2 filter=lfs diff=lfs merge=lfs -text +*.ckpt filter=lfs diff=lfs merge=lfs -text +*.ftz filter=lfs diff=lfs merge=lfs -text +*.gz filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text +*.joblib filter=lfs diff=lfs merge=lfs -text +*.lfs.* filter=lfs diff=lfs merge=lfs -text +*.mlmodel filter=lfs diff=lfs merge=lfs -text +*.model filter=lfs diff=lfs merge=lfs -text +*.msgpack filter=lfs diff=lfs merge=lfs -text +*.npy filter=lfs diff=lfs merge=lfs -text +*.npz filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.ot filter=lfs diff=lfs merge=lfs -text +*.parquet filter=lfs diff=lfs merge=lfs -text +*.pb filter=lfs diff=lfs merge=lfs -text +*.pickle filter=lfs diff=lfs merge=lfs -text +*.pkl filter=lfs diff=lfs merge=lfs -text +*.pt filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.rar filter=lfs diff=lfs merge=lfs -text +*.safetensors filter=lfs diff=lfs merge=lfs -text +saved_model/**/* filter=lfs diff=lfs merge=lfs -text +*.tar.* filter=lfs diff=lfs merge=lfs -text +*.tar filter=lfs diff=lfs merge=lfs -text +*.tflite filter=lfs diff=lfs merge=lfs -text +*.tgz filter=lfs diff=lfs merge=lfs -text +*.wasm filter=lfs diff=lfs merge=lfs -text +*.xz filter=lfs diff=lfs merge=lfs -text +*.zip filter=lfs diff=lfs merge=lfs -text +*.zst filter=lfs diff=lfs merge=lfs -text +*tfevents* filter=lfs diff=lfs merge=lfs -text diff --git a/models/RMBG/SAM/mobile_sam.pt b/models/RMBG/SAM/mobile_sam.pt new file mode 100644 index 0000000000000000000000000000000000000000..59e427ae6522bbb4c640cd8dad4271526d8cdd6c --- /dev/null +++ b/models/RMBG/SAM/mobile_sam.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6dbb90523a35330fedd7f1d3dfc66f995213d81b29a5ca8108dbcdd4e37d6c2f +size 40728226 diff --git a/models/RMBG/SAM/mobile_sam.safetensors b/models/RMBG/SAM/mobile_sam.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..63d5dd32e574fe792b5d173ef25a91a84a9211c0 --- /dev/null +++ b/models/RMBG/SAM/mobile_sam.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:158f98ffe079d9a5f0c6ef6ed8a238ce2a61220754fa972d93156200d693c1bf +size 40613136 diff --git a/models/RMBG/SAM/sam_hq_vit_b.pth b/models/RMBG/SAM/sam_hq_vit_b.pth new file mode 100644 index 0000000000000000000000000000000000000000..358a965085541bd0c2424be5dd2734172b8f52f6 --- /dev/null +++ b/models/RMBG/SAM/sam_hq_vit_b.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:14a9d662cd6f5a9c2dba6d40ab0058d88d287e4a18fd6fdc6ad5fb1a3fdeaa57 +size 379335069 diff --git a/models/RMBG/SAM/sam_hq_vit_h.pth b/models/RMBG/SAM/sam_hq_vit_h.pth new file mode 100644 index 0000000000000000000000000000000000000000..e3d8fbbf24da52ed3ded801fc344786fb3c6fec0 --- /dev/null +++ b/models/RMBG/SAM/sam_hq_vit_h.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a7ac14a085326d9fa6199c8c698c4f0e7280afdbb974d2c4660ec60877b45e35 +size 2570940653 diff --git a/models/RMBG/SAM/sam_hq_vit_l.pth b/models/RMBG/SAM/sam_hq_vit_l.pth new file mode 100644 index 0000000000000000000000000000000000000000..694388067ca30da86a3706035f5c2a813207457f --- /dev/null +++ b/models/RMBG/SAM/sam_hq_vit_l.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e1a6c385d62bf005ded91a54d5ec55c985cfc4103ef89c08d90f39f04934c343 +size 1254865805 diff --git a/models/RMBG/SAM/sam_vit_b.pth b/models/RMBG/SAM/sam_vit_b.pth new file mode 100644 index 0000000000000000000000000000000000000000..ab7d111e57bd052a76fe669986560e3555e9c8f6 --- /dev/null +++ b/models/RMBG/SAM/sam_vit_b.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ec2df62732614e57411cdcf32a23ffdf28910380d03139ee0f4fcbe91eb8c912 +size 375042383 diff --git a/models/RMBG/SAM/sam_vit_h.pth b/models/RMBG/SAM/sam_vit_h.pth new file mode 100644 index 0000000000000000000000000000000000000000..8523acce9ddab1cf7e355628a08b1aab8ce08a72 --- /dev/null +++ b/models/RMBG/SAM/sam_vit_h.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a7bf3b02f3ebf1267aba913ff637d9a2d5c33d3173bb679e46d9f338c26f262e +size 2564550879 diff --git a/models/RMBG/SAM/sam_vit_l.pth b/models/RMBG/SAM/sam_vit_l.pth new file mode 100644 index 0000000000000000000000000000000000000000..87a638d6b789dd2b10fc7414a88dacc34a50769a --- /dev/null +++ b/models/RMBG/SAM/sam_vit_l.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3adcc4315b642a4d2101128f611684e8734c41232a17c648ed1693702a49a622 +size 1249524607 diff --git a/models/RMBG/grounding-dino/GroundingDINO_SwinB.cfg.py b/models/RMBG/grounding-dino/GroundingDINO_SwinB.cfg.py new file mode 100644 index 0000000000000000000000000000000000000000..f490c4bbd598a35de43d36ceafcbd769e7ff21bf --- /dev/null +++ b/models/RMBG/grounding-dino/GroundingDINO_SwinB.cfg.py @@ -0,0 +1,43 @@ +batch_size = 1 +modelname = "groundingdino" +backbone = "swin_B_384_22k" +position_embedding = "sine" +pe_temperatureH = 20 +pe_temperatureW = 20 +return_interm_indices = [1, 2, 3] +backbone_freeze_keywords = None +enc_layers = 6 +dec_layers = 6 +pre_norm = False +dim_feedforward = 2048 +hidden_dim = 256 +dropout = 0.0 +nheads = 8 +num_queries = 900 +query_dim = 4 +num_patterns = 0 +num_feature_levels = 4 +enc_n_points = 4 +dec_n_points = 4 +two_stage_type = "standard" +two_stage_bbox_embed_share = False +two_stage_class_embed_share = False +transformer_activation = "relu" +dec_pred_bbox_embed_share = True +dn_box_noise_scale = 1.0 +dn_label_noise_ratio = 0.5 +dn_label_coef = 1.0 +dn_bbox_coef = 1.0 +embed_init_tgt = True +dn_labelbook_size = 2000 +max_text_len = 256 +text_encoder_type = "bert-base-uncased" +use_text_enhancer = True +use_fusion_layer = True +use_checkpoint = True +use_transformer_ckpt = True +use_text_cross_attention = True +text_dropout = 0.0 +fusion_dropout = 0.0 +fusion_droppath = 0.1 +sub_sentence_present = True diff --git a/models/RMBG/grounding-dino/GroundingDINO_SwinT_OGC.cfg.py b/models/RMBG/grounding-dino/GroundingDINO_SwinT_OGC.cfg.py new file mode 100644 index 0000000000000000000000000000000000000000..9158d5f6260ec74bded95377d382387430d7cd70 --- /dev/null +++ b/models/RMBG/grounding-dino/GroundingDINO_SwinT_OGC.cfg.py @@ -0,0 +1,43 @@ +batch_size = 1 +modelname = "groundingdino" +backbone = "swin_T_224_1k" +position_embedding = "sine" +pe_temperatureH = 20 +pe_temperatureW = 20 +return_interm_indices = [1, 2, 3] +backbone_freeze_keywords = None +enc_layers = 6 +dec_layers = 6 +pre_norm = False +dim_feedforward = 2048 +hidden_dim = 256 +dropout = 0.0 +nheads = 8 +num_queries = 900 +query_dim = 4 +num_patterns = 0 +num_feature_levels = 4 +enc_n_points = 4 +dec_n_points = 4 +two_stage_type = "standard" +two_stage_bbox_embed_share = False +two_stage_class_embed_share = False +transformer_activation = "relu" +dec_pred_bbox_embed_share = True +dn_box_noise_scale = 1.0 +dn_label_noise_ratio = 0.5 +dn_label_coef = 1.0 +dn_bbox_coef = 1.0 +embed_init_tgt = True +dn_labelbook_size = 2000 +max_text_len = 256 +text_encoder_type = "bert-base-uncased" +use_text_enhancer = True +use_fusion_layer = True +use_checkpoint = True +use_transformer_ckpt = True +use_text_cross_attention = True +text_dropout = 0.0 +fusion_dropout = 0.0 +fusion_droppath = 0.1 +sub_sentence_present = True diff --git a/models/RMBG/grounding-dino/gitattributes b/models/RMBG/grounding-dino/gitattributes new file mode 100644 index 0000000000000000000000000000000000000000..a6344aac8c09253b3b630fb776ae94478aa0275b --- /dev/null +++ b/models/RMBG/grounding-dino/gitattributes @@ -0,0 +1,35 @@ +*.7z filter=lfs diff=lfs merge=lfs -text +*.arrow filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +*.bz2 filter=lfs diff=lfs merge=lfs -text +*.ckpt filter=lfs diff=lfs merge=lfs -text +*.ftz filter=lfs diff=lfs merge=lfs -text +*.gz filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text +*.joblib filter=lfs diff=lfs merge=lfs -text +*.lfs.* filter=lfs diff=lfs merge=lfs -text +*.mlmodel filter=lfs diff=lfs merge=lfs -text +*.model filter=lfs diff=lfs merge=lfs -text +*.msgpack filter=lfs diff=lfs merge=lfs -text +*.npy filter=lfs diff=lfs merge=lfs -text +*.npz filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.ot filter=lfs diff=lfs merge=lfs -text +*.parquet filter=lfs diff=lfs merge=lfs -text +*.pb filter=lfs diff=lfs merge=lfs -text +*.pickle filter=lfs diff=lfs merge=lfs -text +*.pkl filter=lfs diff=lfs merge=lfs -text +*.pt filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.rar filter=lfs diff=lfs merge=lfs -text +*.safetensors filter=lfs diff=lfs merge=lfs -text +saved_model/**/* filter=lfs diff=lfs merge=lfs -text +*.tar.* filter=lfs diff=lfs merge=lfs -text +*.tar filter=lfs diff=lfs merge=lfs -text +*.tflite filter=lfs diff=lfs merge=lfs -text +*.tgz filter=lfs diff=lfs merge=lfs -text +*.wasm filter=lfs diff=lfs merge=lfs -text +*.xz filter=lfs diff=lfs merge=lfs -text +*.zip filter=lfs diff=lfs merge=lfs -text +*.zst filter=lfs diff=lfs merge=lfs -text +*tfevents* filter=lfs diff=lfs merge=lfs -text diff --git a/models/RMBG/grounding-dino/groundingdino_swinb_cogcoor.pth b/models/RMBG/grounding-dino/groundingdino_swinb_cogcoor.pth new file mode 100644 index 0000000000000000000000000000000000000000..5f82241224563baebffdc579f20f01a5d9a76c64 --- /dev/null +++ b/models/RMBG/grounding-dino/groundingdino_swinb_cogcoor.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:46270f7a822e6906b655b729c90613e48929d0f2bb8b9b76fd10a856f3ac6ab7 +size 938057991 diff --git a/models/RMBG/grounding-dino/groundingdino_swint_ogc.pth b/models/RMBG/grounding-dino/groundingdino_swint_ogc.pth new file mode 100644 index 0000000000000000000000000000000000000000..5cdf6bcd10d491abf170a78eca4fcebf76aa791a --- /dev/null +++ b/models/RMBG/grounding-dino/groundingdino_swint_ogc.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3b3ca2563c77c69f651d7bd133e97139c186df06231157a64c507099c52bc799 +size 693997677 diff --git a/models/RMBG/segformer_clothes/config.json b/models/RMBG/segformer_clothes/config.json new file mode 100644 index 0000000000000000000000000000000000000000..8352c4562bb0e1f72767dcb170ad6f3f56007836 --- /dev/null +++ b/models/RMBG/segformer_clothes/config.json @@ -0,0 +1,110 @@ +{ + "_name_or_path": "nvidia/mit-b3", + "architectures": [ + "SegformerForSemanticSegmentation" + ], + "attention_probs_dropout_prob": 0.0, + "classifier_dropout_prob": 0.1, + "decoder_hidden_size": 768, + "depths": [ + 3, + 4, + 18, + 3 + ], + "downsampling_rates": [ + 1, + 4, + 8, + 16 + ], + "drop_path_rate": 0.1, + "hidden_act": "gelu", + "hidden_dropout_prob": 0.0, + "hidden_sizes": [ + 64, + 128, + 320, + 512 + ], + "id2label": { + "0": "Background", + "1": "Hat", + "10": "Right-shoe", + "11": "Face", + "12": "Left-leg", + "13": "Right-leg", + "14": "Left-arm", + "15": "Right-arm", + "16": "Bag", + "17": "Scarf", + "2": "Hair", + "3": "Sunglasses", + "4": "Upper-clothes", + "5": "Skirt", + "6": "Pants", + "7": "Dress", + "8": "Belt", + "9": "Left-shoe" + }, + "image_size": 224, + "initializer_range": 0.02, + "label2id": { + "Background": "0", + "Bag": "16", + "Belt": "8", + "Dress": "7", + "Face": "11", + "Hair": "2", + "Hat": "1", + "Left-arm": "14", + "Left-leg": "12", + "Left-shoe": "9", + "Pants": "6", + "Right-arm": "15", + "Right-leg": "13", + "Right-shoe": "10", + "Scarf": "17", + "Skirt": "5", + "Sunglasses": "3", + "Upper-clothes": "4" + }, + "layer_norm_eps": 1e-06, + "mlp_ratios": [ + 4, + 4, + 4, + 4 + ], + "model_type": "segformer", + "num_attention_heads": [ + 1, + 2, + 5, + 8 + ], + "num_channels": 3, + "num_encoder_blocks": 4, + "patch_sizes": [ + 7, + 3, + 3, + 3 + ], + "reshape_last_stage": true, + "semantic_loss_ignore_index": 255, + "sr_ratios": [ + 8, + 4, + 2, + 1 + ], + "strides": [ + 4, + 2, + 2, + 2 + ], + "torch_dtype": "float32", + "transformers_version": "4.38.1" +} diff --git a/models/RMBG/segformer_clothes/gitattributes b/models/RMBG/segformer_clothes/gitattributes new file mode 100644 index 0000000000000000000000000000000000000000..a6344aac8c09253b3b630fb776ae94478aa0275b --- /dev/null +++ b/models/RMBG/segformer_clothes/gitattributes @@ -0,0 +1,35 @@ +*.7z filter=lfs diff=lfs merge=lfs -text +*.arrow filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +*.bz2 filter=lfs diff=lfs merge=lfs -text +*.ckpt filter=lfs diff=lfs merge=lfs -text +*.ftz filter=lfs diff=lfs merge=lfs -text +*.gz filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text +*.joblib filter=lfs diff=lfs merge=lfs -text +*.lfs.* filter=lfs diff=lfs merge=lfs -text +*.mlmodel filter=lfs diff=lfs merge=lfs -text +*.model filter=lfs diff=lfs merge=lfs -text +*.msgpack filter=lfs diff=lfs merge=lfs -text +*.npy filter=lfs diff=lfs merge=lfs -text +*.npz filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.ot filter=lfs diff=lfs merge=lfs -text +*.parquet filter=lfs diff=lfs merge=lfs -text +*.pb filter=lfs diff=lfs merge=lfs -text +*.pickle filter=lfs diff=lfs merge=lfs -text +*.pkl filter=lfs diff=lfs merge=lfs -text +*.pt filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.rar filter=lfs diff=lfs merge=lfs -text +*.safetensors filter=lfs diff=lfs merge=lfs -text +saved_model/**/* filter=lfs diff=lfs merge=lfs -text +*.tar.* filter=lfs diff=lfs merge=lfs -text +*.tar filter=lfs diff=lfs merge=lfs -text +*.tflite filter=lfs diff=lfs merge=lfs -text +*.tgz filter=lfs diff=lfs merge=lfs -text +*.wasm filter=lfs diff=lfs merge=lfs -text +*.xz filter=lfs diff=lfs merge=lfs -text +*.zip filter=lfs diff=lfs merge=lfs -text +*.zst filter=lfs diff=lfs merge=lfs -text +*tfevents* filter=lfs diff=lfs merge=lfs -text diff --git a/models/RMBG/segformer_clothes/model.safetensors b/models/RMBG/segformer_clothes/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..fb5506b45fc6b66daa05add78f13412a4776d0cd --- /dev/null +++ b/models/RMBG/segformer_clothes/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f70ae566c5773fb335796ebaa8acc924ac25eb97222c2b2967d44d2fc11568e6 +size 189029000 diff --git a/models/RMBG/segformer_clothes/preprocessor_config.json b/models/RMBG/segformer_clothes/preprocessor_config.json new file mode 100644 index 0000000000000000000000000000000000000000..b2340cf4e53b37fda4f5b92d28f11c0f33c3d0fd --- /dev/null +++ b/models/RMBG/segformer_clothes/preprocessor_config.json @@ -0,0 +1,23 @@ +{ + "do_normalize": true, + "do_reduce_labels": false, + "do_rescale": true, + "do_resize": true, + "image_mean": [ + 0.485, + 0.456, + 0.406 + ], + "image_processor_type": "SegformerImageProcessor", + "image_std": [ + 0.229, + 0.224, + 0.225 + ], + "resample": 2, + "rescale_factor": 0.00392156862745098, + "size": { + "height": 512, + "width": 512 + } +} diff --git a/models/RMBG/segformer_fashion/config.json b/models/RMBG/segformer_fashion/config.json new file mode 100644 index 0000000000000000000000000000000000000000..8fb2769b75445217cf9bc737f56abd740f313f1d --- /dev/null +++ b/models/RMBG/segformer_fashion/config.json @@ -0,0 +1,168 @@ +{ + "_name_or_path": "nvidia/mit-b3", + "architectures": [ + "SegformerForSemanticSegmentation" + ], + "attention_probs_dropout_prob": 0.0, + "classifier_dropout_prob": 0.1, + "decoder_hidden_size": 768, + "depths": [ + 3, + 4, + 18, + 3 + ], + "downsampling_rates": [ + 1, + 4, + 8, + 16 + ], + "drop_path_rate": 0.1, + "hidden_act": "gelu", + "hidden_dropout_prob": 0.0, + "hidden_sizes": [ + 64, + 128, + 320, + 512 + ], + "id2label": { + "0": "unlabelled", + "1": "shirt, blouse", + "2": "top, t-shirt, sweatshirt", + "3": "sweater", + "4": "cardigan", + "5": "jacket", + "6": "vest", + "7": "pants", + "8": "shorts", + "9": "skirt", + "10": "coat", + "11": "dress", + "12": "jumpsuit", + "13": "cape", + "14": "glasses", + "15": "hat", + "16": "headband, head covering, hair accessory", + "17": "tie", + "18": "glove", + "19": "watch", + "20": "belt", + "21": "leg warmer", + "22": "tights, stockings", + "23": "sock", + "24": "shoe", + "25": "bag, wallet", + "26": "scarf", + "27": "umbrella", + "28": "hood", + "29": "collar", + "30": "lapel", + "31": "epaulette", + "32": "sleeve", + "33": "pocket", + "34": "neckline", + "35": "buckle", + "36": "zipper", + "37": "applique", + "38": "bead", + "39": "bow", + "40": "flower", + "41": "fringe", + "42": "ribbon", + "43": "rivet", + "44": "ruffle", + "45": "sequin", + "46": "tassel" + }, + "image_size": 224, + "initializer_range": 0.02, + "label2id": { + "applique": 37, + "bag, wallet": 25, + "bead": 38, + "belt": 20, + "bow": 39, + "buckle": 35, + "cape": 13, + "cardigan": 4, + "coat": 10, + "collar": 29, + "dress": 11, + "epaulette": 31, + "flower": 40, + "fringe": 41, + "glasses": 14, + "glove": 18, + "hat": 15, + "headband, head covering, hair accessory": 16, + "hood": 28, + "jacket": 5, + "jumpsuit": 12, + "lapel": 30, + "leg warmer": 21, + "neckline": 34, + "pants": 7, + "pocket": 33, + "ribbon": 42, + "rivet": 43, + "ruffle": 44, + "scarf": 26, + "sequin": 45, + "shirt, blouse": 1, + "shoe": 24, + "shorts": 8, + "skirt": 9, + "sleeve": 32, + "sock": 23, + "sweater": 3, + "tassel": 46, + "tie": 17, + "tights, stockings": 22, + "top, t-shirt, sweatshirt": 2, + "umbrella": 27, + "unlabelled": 0, + "vest": 6, + "watch": 19, + "zipper": 36 + }, + "layer_norm_eps": 1e-06, + "mlp_ratios": [ + 4, + 4, + 4, + 4 + ], + "model_type": "segformer", + "num_attention_heads": [ + 1, + 2, + 5, + 8 + ], + "num_channels": 3, + "num_encoder_blocks": 4, + "patch_sizes": [ + 7, + 3, + 3, + 3 + ], + "reshape_last_stage": true, + "semantic_loss_ignore_index": 255, + "sr_ratios": [ + 8, + 4, + 2, + 1 + ], + "strides": [ + 4, + 2, + 2, + 2 + ], + "torch_dtype": "float32", + "transformers_version": "4.30.0" +} diff --git a/models/RMBG/segformer_fashion/gitattributes b/models/RMBG/segformer_fashion/gitattributes new file mode 100644 index 0000000000000000000000000000000000000000..a6344aac8c09253b3b630fb776ae94478aa0275b --- /dev/null +++ b/models/RMBG/segformer_fashion/gitattributes @@ -0,0 +1,35 @@ +*.7z filter=lfs diff=lfs merge=lfs -text +*.arrow filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +*.bz2 filter=lfs diff=lfs merge=lfs -text +*.ckpt filter=lfs diff=lfs merge=lfs -text +*.ftz filter=lfs diff=lfs merge=lfs -text +*.gz filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text +*.joblib filter=lfs diff=lfs merge=lfs -text +*.lfs.* filter=lfs diff=lfs merge=lfs -text +*.mlmodel filter=lfs diff=lfs merge=lfs -text +*.model filter=lfs diff=lfs merge=lfs -text +*.msgpack filter=lfs diff=lfs merge=lfs -text +*.npy filter=lfs diff=lfs merge=lfs -text +*.npz filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.ot filter=lfs diff=lfs merge=lfs -text +*.parquet filter=lfs diff=lfs merge=lfs -text +*.pb filter=lfs diff=lfs merge=lfs -text +*.pickle filter=lfs diff=lfs merge=lfs -text +*.pkl filter=lfs diff=lfs merge=lfs -text +*.pt filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.rar filter=lfs diff=lfs merge=lfs -text +*.safetensors filter=lfs diff=lfs merge=lfs -text +saved_model/**/* filter=lfs diff=lfs merge=lfs -text +*.tar.* filter=lfs diff=lfs merge=lfs -text +*.tar filter=lfs diff=lfs merge=lfs -text +*.tflite filter=lfs diff=lfs merge=lfs -text +*.tgz filter=lfs diff=lfs merge=lfs -text +*.wasm filter=lfs diff=lfs merge=lfs -text +*.xz filter=lfs diff=lfs merge=lfs -text +*.zip filter=lfs diff=lfs merge=lfs -text +*.zst filter=lfs diff=lfs merge=lfs -text +*tfevents* filter=lfs diff=lfs merge=lfs -text diff --git a/models/RMBG/segformer_fashion/model.safetensors b/models/RMBG/segformer_fashion/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..0644c245dfa3ba121e7330a75d9a908ce8e47e60 --- /dev/null +++ b/models/RMBG/segformer_fashion/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f3f5b30179f1480d329224d089f6d286580142c2b12846d08de814a48a81f42f +size 189118204 diff --git a/models/RMBG/segformer_fashion/preprocessor_config.json b/models/RMBG/segformer_fashion/preprocessor_config.json new file mode 100644 index 0000000000000000000000000000000000000000..b2340cf4e53b37fda4f5b92d28f11c0f33c3d0fd --- /dev/null +++ b/models/RMBG/segformer_fashion/preprocessor_config.json @@ -0,0 +1,23 @@ +{ + "do_normalize": true, + "do_reduce_labels": false, + "do_rescale": true, + "do_resize": true, + "image_mean": [ + 0.485, + 0.456, + 0.406 + ], + "image_processor_type": "SegformerImageProcessor", + "image_std": [ + 0.229, + 0.224, + 0.225 + ], + "resample": 2, + "rescale_factor": 0.00392156862745098, + "size": { + "height": 512, + "width": 512 + } +}