def __init__(self, config=None):
super(U2NetBackbone, self).__init__()
self.cfg = Dict(config)
in_ch = self.cfg.get("in_ch", 3)
self.model_name = self.cfg.get("model_name", "small")
if self.model_name == "small":
cfg = [[in_ch, 16, 64], [64, 16, 64], [64, 16, 64], [64, 16, 64],
[64, 16, 64], [64, 16, 64]]
else:
cfg = [[in_ch, 32, 64], [64, 32, 128], [128, 64, 256],
[256, 128, 512], [512, 256, 512], [512, 256, 512]]
self.stage1 = RSU7(cfg[0][0], cfg[0][1], cfg[0][2])
self.pool12 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
self.stage2 = RSU6(cfg[1][0], cfg[1][1], cfg[1][2])
self.pool23 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
self.stage3 = RSU5(cfg[2][0], cfg[2][1], cfg[2][2])
self.pool34 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
self.stage4 = RSU4(cfg[3][0], cfg[3][1], cfg[3][2])
self.pool45 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
self.stage5 = RSU4F(cfg[4][0], cfg[4][1], cfg[4][2])
self.pool56 = nn.MaxPool2d(2, stride=2, ceil_mode=True)
self.stage6 = RSU4F(cfg[5][0], cfg[5][1], cfg[5][2])
def __init__(self, config=None):
super(U2NetHead, self).__init__()
self.cfg = Dict(config)
out_ch = self.cfg.get("out_ch", 3)
self.model_name = self.cfg.get("model_name", "small")
if self.model_name == "small":
cfg = [[128, 16, 64], [128, 16, 64], [128, 16, 64], [128, 16, 64],
[128, 16, 64]]
else:
cfg = [
[1024, 256, 512],
[1024, 128, 256],
[512, 64, 128],
[256, 32, 64],
[128, 16, 64],
]
# decoder
self.stage5d = RSU4F(cfg[0][0], cfg[0][1], cfg[0][2])
self.stage4d = RSU4(cfg[1][0], cfg[1][1], cfg[1][2])
self.stage3d = RSU5(cfg[2][0], cfg[2][1], cfg[2][2])
self.stage2d = RSU6(cfg[3][0], cfg[3][1], cfg[3][2])
self.stage1d = RSU7(cfg[4][0], cfg[4][1], cfg[4][2])
self.side1 = nn.Conv2d(cfg[4][2], out_ch, 3, padding=1) # 16
self.side2 = nn.Conv2d(cfg[3][2], out_ch, 3, padding=1) # 32
self.side3 = nn.Conv2d(cfg[2][2], out_ch, 3, padding=1) # 64
self.side4 = nn.Conv2d(cfg[1][2], out_ch, 3, padding=1) # 128
self.side5 = nn.Conv2d(cfg[0][2], out_ch, 3, padding=1) # 512
self.side6 = nn.Conv2d(cfg[0][2], out_ch, 3, padding=1) # 512
self.outconv = nn.Conv2d(6, out_ch, 1)
def __init__(self, config=None):
super(U2NetHeadV3, self).__init__()
self.cfg = Dict(config)
out_ch = self.cfg.get("out_ch", 1)
self.model_name = self.cfg.get("model_name", "small")
if self.model_name == "small":
cfg = [[80, 16, 64], [80, 16, 64], [80, 16, 64], [80, 16, 64]]
else:
cfg = [[32, 16, 32], [64, 32, 64], [128, 64, 128], [256, 128, 256],
[512, 512, 512]]
# decoder
self.stage1d = RSU7(cfg[0][0], cfg[0][1], cfg[0][2])
self.stage2d = RSU6(cfg[1][0], cfg[1][1], cfg[1][2])
self.stage3d = RSU5(cfg[2][0], cfg[2][1], cfg[2][2])
self.stage4d = RSU4(cfg[3][0], cfg[3][1], cfg[3][2])
# compress
self.c1 = nn.Conv2d(cfg[0][2] * 2, cfg[0][1], 1)
self.c2 = nn.Conv2d(cfg[1][2] * 2, cfg[1][1], 1)
self.c3 = nn.Conv2d(cfg[2][2] * 2, cfg[2][1], 1)
self.c4 = nn.Conv2d(cfg[3][2] * 2, cfg[3][1], 1)
self.se5 = SEBlock(cfg[4][2], cfg[4][1])
# side
self.side1 = nn.Conv2d(cfg[0][2], 4 * out_ch, 1) # 64
self.side2 = nn.Conv2d(cfg[1][2], 4 * out_ch, 1) # 128
self.side3 = nn.Conv2d(cfg[2][2], 4 * out_ch, 1) # 256
self.side4 = nn.Conv2d(cfg[3][2], 4 * out_ch, 1) # 512
# pixel shuffle
self.ps0 = nn.PixelShuffle(2)
self.ps1 = nn.PixelShuffle(2)
self.ps2 = nn.PixelShuffle(2)
self.ps3 = nn.PixelShuffle(2)
self.ps4 = nn.PixelShuffle(2)
self.ps5 = nn.PixelShuffle(2)
# attention
self.attn1 = DFMAtt(cfg[0][1], cfg[0][1], k=4)
self.attn2 = DFMAtt(cfg[1][1], cfg[1][1], k=4)
self.attn3 = DFMAtt(cfg[2][1], cfg[2][1], k=4)
self.attn4 = DFMAtt(cfg[3][1], cfg[3][1], k=4)
def __init__(self, config=None):
super(U2NetHeadV2, self).__init__()
self.cfg = Dict(config)
out_ch = self.cfg.get("out_ch", 1)
self.model_name = self.cfg.get("model_name", "small")
if self.model_name == "small":
cfg = [[80, 16, 64], [80, 16, 64], [80, 16, 64], [80, 16, 64]]
else:
cfg = [
[768, 128, 512],
[384, 64, 256],
[192, 32, 128],
[96, 16, 64],
]
# cfg = [
# [1024, 16, 64],
# [512, 16, 64],
# [256, 16, 64],
# [128, 16, 64],
# ]
# decoder
self.stage4d = RSU4(cfg[0][0], cfg[0][1], cfg[0][2])
self.stage3d = RSU5(cfg[1][0], cfg[1][1], cfg[1][2])
self.stage2d = RSU6(cfg[2][0], cfg[2][1], cfg[2][2])
self.stage1d = RSU7(cfg[3][0], cfg[3][1], cfg[3][2])
# side
self.side1 = nn.Conv2d(cfg[3][2], 4 * out_ch, 1) # 64
self.side2 = nn.Conv2d(cfg[2][2], 4 * out_ch, 1) # 128
self.side3 = nn.Conv2d(cfg[1][2], 4 * out_ch, 1) # 256
self.side4 = nn.Conv2d(cfg[0][2], 4 * out_ch, 1) # 512
# self.side1 = nn.Sequential(
# nn.Conv2d(cfg[3][2], cfg[3][2], kernel_size=1, stride=1, padding=0),
# nn.BatchNorm2d(cfg[3][2]),
# nn.ReLU(inplace=True),
# nn.Conv2d(cfg[3][2], 4 * out_ch, kernel_size=1, stride=1, padding=0, bias=True)
# )
# self.side2 = nn.Sequential(
# nn.Conv2d(cfg[2][2], cfg[2][2], kernel_size=1, stride=1, padding=0),
# nn.BatchNorm2d(cfg[2][2]),
# nn.ReLU(inplace=True),
# nn.Conv2d(cfg[2][2], 4 * out_ch, kernel_size=1, stride=1, padding=0, bias=True)
# )
# self.side3 = nn.Sequential(
# nn.Conv2d(cfg[1][2], cfg[1][2], kernel_size=1, stride=1, padding=0),
# nn.BatchNorm2d(cfg[1][2]),
# nn.ReLU(inplace=True),
# nn.Conv2d(cfg[1][2], 4 * out_ch, kernel_size=1, stride=1, padding=0, bias=True)
# )
# self.side4 = nn.Sequential(
# nn.Conv2d(cfg[0][2], cfg[0][2], kernel_size=1, stride=1, padding=0),
# nn.BatchNorm2d(cfg[0][2]),
# nn.ReLU(inplace=True),
# nn.Conv2d(cfg[0][2], 4 * out_ch, kernel_size=1, stride=1, padding=0, bias=True)
# )
# pixel shuffle
self.ps1 = nn.PixelShuffle(2)
self.ps2 = nn.PixelShuffle(2)
self.ps3 = nn.PixelShuffle(2)
self.ps4 = nn.PixelShuffle(2)
self.ps5 = nn.PixelShuffle(2)
# attention
self.attn1 = DFMAtt(cfg[3][2], cfg[3][2], k=1)
self.attn2 = DFMAtt(cfg[2][2], cfg[2][2], k=1)
self.attn3 = DFMAtt(cfg[1][2], cfg[1][2], k=1)
self.attn4 = DFMAtt(cfg[0][2], cfg[0][2], k=1)
self.attn5 = DFMAtt(cfg[0][2] * 2, cfg[0][2] * 2, k=1)
self.outconv = nn.Conv2d(4, out_ch, 1)