|
import torch |
|
import torch.nn as nn |
|
from backbones_unet.model.unet import Unet |
|
import torch.nn.functional as F |
|
from utils import * |
|
__all__ = ['UNext'] |
|
|
|
from timm.models.layers import DropPath, to_2tuple, trunc_normal_ |
|
import math |
|
|
|
class UNet(nn.Module): |
|
def __init__(self, in_channels = 3, out_channels = 1, init_features = 32, pretrained=True , back_bone=None): |
|
super().__init__() |
|
if back_bone is None: |
|
self.model = torch.hub.load( |
|
'mateuszbuda/brain-segmentation-pytorch', 'unet', in_channels=in_channels, out_channels=out_channels, |
|
init_features=init_features, pretrained=pretrained |
|
) |
|
else: |
|
self.model = UNet( |
|
in_channels= in_channels, |
|
out_channels= out_channels, |
|
backbone=back_bone |
|
) |
|
|
|
self.soft = nn.Softmax(dim =1) |
|
def forward(self, x, text_dummy): |
|
return self.soft(self.model(x)),0 |
|
|
|
|
|
def conv1x1(in_planes: int, out_planes: int, stride: int = 1) -> nn.Conv2d: |
|
"""1x1 convolution""" |
|
return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=1, bias=False) |
|
|
|
class shiftmlp(nn.Module): |
|
def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0., shift_size=5): |
|
super().__init__() |
|
out_features = out_features or in_features |
|
hidden_features = hidden_features or in_features |
|
self.dim = 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.shift_size = shift_size |
|
self.pad = shift_size // 2 |
|
|
|
|
|
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 |
|
|
|
xn = x.transpose(1, 2).view(B, C, H, W).contiguous() |
|
xn = F.pad(xn, (self.pad, self.pad, self.pad, self.pad) , "constant", 0) |
|
xs = torch.chunk(xn, self.shift_size, 1) |
|
x_shift = [torch.roll(x_c, shift, 2) for x_c, shift in zip(xs, range(-self.pad, self.pad+1))] |
|
x_cat = torch.cat(x_shift, 1) |
|
x_cat = torch.narrow(x_cat, 2, self.pad, H) |
|
x_s = torch.narrow(x_cat, 3, self.pad, W) |
|
|
|
|
|
x_s = x_s.reshape(B,C,H*W).contiguous() |
|
x_shift_r = x_s.transpose(1,2) |
|
|
|
|
|
x = self.fc1(x_shift_r) |
|
|
|
x = self.dwconv(x, H, W) |
|
x = self.act(x) |
|
x = self.drop(x) |
|
|
|
xn = x.transpose(1, 2).view(B, C, H, W).contiguous() |
|
xn = F.pad(xn, (self.pad, self.pad, self.pad, self.pad) , "constant", 0) |
|
xs = torch.chunk(xn, self.shift_size, 1) |
|
x_shift = [torch.roll(x_c, shift, 3) for x_c, shift in zip(xs, range(-self.pad, self.pad+1))] |
|
x_cat = torch.cat(x_shift, 1) |
|
x_cat = torch.narrow(x_cat, 2, self.pad, H) |
|
x_s = torch.narrow(x_cat, 3, self.pad, W) |
|
x_s = x_s.reshape(B,C,H*W).contiguous() |
|
x_shift_c = x_s.transpose(1,2) |
|
|
|
x = self.fc2(x_shift_c) |
|
x = self.drop(x) |
|
return x |
|
|
|
|
|
|
|
class shiftedBlock(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.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 = shiftmlp(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.mlp(self.norm2(x), H, W)) |
|
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) |
|
x = self.dwconv(x) |
|
x = x.flatten(2).transpose(1, 2) |
|
|
|
return x |
|
|
|
class OverlapPatchEmbed(nn.Module): |
|
""" Image to Patch Embedding |
|
""" |
|
|
|
def __init__(self, img_size=224, patch_size=7, stride=4, in_chans=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_chans, 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 UNext(nn.Module): |
|
|
|
|
|
|
|
def __init__(self, num_classes, input_channels=3, deep_supervision=False,img_size=256, patch_size=16, in_chans=3, embed_dims=[ 128, 160, 256], |
|
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=[1, 1, 1], sr_ratios=[8, 4, 2, 1], **kwargs): |
|
super().__init__() |
|
|
|
self.encoder1 = nn.Conv2d(3, 16, 3, stride=1, padding=1) |
|
self.encoder2 = nn.Conv2d(16, 32, 3, stride=1, padding=1) |
|
self.encoder3 = nn.Conv2d(32, 128, 3, stride=1, padding=1) |
|
|
|
self.ebn1 = nn.BatchNorm2d(16) |
|
self.ebn2 = nn.BatchNorm2d(32) |
|
self.ebn3 = nn.BatchNorm2d(128) |
|
|
|
self.norm3 = norm_layer(embed_dims[1]) |
|
self.norm4 = norm_layer(embed_dims[2]) |
|
|
|
self.dnorm3 = norm_layer(160) |
|
self.dnorm4 = norm_layer(128) |
|
|
|
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] |
|
|
|
self.block1 = nn.ModuleList([shiftedBlock( |
|
dim=embed_dims[1], num_heads=num_heads[0], mlp_ratio=1, qkv_bias=qkv_bias, qk_scale=qk_scale, |
|
drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[0], norm_layer=norm_layer, |
|
sr_ratio=sr_ratios[0])]) |
|
|
|
self.block2 = nn.ModuleList([shiftedBlock( |
|
dim=embed_dims[2], num_heads=num_heads[0], mlp_ratio=1, qkv_bias=qkv_bias, qk_scale=qk_scale, |
|
drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[1], norm_layer=norm_layer, |
|
sr_ratio=sr_ratios[0])]) |
|
|
|
self.dblock1 = nn.ModuleList([shiftedBlock( |
|
dim=embed_dims[1], num_heads=num_heads[0], mlp_ratio=1, qkv_bias=qkv_bias, qk_scale=qk_scale, |
|
drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[0], norm_layer=norm_layer, |
|
sr_ratio=sr_ratios[0])]) |
|
|
|
self.dblock2 = nn.ModuleList([shiftedBlock( |
|
dim=embed_dims[0], num_heads=num_heads[0], mlp_ratio=1, qkv_bias=qkv_bias, qk_scale=qk_scale, |
|
drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[1], norm_layer=norm_layer, |
|
sr_ratio=sr_ratios[0])]) |
|
|
|
self.patch_embed3 = OverlapPatchEmbed(img_size=img_size // 4, patch_size=3, stride=2, in_chans=embed_dims[0], |
|
embed_dim=embed_dims[1]) |
|
self.patch_embed4 = OverlapPatchEmbed(img_size=img_size // 8, patch_size=3, stride=2, in_chans=embed_dims[1], |
|
embed_dim=embed_dims[2]) |
|
|
|
self.decoder1 = nn.Conv2d(256, 160, 3, stride=1,padding=1) |
|
self.decoder2 = nn.Conv2d(160, 128, 3, stride=1, padding=1) |
|
self.decoder3 = nn.Conv2d(128, 32, 3, stride=1, padding=1) |
|
self.decoder4 = nn.Conv2d(32, 16, 3, stride=1, padding=1) |
|
self.decoder5 = nn.Conv2d(16, 16, 3, stride=1, padding=1) |
|
|
|
self.dbn1 = nn.BatchNorm2d(160) |
|
self.dbn2 = nn.BatchNorm2d(128) |
|
self.dbn3 = nn.BatchNorm2d(32) |
|
self.dbn4 = nn.BatchNorm2d(16) |
|
|
|
self.final = nn.Conv2d(16, num_classes, kernel_size=1) |
|
|
|
self.soft = nn.Softmax(dim =1) |
|
|
|
def forward(self, x, text_dummy): |
|
|
|
B = x.shape[0] |
|
|
|
|
|
|
|
|
|
out = F.relu(F.max_pool2d(self.ebn1(self.encoder1(x)),2,2)) |
|
t1 = out |
|
|
|
out = F.relu(F.max_pool2d(self.ebn2(self.encoder2(out)),2,2)) |
|
t2 = out |
|
|
|
out = F.relu(F.max_pool2d(self.ebn3(self.encoder3(out)),2,2)) |
|
t3 = out |
|
|
|
|
|
|
|
|
|
out,H,W = self.patch_embed3(out) |
|
for i, blk in enumerate(self.block1): |
|
out = blk(out, H, W) |
|
out = self.norm3(out) |
|
out = out.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() |
|
t4 = out |
|
|
|
|
|
|
|
out ,H,W= self.patch_embed4(out) |
|
for i, blk in enumerate(self.block2): |
|
out = blk(out, H, W) |
|
out = self.norm4(out) |
|
out = out.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() |
|
|
|
|
|
|
|
out = F.relu(F.interpolate(self.dbn1(self.decoder1(out)),scale_factor=(2,2),mode ='bilinear')) |
|
|
|
out = torch.add(out,t4) |
|
_,_,H,W = out.shape |
|
out = out.flatten(2).transpose(1,2) |
|
for i, blk in enumerate(self.dblock1): |
|
out = blk(out, H, W) |
|
|
|
|
|
|
|
out = self.dnorm3(out) |
|
out = out.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() |
|
out = F.relu(F.interpolate(self.dbn2(self.decoder2(out)),scale_factor=(2,2),mode ='bilinear')) |
|
out = torch.add(out,t3) |
|
_,_,H,W = out.shape |
|
out = out.flatten(2).transpose(1,2) |
|
|
|
for i, blk in enumerate(self.dblock2): |
|
out = blk(out, H, W) |
|
|
|
out = self.dnorm4(out) |
|
out = out.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() |
|
|
|
out = F.relu(F.interpolate(self.dbn3(self.decoder3(out)),scale_factor=(2,2),mode ='bilinear')) |
|
out = torch.add(out,t2) |
|
out = F.relu(F.interpolate(self.dbn4(self.decoder4(out)),scale_factor=(2,2),mode ='bilinear')) |
|
out = torch.add(out,t1) |
|
out = F.relu(F.interpolate(self.decoder5(out),scale_factor=(2,2),mode ='bilinear')) |
|
|
|
return self.soft(self.final(out)),0 |
|
|
|
|
|
class UNext_S(nn.Module): |
|
|
|
|
|
|
|
def __init__(self, num_classes, input_channels=3, deep_supervision=False,img_size=256, patch_size=16, in_chans=3, embed_dims=[32, 64, 128, 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=[1, 1, 1], sr_ratios=[8, 4, 2, 1], **kwargs): |
|
super().__init__() |
|
|
|
self.encoder1 = nn.Conv2d(3, 8, 3, stride=1, padding=1) |
|
self.encoder2 = nn.Conv2d(8, 16, 3, stride=1, padding=1) |
|
self.encoder3 = nn.Conv2d(16, 32, 3, stride=1, padding=1) |
|
|
|
self.ebn1 = nn.BatchNorm2d(8) |
|
self.ebn2 = nn.BatchNorm2d(16) |
|
self.ebn3 = nn.BatchNorm2d(32) |
|
|
|
self.norm3 = norm_layer(embed_dims[1]) |
|
self.norm4 = norm_layer(embed_dims[2]) |
|
|
|
self.dnorm3 = norm_layer(64) |
|
self.dnorm4 = norm_layer(32) |
|
|
|
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] |
|
|
|
self.block1 = nn.ModuleList([shiftedBlock( |
|
dim=embed_dims[1], num_heads=num_heads[0], mlp_ratio=1, qkv_bias=qkv_bias, qk_scale=qk_scale, |
|
drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[0], norm_layer=norm_layer, |
|
sr_ratio=sr_ratios[0])]) |
|
|
|
self.block2 = nn.ModuleList([shiftedBlock( |
|
dim=embed_dims[2], num_heads=num_heads[0], mlp_ratio=1, qkv_bias=qkv_bias, qk_scale=qk_scale, |
|
drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[1], norm_layer=norm_layer, |
|
sr_ratio=sr_ratios[0])]) |
|
|
|
self.dblock1 = nn.ModuleList([shiftedBlock( |
|
dim=embed_dims[1], num_heads=num_heads[0], mlp_ratio=1, qkv_bias=qkv_bias, qk_scale=qk_scale, |
|
drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[0], norm_layer=norm_layer, |
|
sr_ratio=sr_ratios[0])]) |
|
|
|
self.dblock2 = nn.ModuleList([shiftedBlock( |
|
dim=embed_dims[0], num_heads=num_heads[0], mlp_ratio=1, qkv_bias=qkv_bias, qk_scale=qk_scale, |
|
drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[1], norm_layer=norm_layer, |
|
sr_ratio=sr_ratios[0])]) |
|
|
|
self.patch_embed3 = OverlapPatchEmbed(img_size=img_size // 4, patch_size=3, stride=2, in_chans=embed_dims[0], |
|
embed_dim=embed_dims[1]) |
|
self.patch_embed4 = OverlapPatchEmbed(img_size=img_size // 8, patch_size=3, stride=2, in_chans=embed_dims[1], |
|
embed_dim=embed_dims[2]) |
|
|
|
self.decoder1 = nn.Conv2d(128, 64, 3, stride=1,padding=1) |
|
self.decoder2 = nn.Conv2d(64, 32, 3, stride=1, padding=1) |
|
self.decoder3 = nn.Conv2d(32, 16, 3, stride=1, padding=1) |
|
self.decoder4 = nn.Conv2d(16, 8, 3, stride=1, padding=1) |
|
self.decoder5 = nn.Conv2d(8, 8, 3, stride=1, padding=1) |
|
|
|
self.dbn1 = nn.BatchNorm2d(64) |
|
self.dbn2 = nn.BatchNorm2d(32) |
|
self.dbn3 = nn.BatchNorm2d(16) |
|
self.dbn4 = nn.BatchNorm2d(8) |
|
|
|
self.final = nn.Conv2d(8, num_classes, kernel_size=1) |
|
|
|
self.soft = nn.Softmax(dim =1) |
|
|
|
def forward(self, x, text_dummy): |
|
|
|
B = x.shape[0] |
|
|
|
|
|
|
|
|
|
out = F.relu(F.max_pool2d(self.ebn1(self.encoder1(x)),2,2)) |
|
t1 = out |
|
|
|
out = F.relu(F.max_pool2d(self.ebn2(self.encoder2(out)),2,2)) |
|
t2 = out |
|
|
|
out = F.relu(F.max_pool2d(self.ebn3(self.encoder3(out)),2,2)) |
|
t3 = out |
|
|
|
|
|
|
|
|
|
out,H,W = self.patch_embed3(out) |
|
for i, blk in enumerate(self.block1): |
|
out = blk(out, H, W) |
|
out = self.norm3(out) |
|
out = out.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() |
|
t4 = out |
|
|
|
|
|
|
|
out ,H,W= self.patch_embed4(out) |
|
for i, blk in enumerate(self.block2): |
|
out = blk(out, H, W) |
|
out = self.norm4(out) |
|
out = out.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() |
|
|
|
|
|
|
|
out = F.relu(F.interpolate(self.dbn1(self.decoder1(out)),scale_factor=(2,2),mode ='bilinear')) |
|
|
|
out = torch.add(out,t4) |
|
_,_,H,W = out.shape |
|
out = out.flatten(2).transpose(1,2) |
|
for i, blk in enumerate(self.dblock1): |
|
out = blk(out, H, W) |
|
|
|
|
|
|
|
out = self.dnorm3(out) |
|
out = out.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() |
|
out = F.relu(F.interpolate(self.dbn2(self.decoder2(out)),scale_factor=(2,2),mode ='bilinear')) |
|
out = torch.add(out,t3) |
|
_,_,H,W = out.shape |
|
out = out.flatten(2).transpose(1,2) |
|
|
|
for i, blk in enumerate(self.dblock2): |
|
out = blk(out, H, W) |
|
|
|
out = self.dnorm4(out) |
|
out = out.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() |
|
|
|
out = F.relu(F.interpolate(self.dbn3(self.decoder3(out)),scale_factor=(2,2),mode ='bilinear')) |
|
out = torch.add(out,t2) |
|
out = F.relu(F.interpolate(self.dbn4(self.decoder4(out)),scale_factor=(2,2),mode ='bilinear')) |
|
out = torch.add(out,t1) |
|
out = F.relu(F.interpolate(self.decoder5(out),scale_factor=(2,2),mode ='bilinear')) |
|
|
|
return self.final(out) |
|
|
|
|
|
class medt_net(nn.Module): |
|
|
|
def __init__(self, block, block_2, layers, num_classes=2, zero_init_residual=True, |
|
groups=8, width_per_group=64, replace_stride_with_dilation=None, |
|
norm_layer=None, s=0.125, img_size = 128,imgchan = 3): |
|
super(medt_net, self).__init__() |
|
if norm_layer is None: |
|
norm_layer = nn.BatchNorm2d |
|
self._norm_layer = norm_layer |
|
|
|
self.inplanes = int(64 * s) |
|
self.dilation = 1 |
|
if replace_stride_with_dilation is None: |
|
replace_stride_with_dilation = [False, False, False] |
|
if len(replace_stride_with_dilation) != 3: |
|
raise ValueError("replace_stride_with_dilation should be None " |
|
"or a 3-element tuple, got {}".format(replace_stride_with_dilation)) |
|
self.groups = groups |
|
self.base_width = width_per_group |
|
self.conv1 = nn.Conv2d(imgchan, self.inplanes, kernel_size=7, stride=2, padding=3, |
|
bias=False) |
|
self.conv2 = nn.Conv2d(self.inplanes, 128, kernel_size=3, stride=1, padding=1, bias=False) |
|
self.conv3 = nn.Conv2d(128, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False) |
|
self.bn1 = norm_layer(self.inplanes) |
|
self.bn2 = norm_layer(128) |
|
self.bn3 = norm_layer(self.inplanes) |
|
self.bn1 = norm_layer(self.inplanes) |
|
self.relu = nn.ReLU(inplace=True) |
|
self.layer1 = self._make_layer(block, int(128 * s), layers[0], kernel_size= (img_size//2)) |
|
self.layer2 = self._make_layer(block, int(256 * s), layers[1], stride=2, kernel_size=(img_size//2), |
|
dilate=replace_stride_with_dilation[0]) |
|
|
|
self.decoder4 = nn.Conv2d(int(512*s) , int(256*s), kernel_size=3, stride=1, padding=1) |
|
self.decoder5 = nn.Conv2d(int(256*s) , int(128*s) , kernel_size=3, stride=1, padding=1) |
|
self.adjust = nn.Conv2d(int(128*s) , num_classes, kernel_size=1, stride=1, padding=0) |
|
self.soft = nn.Softmax(dim=1) |
|
|
|
|
|
self.conv1_p = nn.Conv2d(imgchan, self.inplanes, kernel_size=7, stride=2, padding=3, |
|
bias=False) |
|
self.conv2_p = nn.Conv2d(self.inplanes,128, kernel_size=3, stride=1, padding=1, |
|
bias=False) |
|
self.conv3_p = nn.Conv2d(128, self.inplanes, kernel_size=3, stride=1, padding=1, |
|
bias=False) |
|
self.bn1_p = norm_layer(self.inplanes) |
|
self.bn2_p = norm_layer(128) |
|
self.bn3_p = norm_layer(self.inplanes) |
|
|
|
self.relu_p = nn.ReLU(inplace=True) |
|
|
|
img_size_p = img_size // 4 |
|
|
|
self.layer1_p = self._make_layer(block_2, int(128 * s), layers[0], kernel_size= (img_size_p//2)) |
|
self.layer2_p = self._make_layer(block_2, int(256 * s), layers[1], stride=2, kernel_size=(img_size_p//2), |
|
dilate=replace_stride_with_dilation[0]) |
|
self.layer3_p = self._make_layer(block_2, int(512 * s), layers[2], stride=2, kernel_size=(img_size_p//4), |
|
dilate=replace_stride_with_dilation[1]) |
|
self.layer4_p = self._make_layer(block_2, int(1024 * s), layers[3], stride=2, kernel_size=(img_size_p//8), |
|
dilate=replace_stride_with_dilation[2]) |
|
|
|
|
|
self.decoder1_p = nn.Conv2d(int(1024 *2*s) , int(1024*2*s), kernel_size=3, stride=2, padding=1) |
|
self.decoder2_p = nn.Conv2d(int(1024 *2*s) , int(1024*s), kernel_size=3, stride=1, padding=1) |
|
self.decoder3_p = nn.Conv2d(int(1024*s), int(512*s), kernel_size=3, stride=1, padding=1) |
|
self.decoder4_p = nn.Conv2d(int(512*s) , int(256*s), kernel_size=3, stride=1, padding=1) |
|
self.decoder5_p = nn.Conv2d(int(256*s) , int(128*s) , kernel_size=3, stride=1, padding=1) |
|
|
|
self.decoderf = nn.Conv2d(int(128*s) , int(128*s) , kernel_size=3, stride=1, padding=1) |
|
self.adjust_p = nn.Conv2d(int(128*s) , num_classes, kernel_size=1, stride=1, padding=0) |
|
self.soft_p = nn.Softmax(dim=1) |
|
|
|
|
|
def _make_layer(self, block, planes, blocks, kernel_size=56, stride=1, dilate=False): |
|
norm_layer = self._norm_layer |
|
downsample = None |
|
previous_dilation = self.dilation |
|
if dilate: |
|
self.dilation *= stride |
|
stride = 1 |
|
if stride != 1 or self.inplanes != planes * block.expansion: |
|
downsample = nn.Sequential( |
|
conv1x1(self.inplanes, planes * block.expansion, stride), |
|
norm_layer(planes * block.expansion), |
|
) |
|
|
|
layers = [] |
|
layers.append(block(self.inplanes, planes, stride, downsample, groups=self.groups, |
|
base_width=self.base_width, dilation=previous_dilation, |
|
norm_layer=norm_layer, kernel_size=kernel_size)) |
|
self.inplanes = planes * block.expansion |
|
if stride != 1: |
|
kernel_size = kernel_size // 2 |
|
|
|
for _ in range(1, blocks): |
|
layers.append(block(self.inplanes, planes, groups=self.groups, |
|
base_width=self.base_width, dilation=self.dilation, |
|
norm_layer=norm_layer, kernel_size=kernel_size)) |
|
|
|
return nn.Sequential(*layers) |
|
|
|
def _forward_impl(self, x): |
|
|
|
xin = x.clone() |
|
x = self.conv1(x) |
|
x = self.bn1(x) |
|
x = self.relu(x) |
|
x = self.conv2(x) |
|
x = self.bn2(x) |
|
x = self.relu(x) |
|
x = self.conv3(x) |
|
x = self.bn3(x) |
|
x = self.relu(x) |
|
|
|
x1 = self.layer1(x) |
|
x2 = self.layer2(x1) |
|
|
|
x = F.relu(F.interpolate(self.decoder4(x2) , scale_factor=(2,2), mode ='bilinear')) |
|
x = torch.add(x, x1) |
|
x = F.relu(F.interpolate(self.decoder5(x) , scale_factor=(2,2), mode ='bilinear')) |
|
|
|
|
|
|
|
x_loc = x.clone() |
|
|
|
for i in range(0,4): |
|
for j in range(0,4): |
|
|
|
x_p = xin[:,:,32*i:32*(i+1),32*j:32*(j+1)] |
|
|
|
x_p = self.conv1_p(x_p) |
|
x_p = self.bn1_p(x_p) |
|
x_p = self.relu(x_p) |
|
|
|
x_p = self.conv2_p(x_p) |
|
x_p = self.bn2_p(x_p) |
|
x_p = self.relu(x_p) |
|
x_p = self.conv3_p(x_p) |
|
x_p = self.bn3_p(x_p) |
|
x_p = self.relu(x_p) |
|
|
|
x1_p = self.layer1_p(x_p) |
|
x2_p = self.layer2_p(x1_p) |
|
x3_p = self.layer3_p(x2_p) |
|
x4_p = self.layer4_p(x3_p) |
|
|
|
x_p = F.relu(F.interpolate(self.decoder1_p(x4_p), scale_factor=(2,2), mode ='bilinear')) |
|
x_p = torch.add(x_p, x4_p) |
|
x_p = F.relu(F.interpolate(self.decoder2_p(x_p) , scale_factor=(2,2), mode ='bilinear')) |
|
x_p = torch.add(x_p, x3_p) |
|
x_p = F.relu(F.interpolate(self.decoder3_p(x_p) , scale_factor=(2,2), mode ='bilinear')) |
|
x_p = torch.add(x_p, x2_p) |
|
x_p = F.relu(F.interpolate(self.decoder4_p(x_p) , scale_factor=(2,2), mode ='bilinear')) |
|
x_p = torch.add(x_p, x1_p) |
|
x_p = F.relu(F.interpolate(self.decoder5_p(x_p) , scale_factor=(2,2), mode ='bilinear')) |
|
|
|
x_loc[:,:,32*i:32*(i+1),32*j:32*(j+1)] = x_p |
|
|
|
x = torch.add(x,x_loc) |
|
x = F.relu(self.decoderf(x)) |
|
|
|
x = self.adjust(F.relu(x)) |
|
|
|
return x |
|
|
|
def forward(self, x, text_dummy): |
|
return self._forward_impl(x) |
|
|
