def __init__(self,
num_classes,
trunk='wrn38',
criterion=None,
fuse_aspp=False,
attn_2b=False):
super(MscaleDeeper, self).__init__()
self.criterion = criterion
self.fuse_aspp = fuse_aspp
self.attn_2b = attn_2b
self.backbone, s2_ch, s4_ch, high_level_ch = get_trunk(
trunk_name=trunk, output_stride=8)
self.aspp, aspp_out_ch = get_aspp(high_level_ch,
bottleneck_ch=256,
output_stride=8)
self.convs2 = nn.Conv2d(s2_ch, 32, kernel_size=1, bias=False)
self.convs4 = nn.Conv2d(s4_ch, 64, kernel_size=1, bias=False)
self.conv_up1 = nn.Conv2d(aspp_out_ch, 256, kernel_size=1, bias=False)
self.conv_up2 = ConvBnRelu(256 + 64, 256, kernel_size=5, padding=2)
self.conv_up3 = ConvBnRelu(256 + 32, 256, kernel_size=5, padding=2)
self.conv_up5 = nn.Conv2d(256, num_classes, kernel_size=1, bias=False)
# Scale-attention prediction head
if self.attn_2b:
attn_ch = 2
else:
attn_ch = 1
self.scale_attn = make_attn_head(in_ch=256, out_ch=attn_ch)
if cfg.OPTIONS.INIT_DECODER:
initialize_weights(self.convs2, self.convs4, self.conv_up1,
self.conv_up2, self.conv_up3, self.conv_up5,
self.scale_attn)
def __init__(self, num_classes, trunk='hrnetv2', criterion=None):
super(MscaleOCR, self).__init__()
self.criterion = criterion
self.backbone, _, _, high_level_ch = get_trunk(trunk)
self.ocr = OCR_block(high_level_ch)
self.scale_attn = make_attn_head(
in_ch=cfg.MODEL.OCR.MID_CHANNELS, out_ch=1)
def __init__(self, num_classes, trunk='hrnetv2', criterion=None):
super(MscaleBasic, self).__init__()
self.criterion = criterion
self.backbone, _, _, high_level_ch = get_trunk(trunk_name=trunk,
output_stride=8)
self.cls_head = make_seg_head(in_ch=high_level_ch, out_ch=num_classes)
self.scale_attn = make_attn_head(in_ch=high_level_ch, out_ch=1)
def __init__(self, num_classes, criterion=None):
super(MscaleOCR, self).__init__()
self.criterion = criterion
self.backbone = get_seg_model()
high_level_ch = self.backbone.high_level_ch
self.ocr = OCR_block(high_level_ch)
self.scale_attn = make_attn_head(in_ch=cfg.MODEL.OCR.MID_CHANNELS,
out_ch=1)
def __init__(self, num_classes, trunk='hrnetv2', criterion=None):
super(ASPP, self).__init__()
self.criterion = criterion
self.backbone, s2_ch, _s4_ch, high_level_ch = get_trunk(trunk)
self.aspp, aspp_out_ch = get_aspp(high_level_ch,
bottleneck_ch=cfg.MODEL.ASPP_BOT_CH,
output_stride=8)
self.bot_aspp = nn.Conv2d(aspp_out_ch, 256, kernel_size=1, bias=False)
self.final = make_seg_head(in_ch=256, out_ch=num_classes)
self.scale_attn = make_attn_head(in_ch=256, out_ch=1)
initialize_weights(self.final)
def __init__(self, num_classes, trunk='wrn38', criterion=None,
use_dpc=False, fuse_aspp=False, attn_2b=False, attn_cls=False):
super(MscaleV3Plus, self).__init__()
self.criterion = criterion
self.fuse_aspp = fuse_aspp
self.attn_2b = attn_2b
self.attn_cls = attn_cls
self.backbone, s2_ch, _s4_ch, high_level_ch = get_trunk(trunk)
self.aspp, aspp_out_ch = get_aspp(high_level_ch,
bottleneck_ch=256,
output_stride=8,
dpc=use_dpc)
self.bot_fine = nn.Conv2d(s2_ch, 48, kernel_size=1, bias=False)
self.bot_aspp = nn.Conv2d(aspp_out_ch, 256, kernel_size=1, bias=False)
# Semantic segmentation prediction head
bot_ch = cfg.MODEL.SEGATTN_BOT_CH
self.final = nn.Sequential(
nn.Conv2d(256 + 48, bot_ch, kernel_size=3, padding=1, bias=False),
Norm2d(bot_ch),
nn.ReLU(inplace=True),
nn.Conv2d(bot_ch, bot_ch, kernel_size=3, padding=1, bias=False),
Norm2d(bot_ch),
nn.ReLU(inplace=True),
nn.Conv2d(bot_ch, num_classes, kernel_size=1, bias=False))
# Scale-attention prediction head
if self.attn_2b:
attn_ch = 2
else:
if self.attn_cls:
attn_ch = num_classes
else:
attn_ch = 1
scale_in_ch = 256 + 48
self.scale_attn = make_attn_head(in_ch=scale_in_ch,
out_ch=attn_ch)
if cfg.OPTIONS.INIT_DECODER:
initialize_weights(self.bot_fine)
initialize_weights(self.bot_aspp)
initialize_weights(self.scale_attn)
initialize_weights(self.final)
else:
initialize_weights(self.final)
def __init__(self, num_classes, trunk='wrn38', criterion=None,
use_dpc=False, fuse_aspp=False, attn_2b=False):
super(MscaleV3Plus, self).__init__()
self.criterion = criterion
self.fuse_aspp = fuse_aspp
self.attn_2b = attn_2b
self.backbone, s2_ch, _s4_ch, high_level_ch = get_trunk(trunk)
self.aspp, aspp_out_ch = get_aspp(high_level_ch,
bottleneck_ch=256,
output_stride=8,
dpc=use_dpc,
img_norm = False)
self.bot_fine = nn.Conv2d(s2_ch, 48, kernel_size=1, bias=False)
self.bot_aspp = nn.Conv2d(aspp_out_ch, 256, kernel_size=1, bias=False)
#self.asnb = ASNB(low_in_channels = 48, high_in_channels=256, out_channels=256, key_channels=64, value_channels=256, dropout=0., sizes=([1]), norm_type='batchnorm',attn_scale=0.25)
self.adnb = ADNB(d_model=256, nhead=8, num_encoder_layers=2, dim_feedforward=256, dropout=0.5, activation="relu", num_feature_levels=1, enc_n_points=4)
# Semantic segmentation prediction head
bot_ch = cfg.MODEL.SEGATTN_BOT_CH
self.final = nn.Sequential(
nn.Conv2d(256 + 48, bot_ch, kernel_size=3, padding=1, bias=False),
Norm2d(bot_ch),
nn.ReLU(inplace=True),
nn.Conv2d(bot_ch, bot_ch, kernel_size=3, padding=1, bias=False),
Norm2d(bot_ch),
nn.ReLU(inplace=True),
nn.Conv2d(bot_ch, num_classes, kernel_size=1, bias=False))
# Scale-attention prediction head
if self.attn_2b:
attn_ch = 2
else:
attn_ch = 1
scale_in_ch = 256 + 48
self.scale_attn = make_attn_head(in_ch=scale_in_ch,
out_ch=attn_ch)
if cfg.OPTIONS.INIT_DECODER:
initialize_weights(self.bot_fine)
initialize_weights(self.bot_aspp)
initialize_weights(self.scale_attn)
initialize_weights(self.final)
else:
initialize_weights(self.final)
