|
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) |
|
|
|
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) |
|
|
|
|
|
|
|
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) |
|
|
|
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) |
|
|
|
|
|
|
|
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): |
|
|
|
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 |
|
|
