| from model import common |
| |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| def make_model(args, parent=False): |
| return SMGARN(args) |
|
|
|
|
| class SnowMaskBlock(nn.Module): |
| def __init__(self, embed_dim): |
| super(SnowMaskBlock, self).__init__() |
| self.smblock = common.MaskBlock(embed_dim) |
| self.conv3 = common.default_conv(embed_dim, embed_dim, 3) |
| self.norm = nn.LayerNorm(embed_dim) |
|
|
| def forward(self, x): |
| B, H, W = x.shape[0], x.shape[2], x.shape[3] |
| shortcut = x |
| x = self.smblock(x) |
| x = self.norm(x.flatten(2).transpose(-1, -2)) |
| x = self.conv3(x.transpose(-1, -2).view(B, -1, H, W)) |
| return x + shortcut |
|
|
|
|
| class Mask_Net(nn.Module): |
| def __init__(self, n_colors, embed_dim, conv): |
| super(Mask_Net, self).__init__() |
| h = [] |
| h.append(conv(n_colors, embed_dim, 3)) |
| h.append(conv(embed_dim, embed_dim, 3)) |
| self.head = nn.Sequential(*h) |
| self.g_mp1 = SnowMaskBlock(embed_dim) |
|
|
| self.conv_out1 = common.default_conv(embed_dim, embed_dim, 3) |
| self.conv_out2 = common.default_conv(embed_dim, 3, 3) |
|
|
| def forward(self, x): |
| x = self.head(x) |
| out_1 = self.g_mp1(x) |
| out_1 = self.conv_out1(out_1) |
|
|
| out = self.conv_out2(out_1) |
| return out, out_1 |
|
|
| class ReconstructNet(nn.Module): |
| def __init__(self, n_colors, dim, depth): |
| super(ReconstructNet, self).__init__() |
| self.fusion = common.FusionBlock(n_colors, dim) |
| block = [] |
| for i in range(depth): |
| block.append(common.MARB(dim)) |
| self.recon = nn.Sequential(*block) |
| t = [] |
| t.append(common.default_conv(dim, dim, 3)) |
| t.append(nn.ReLU(True)) |
| t.append(common.default_conv(dim, n_colors, 3)) |
| self.tail = nn.Sequential(*t) |
|
|
| def forward(self, x, mask): |
| x = self.fusion(x, mask) |
| out = self.recon(x) + x |
| out = self.tail(out) |
| return out |
|
|
|
|
|
|
| class SMGARN(nn.Module): |
| def __init__(self, args, conv=common.default_conv): |
| super(SMGARN, self).__init__() |
| print("SMGARN") |
| n_colors = 3 |
| dim = 112 |
| ReconNet_num = 3 |
|
|
| self.Stage_I = Mask_Net(n_colors=n_colors, embed_dim=dim, conv=conv) |
|
|
| self.Stage_II = ReconstructNet(n_colors, dim, ReconNet_num) |
|
|
| def forward(self, x): |
| mask, mask_feature = self.Stage_I(x) |
| x = self.Stage_II(x, mask_feature) |
| return x, mask |
