137 lines
7.1 KiB
Python
137 lines
7.1 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from .models.swin import SwinTransformer
|
|
from .models.resnet import ResNet
|
|
from .models.t2t_vit import T2t_vit_t_14
|
|
from .models.EfficientNet import EfficientNet
|
|
from .multistage_fusion import decoder
|
|
from .multilevel_interaction import MultilevelInteractionBlock
|
|
|
|
|
|
class M3Net(nn.Module):
|
|
r""" Multilevel, Mixed and Multistage Attention Network for Salient Object Detection.
|
|
|
|
Args:
|
|
embed_dim (int): Dimension for attention. Default 384
|
|
dim (int): Patch embedding dimension. Default 96
|
|
img_size (int): Input image size. Default 224
|
|
method (string): Backbone used as the encoder.
|
|
"""
|
|
|
|
def __init__(self, embed_dim=384, dim=96, img_size=224, method='M3Net-S'):
|
|
super(M3Net, self).__init__()
|
|
self.img_size = img_size
|
|
self.feature_dims = []
|
|
self.method = method
|
|
self.dim = dim
|
|
if method == 'M3Net-S':
|
|
self.encoder = SwinTransformer(img_size=img_size,
|
|
embed_dim=128,
|
|
depths=[2, 2, 18, 2],
|
|
num_heads=[4, 8, 16, 32],
|
|
window_size=img_size // 32)
|
|
self.proj1 = nn.Linear(128, dim)
|
|
self.proj2 = nn.Linear(256, dim * 2)
|
|
self.proj3 = nn.Linear(512, dim * 4)
|
|
self.proj4 = nn.Linear(1024, dim * 8)
|
|
|
|
self.interact1 = MultilevelInteractionBlock(dim=dim * 4, dim1=dim * 8, embed_dim=embed_dim, num_heads=4,
|
|
mlp_ratio=3)
|
|
self.interact2 = MultilevelInteractionBlock(dim=dim * 2, dim1=dim * 4, dim2=dim * 8, embed_dim=embed_dim,
|
|
num_heads=2, mlp_ratio=3)
|
|
self.interact3 = MultilevelInteractionBlock(dim=dim, dim1=dim * 2, dim2=dim * 4, embed_dim=embed_dim,
|
|
num_heads=1, mlp_ratio=3)
|
|
feature_dims = [dim, dim * 2, dim * 4]
|
|
|
|
elif method == 'M3Net-R':
|
|
self.encoder = ResNet()
|
|
self.proj1 = nn.Conv2d(256, dim, 1)
|
|
self.proj2 = nn.Conv2d(512, dim * 2, 1)
|
|
self.proj3 = nn.Conv2d(1024, dim * 4, 1)
|
|
self.proj4 = nn.Conv2d(2048, dim * 8, 1)
|
|
|
|
self.interact1 = MultilevelInteractionBlock(dim=dim * 4, dim1=dim * 8, embed_dim=embed_dim, num_heads=4,
|
|
mlp_ratio=3)
|
|
self.interact2 = MultilevelInteractionBlock(dim=dim * 2, dim1=dim * 4, dim2=dim * 8, embed_dim=embed_dim,
|
|
num_heads=2, mlp_ratio=3)
|
|
self.interact3 = MultilevelInteractionBlock(dim=dim, dim1=dim * 2, dim2=dim * 4, embed_dim=embed_dim,
|
|
num_heads=1, mlp_ratio=3)
|
|
feature_dims = [dim, dim * 2, dim * 4]
|
|
|
|
|
|
elif method == 'M3Net-T':
|
|
self.encoder = T2t_vit_t_14(pretrained=False)
|
|
self.interact2 = MultilevelInteractionBlock(dim=dim, dim1=embed_dim, embed_dim=embed_dim, num_heads=2,
|
|
mlp_ratio=3)
|
|
self.interact3 = MultilevelInteractionBlock(dim=dim, dim1=dim, dim2=embed_dim, embed_dim=embed_dim,
|
|
num_heads=1, mlp_ratio=3)
|
|
feature_dims = [dim, dim, embed_dim]
|
|
|
|
elif method == 'M3Net-E':
|
|
self.encoder = EfficientNet.from_name(f'efficientnet-b7')
|
|
self.proj1 = nn.Conv2d(48, dim, 1)
|
|
self.proj2 = nn.Conv2d(80, dim * 2, 1)
|
|
self.proj3 = nn.Conv2d(224, dim * 4, 1)
|
|
self.proj4 = nn.Conv2d(640, dim * 8, 1)
|
|
|
|
self.interact1 = MultilevelInteractionBlock(dim=dim * 4, dim1=dim * 8, embed_dim=embed_dim, num_heads=4,
|
|
mlp_ratio=3)
|
|
self.interact2 = MultilevelInteractionBlock(dim=dim * 2, dim1=dim * 4, dim2=dim * 8, embed_dim=embed_dim,
|
|
num_heads=2, mlp_ratio=3)
|
|
self.interact3 = MultilevelInteractionBlock(dim=dim, dim1=dim * 2, dim2=dim * 4, embed_dim=embed_dim,
|
|
num_heads=1, mlp_ratio=3)
|
|
feature_dims = [dim, dim * 2, dim * 4]
|
|
|
|
self.decoder = decoder(embed_dim=embed_dim, dims=feature_dims, img_size=img_size, mlp_ratio=1)
|
|
|
|
def forward(self, x):
|
|
fea = self.encoder(x)
|
|
if self.method == 'M3Net-S':
|
|
fea_1_4, fea_1_8, fea_1_16, fea_1_32 = fea
|
|
fea_1_4 = self.proj1(fea_1_4)
|
|
fea_1_8 = self.proj2(fea_1_8)
|
|
fea_1_16 = self.proj3(fea_1_16)
|
|
fea_1_32 = self.proj4(fea_1_32)
|
|
fea_1_16_ = self.interact1(fea_1_16, fea_1_32)
|
|
fea_1_8_ = self.interact2(fea_1_8, fea_1_16_, fea_1_32)
|
|
fea_1_4_ = self.interact3(fea_1_4, fea_1_8_, fea_1_16_)
|
|
elif self.method == 'M3Net-R':
|
|
fea_1_4, fea_1_8, fea_1_16, fea_1_32 = fea
|
|
B, _, _, _ = fea_1_4.shape
|
|
fea_1_4 = self.proj1(fea_1_4).reshape(B, self.dim, -1).transpose(1, 2)
|
|
fea_1_8 = self.proj2(fea_1_8).reshape(B, self.dim * 2, -1).transpose(1, 2)
|
|
fea_1_16 = self.proj3(fea_1_16).reshape(B, self.dim * 4, -1).transpose(1, 2)
|
|
fea_1_32 = self.proj4(fea_1_32).reshape(B, self.dim * 8, -1).transpose(1, 2)
|
|
fea_1_16_ = self.interact1(fea_1_16, fea_1_32)
|
|
fea_1_8_ = self.interact2(fea_1_8, fea_1_16_, fea_1_32)
|
|
fea_1_4_ = self.interact3(fea_1_4, fea_1_8_, fea_1_16_)
|
|
elif self.method == 'M3Net-T':
|
|
fea_1_4, fea_1_8, fea_1_16_ = fea
|
|
fea_1_8_ = self.interact2(fea_1_8, fea_1_16_)
|
|
fea_1_4_ = self.interact3(fea_1_4, fea_1_8_, fea_1_16_)
|
|
elif self.method == 'M3Net-E':
|
|
fea_1_4, fea_1_8, fea_1_16, fea_1_32 = fea
|
|
B, _, _, _ = fea_1_4.shape
|
|
fea_1_4 = self.proj1(fea_1_4).reshape(B, self.dim, -1).transpose(1, 2)
|
|
fea_1_8 = self.proj2(fea_1_8).reshape(B, self.dim * 2, -1).transpose(1, 2)
|
|
fea_1_16 = self.proj3(fea_1_16).reshape(B, self.dim * 4, -1).transpose(1, 2)
|
|
fea_1_32 = self.proj4(fea_1_32).reshape(B, self.dim * 8, -1).transpose(1, 2)
|
|
fea_1_16_ = self.interact1(fea_1_16, fea_1_32)
|
|
fea_1_8_ = self.interact2(fea_1_8, fea_1_16_, fea_1_32)
|
|
fea_1_4_ = self.interact3(fea_1_4, fea_1_8_, fea_1_16_)
|
|
mask = self.decoder([fea_1_16_, fea_1_8_, fea_1_4_])
|
|
return mask
|
|
|
|
def flops(self):
|
|
flops = 0
|
|
flops += self.encoder.flops()
|
|
N1 = self.img_size // 4 * self.img_size // 4
|
|
N2 = self.img_size // 8 * self.img_size // 8
|
|
N3 = self.img_size // 16 * self.img_size // 16
|
|
N4 = self.img_size // 32 * self.img_size // 32
|
|
flops += self.interact1.flops(N3, N4)
|
|
flops += self.interact2.flops(N2, N3, N4)
|
|
flops += self.interact3.flops(N1, N2, N3)
|
|
flops += self.decoder.flops()
|
|
return flops
|