def __init__(self, in_channels, key_channels, value_channels, out_channels=None, scale=1, bn_type=None):
super(_SelfAttentionBlock, self).__init__()
self.scale = scale
self.in_channels = in_channels
self.out_channels = out_channels
self.key_channels = key_channels
self.value_channels = value_channels
if out_channels == None:
self.out_channels = in_channels
self.pool = nn.MaxPool2d(kernel_size=(scale, scale))
self.f_key = nn.Sequential(
nn.Conv2d(in_channels=self.in_channels, out_channels=self.key_channels,
kernel_size=1, stride=1, padding=0),
ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type),
nn.Conv2d(in_channels=self.key_channels, out_channels=self.key_channels,
kernel_size=1, stride=1, padding=0),
ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type),
)
self.f_query = nn.Sequential(
nn.Conv2d(in_channels=self.in_channels, out_channels=self.key_channels,
kernel_size=1, stride=1, padding=0),
ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type),
nn.Conv2d(in_channels=self.key_channels, out_channels=self.key_channels,
kernel_size=1, stride=1, padding=0),
ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type),
)
self.f_value = nn.Conv2d(in_channels=self.in_channels, out_channels=self.value_channels,
kernel_size=1, stride=1, padding=0)
self.W = nn.Conv2d(in_channels=self.value_channels, out_channels=self.out_channels,
kernel_size=1, stride=1, padding=0)
nn.init.constant_(self.W.weight, 0)
nn.init.constant_(self.W.bias, 0)
def __init__(self, configer):
self.inplanes = 128
super(BaseOCNet, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
self.backbone = BackboneSelector(configer).get_backbone()
# extra added layers
if "wide_resnet38" in self.configer.get('network', 'backbone'):
in_channels = [2048, 4096]
else:
in_channels = [1024, 2048]
self.oc_module_pre = nn.Sequential(
nn.Conv2d(in_channels[1], 512, kernel_size=3, stride=1, padding=1),
ModuleHelper.BNReLU(512, bn_type=self.configer.get('network', 'bn_type')),
)
from lib.models.modules.base_oc_block import BaseOC_Module
self.oc_module = BaseOC_Module(in_channels=512,
out_channels=512,
key_channels=256,
value_channels=256,
dropout=0.05,
sizes=([1]),
bn_type=self.configer.get('network', 'bn_type'))
self.cls = nn.Conv2d(512, self.num_classes, kernel_size=1, stride=1, padding=0, bias=True)
self.dsn = nn.Sequential(
nn.Conv2d(in_channels[0], 512, kernel_size=3, stride=1, padding=1),
ModuleHelper.BNReLU(512, bn_type=self.configer.get('network', 'bn_type')),
nn.Conv2d(512, self.num_classes, kernel_size=1, stride=1, padding=0, bias=True)
)
def __init__(self, configer):
self.inplanes = 128
super(SpatialOCRNet, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
self.backbone = BackboneSelector(configer).get_backbone()
# extra added layers
if "wide_resnet38" in self.configer.get('network', 'backbone'):
in_channels = [2048, 4096]
else:
in_channels = [1024, 2048]
self.conv_3x3 = nn.Sequential(
nn.Conv2d(in_channels[1], 512, kernel_size=3, stride=1, padding=1),
ModuleHelper.BNReLU(512, bn_type=self.configer.get('network', 'bn_type')),
)
from lib.models.modules.spatial_ocr_block import SpatialGather_Module, SpatialOCR_Module
self.spatial_context_head = SpatialGather_Module(self.num_classes)
self.spatial_ocr_head = SpatialOCR_Module(in_channels=512,
key_channels=256,
out_channels=512,
scale=1,
dropout=0.05,
bn_type=self.configer.get('network', 'bn_type'))
self.head = nn.Conv2d(512, self.num_classes, kernel_size=1, stride=1, padding=0, bias=True)
self.dsn_head = nn.Sequential(
nn.Conv2d(in_channels[0], 512, kernel_size=3, stride=1, padding=1),
ModuleHelper.BNReLU(512, bn_type=self.configer.get('network', 'bn_type')),
nn.Dropout2d(0.05),
nn.Conv2d(512, self.num_classes, kernel_size=1, stride=1, padding=0, bias=True)
)
def __init__(self,
in_channels,
out_channels,
dropout,
sizes=([1]),
bn_type=None):
super(Pyramid_OC_Module, self).__init__()
self.group = len(sizes)
self.stages = []
self.stages = nn.ModuleList([
self._make_stage(in_channels, out_channels, in_channels // 2,
in_channels, size, bn_type) for size in sizes
])
self.conv_bn_dropout = nn.Sequential(
nn.Conv2d(2 * in_channels * self.group,
out_channels,
kernel_size=1,
padding=0),
ModuleHelper.BNReLU(out_channels, bn_type=bn_type),
nn.Dropout2d(dropout))
self.up_dr = nn.Sequential(
nn.Conv2d(in_channels,
in_channels * self.group,
kernel_size=1,
padding=0),
ModuleHelper.BNReLU(in_channels * self.group, bn_type=bn_type))
def __init__(self,
num_classes,
dropout=0,
bn_type=None,
inplane1=512,
inplane2=256):
super(CE2P_Decoder_Module, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(inplane1,
256,
kernel_size=1,
padding=0,
dilation=1,
bias=False),
ModuleHelper.BNReLU(256, bn_type=bn_type),
)
self.conv2 = nn.Sequential(
nn.Conv2d(inplane2,
48,
kernel_size=1,
stride=1,
padding=0,
dilation=1,
bias=False),
ModuleHelper.BNReLU(48, bn_type=bn_type),
)
self.conv3 = nn.Sequential(
nn.Conv2d(304,
256,
kernel_size=1,
padding=0,
dilation=1,
bias=False),
ModuleHelper.BNReLU(256, bn_type=bn_type),
nn.Conv2d(256,
256,
kernel_size=1,
padding=0,
dilation=1,
bias=False),
ModuleHelper.BNReLU(256, bn_type=bn_type),
nn.Dropout2d(dropout),
)
self.conv4 = nn.Conv2d(256,
num_classes,
kernel_size=1,
padding=0,
dilation=1,
bias=True)
def __init__(self,
in_channels,
key_channels,
value_channels,
out_channels=None,
bn_type=None):
super(SelfAttentionBlock2D, self).__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.key_channels = key_channels
self.value_channels = value_channels
if out_channels == None:
self.out_channels = in_channels
self.f_key = nn.Sequential(
nn.Conv2d(self.in_channels,
self.key_channels,
kernel_size=1,
bias=False),
ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type),
nn.Conv2d(self.key_channels,
self.key_channels,
kernel_size=1,
bias=False),
ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type),
)
self.f_query = nn.Sequential(
nn.Conv2d(self.in_channels,
self.key_channels,
kernel_size=1,
bias=False),
ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type),
nn.Conv2d(self.key_channels,
self.key_channels,
kernel_size=1,
bias=False),
ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type),
)
self.f_value = nn.Conv2d(self.in_channels,
self.value_channels,
kernel_size=1,
bias=False)
self.W = nn.Sequential(
nn.Conv2d(self.value_channels,
self.out_channels,
kernel_size=1,
bias=False),
ModuleHelper.BNReLU(self.out_channels, bn_type=bn_type))
def __init__(self, configer):
super(SegFix_HRNet, self).__init__()
self.configer = configer
self.backbone = BackboneSelector(configer).get_backbone()
backbone_name = self.configer.get('network', 'backbone')
width = int(backbone_name[-2:])
if 'hrnet' in backbone_name:
in_channels = width * 15
else:
in_channels = width * 31
num_masks = 2
num_directions = DTOffsetConfig.num_classes
mid_channels = 256
self.dir_head = nn.Sequential(
nn.Conv2d(in_channels,
mid_channels,
kernel_size=1,
stride=1,
padding=0,
bias=False),
ModuleHelper.BNReLU(mid_channels,
bn_type=self.configer.get(
'network', 'bn_type')),
nn.Conv2d(mid_channels,
num_directions,
kernel_size=1,
stride=1,
padding=0,
bias=False))
self.mask_head = nn.Sequential(
nn.Conv2d(in_channels,
mid_channels,
kernel_size=1,
stride=1,
padding=0,
bias=False),
ModuleHelper.BNReLU(mid_channels,
bn_type=self.configer.get(
'network', 'bn_type')),
nn.Conv2d(mid_channels,
num_masks,
kernel_size=1,
stride=1,
padding=0,
bias=False))
def __init__(self, configer):
self.inplanes = 128
super(DeepLabV3, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
self.backbone = BackboneSelector(configer).get_backbone()
self.head = nn.Sequential(
ASPPModule(2048, bn_type=self.configer.get('network', 'bn_type')),
nn.Conv2d(512,
self.num_classes,
kernel_size=1,
stride=1,
padding=0,
bias=True))
self.dsn = nn.Sequential(
nn.Conv2d(1024, 512, kernel_size=3, stride=1, padding=1),
ModuleHelper.BNReLU(512,
bn_type=self.configer.get(
'network', 'bn_type')), nn.Dropout2d(0.1),
nn.Conv2d(512,
self.num_classes,
kernel_size=1,
stride=1,
padding=0,
bias=True))
def __init__(self, configer):
self.inplanes = 128
super(ASPOCRNet, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
self.backbone = BackboneSelector(configer).get_backbone()
# extra added layers
if "wide_resnet38" in self.configer.get('network', 'backbone'):
in_channels = [2048, 4096]
else:
in_channels = [1024, 2048]
# we should increase the dilation rates as the output stride is larger
from lib.models.modules.spatial_ocr_block import SpatialOCR_ASP_Module
self.asp_ocr_head = SpatialOCR_ASP_Module(features=2048,
hidden_features=256,
out_features=256,
num_classes=self.num_classes,
bn_type=self.configer.get('network', 'bn_type'))
self.head = nn.Conv2d(256, self.num_classes, kernel_size=1, stride=1, padding=0, bias=True)
self.dsn_head = nn.Sequential(
nn.Conv2d(in_channels[0], 512, kernel_size=3, stride=1, padding=1),
ModuleHelper.BNReLU(512, bn_type=self.configer.get('network', 'bn_type')),
nn.Dropout2d(0.1),
nn.Conv2d(512, self.num_classes, kernel_size=1, stride=1, padding=0, bias=True)
)
def __init__(self, in_channels, key_channels, scale, bn_type=None):
super(FastBaseOC_Context_Module, self).__init__()
self.object_context_block = ObjectAttentionBlock2D(in_channels, key_channels, scale=1, bn_type=bn_type)
self.conv_bn_dropout = nn.Sequential(
nn.Conv2d(in_channels, in_channels, kernel_size=1, padding=0),
ModuleHelper.BNReLU(in_channels, bn_type=bn_type),
)
def __init__(self, bn_type=None, inplane1=512, inplane2=256, outplane=128):
super(Decoder_Module, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(inplane1, 256, kernel_size=1, padding=0, dilation=1, bias=False),
ModuleHelper.BNReLU(256, bn_type=bn_type),
)
self.conv2 = nn.Sequential(
nn.Conv2d(inplane2, 48, kernel_size=1, padding=0, dilation=1, bias=False),
ModuleHelper.BNReLU(48, bn_type=bn_type),
)
self.conv3 = nn.Sequential(
nn.Conv2d(304, outplane, kernel_size=1, padding=0, dilation=1, bias=False),
ModuleHelper.BNReLU(outplane, bn_type=bn_type),
nn.Conv2d(outplane, outplane, kernel_size=1, padding=0, dilation=1, bias=False),
ModuleHelper.BNReLU(outplane, bn_type=bn_type),
)
def __init__(self, configer):
super(HRNet_W48_ASPOCR, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
self.backbone = BackboneSelector(configer).get_backbone()
# extra added layers
in_channels = 720 # 48 + 96 + 192 + 384
from lib.models.modules.spatial_ocr_block import SpatialOCR_ASP_Module
self.asp_ocr_head = SpatialOCR_ASP_Module(features=720,
hidden_features=256,
out_features=256,
dilations=(24, 48, 72),
num_classes=self.num_classes,
bn_type=self.configer.get(
'network', 'bn_type'))
self.cls_head = nn.Conv2d(256,
self.num_classes,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.aux_head = nn.Sequential(
nn.Conv2d(in_channels, 512, kernel_size=3, stride=1, padding=1),
ModuleHelper.BNReLU(512,
bn_type=self.configer.get(
'network', 'bn_type')),
nn.Conv2d(512,
self.num_classes,
kernel_size=1,
stride=1,
padding=0,
bias=False))
def __init__(self,
in_channels,
key_channels,
out_channels,
scale=1,
dropout=0.1,
use_gt=False,
use_bg=False,
use_oc=True,
fetch_attention=False,
bn_type=None):
super(SpatialOCR_Module, self).__init__()
self.use_gt = use_gt
self.use_bg = use_bg
self.use_oc = use_oc
self.fetch_attention = fetch_attention
self.object_context_block = ObjectAttentionBlock2D(
in_channels, key_channels, scale, use_gt, use_bg, fetch_attention,
bn_type)
if self.use_bg:
if self.use_oc:
_in_channels = 3 * in_channels
else:
_in_channels = 2 * in_channels
else:
_in_channels = 2 * in_channels
self.conv_bn_dropout = nn.Sequential(
nn.Conv2d(_in_channels, out_channels, kernel_size=1, padding=0),
ModuleHelper.BNReLU(out_channels, bn_type=bn_type),
nn.Dropout2d(dropout))
def __init__(self, in_channels, out_channels, key_channels, value_channels, dropout=0, sizes=([1]), bn_type=None):
super(BaseOC_Context_Module, self).__init__()
self.stages = []
self.stages = nn.ModuleList([self._make_stage(in_channels, out_channels,
key_channels, value_channels, size, bn_type) for size in sizes])
self.conv_bn_dropout = nn.Sequential(
ModuleHelper.BNReLU(out_channels, bn_type=bn_type),
nn.Dropout2d(dropout),
)
def __init__(self, configer):
super(HRNet_W48_OCR, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
self.backbone = BackboneSelector(configer).get_backbone()
in_channels = 720
self.conv3x3 = nn.Sequential(
nn.Conv2d(in_channels, 512, kernel_size=3, stride=1, padding=1),
ModuleHelper.BNReLU(512,
bn_type=self.configer.get(
'network', 'bn_type')),
)
from lib.models.modules.spatial_ocr_block import SpatialGather_Module
self.ocr_gather_head = SpatialGather_Module(self.num_classes)
from lib.models.modules.spatial_ocr_block import SpatialOCR_Module
self.ocr_distri_head = SpatialOCR_Module(in_channels=512,
key_channels=256,
out_channels=512,
scale=1,
dropout=0.05,
bn_type=self.configer.get(
'network', 'bn_type'))
self.cls_head = nn.Conv2d(512,
self.num_classes,
kernel_size=1,
stride=1,
padding=0,
bias=True)
self.aux_head = nn.Sequential(
nn.Conv2d(in_channels,
in_channels,
kernel_size=3,
stride=1,
padding=1),
ModuleHelper.BNReLU(in_channels,
bn_type=self.configer.get(
'network', 'bn_type')),
nn.Conv2d(in_channels,
self.num_classes,
kernel_size=1,
stride=1,
padding=0,
bias=True))
def __init__(self, configer):
self.inplanes = 128
super(FcnNet_wo_dsn, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
self.backbone = BackboneSelector(configer).get_backbone()
# extra added layers
if "wide_resnet38" in self.configer.get('network', 'backbone'):
in_channels = [2048, 4096]
elif "mobilenetv2" in self.configer.get('network', 'backbone'):
in_channels = [160, 320]
else:
in_channels = [1024, 2048]
self.cls_head = nn.Sequential(
nn.Conv2d(in_channels[1], 512, kernel_size=3, stride=1, padding=1),
ModuleHelper.BNReLU(512,
bn_type=self.configer.get(
'network', 'bn_type')), nn.Dropout2d(0.10),
nn.Conv2d(512,
self.num_classes,
kernel_size=1,
stride=1,
padding=0,
bias=True))
if "mobilenetv2" in self.configer.get('network', 'backbone'):
self.cls_head = nn.Sequential(
nn.Conv2d(in_channels[1],
256,
kernel_size=3,
stride=1,
padding=1),
ModuleHelper.BNReLU(256,
bn_type=self.configer.get(
'network', 'bn_type')),
nn.Dropout2d(0.10),
nn.Conv2d(256,
self.num_classes,
kernel_size=1,
stride=1,
padding=0,
bias=False))
def __init__(self, mid_fea, out_fea, bn_type=None, factor=1):
super(Edge_Module, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(factor * 256,
mid_fea,
kernel_size=1,
padding=0,
dilation=1,
bias=False),
ModuleHelper.BNReLU(mid_fea, bn_type=bn_type),
)
self.conv2 = nn.Sequential(
nn.Conv2d(factor * 512,
mid_fea,
kernel_size=1,
padding=0,
dilation=1,
bias=False),
ModuleHelper.BNReLU(mid_fea, bn_type=bn_type),
)
self.conv3 = nn.Sequential(
nn.Conv2d(factor * 1024,
mid_fea,
kernel_size=1,
padding=0,
dilation=1,
bias=False),
ModuleHelper.BNReLU(mid_fea, bn_type=bn_type),
)
self.conv4 = nn.Conv2d(mid_fea,
out_fea,
kernel_size=3,
padding=1,
dilation=1,
bias=True)
self.conv5 = nn.Conv2d(out_fea * 3,
out_fea,
kernel_size=1,
padding=0,
dilation=1,
bias=True)
def __init__(self, features, inner_features=512, out_features=512, dilations=(12, 24, 36), bn_type=None):
super(ASP_OC_Module_v5, self).__init__()
self.conv1 = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)),
nn.Conv2d(features, inner_features, kernel_size=1, padding=0, dilation=1,
bias=False),
ModuleHelper.BNReLU(inner_features, bn_type=bn_type))
self.conv2 = nn.Sequential(
nn.Conv2d(features, inner_features, kernel_size=1, padding=0, dilation=1, bias=False),
ModuleHelper.BNReLU(inner_features, bn_type=bn_type))
self.conv3 = nn.Sequential(SelfAttentionModuleV2(in_channels=features, key_channels=out_features//2,
value_channels=out_features//2, out_channels=inner_features,
kernel_size=3, padding_list=dilations, dilation_list=dilations,
bn_type=bn_type),
ModuleHelper.BNReLU(inner_features, bn_type=bn_type))
self.conv_bn_dropout = nn.Sequential(
nn.Conv2d(inner_features * 3, out_features, kernel_size=1, padding=0, dilation=1, bias=False),
ModuleHelper.BNReLU(out_features, bn_type=bn_type),
nn.Dropout2d(0.1)
)
def __init__(self, configer):
self.inplanes = 128
super(AspOCNet, self).__init__()
self.configer = configer
self.num_classes = self.configer.get("data", "num_classes")
self.backbone = BackboneSelector(configer).get_backbone()
# extra added layers
if "wide_resnet38" in self.configer.get("network", "backbone"):
in_channels = [2048, 4096]
else:
in_channels = [1024, 2048]
from lib.models.modules.asp_oc_block import ASP_OC_Module
self.context = nn.Sequential(
nn.Conv2d(in_channels[1], 512, kernel_size=3, stride=1, padding=1),
ModuleHelper.BNReLU(512,
bn_type=self.configer.get(
"network", "bn_type")),
ASP_OC_Module(512,
256,
bn_type=self.configer.get("network", "bn_type")),
)
self.cls = nn.Conv2d(512,
self.num_classes,
kernel_size=1,
stride=1,
padding=0,
bias=True)
self.dsn = nn.Sequential(
nn.Conv2d(in_channels[0], 512, kernel_size=3, stride=1, padding=1),
ModuleHelper.BNReLU(512,
bn_type=self.configer.get(
"network", "bn_type")),
nn.Conv2d(512,
self.num_classes,
kernel_size=1,
stride=1,
padding=0,
bias=True),
)
def __init__(self,
in_channels,
key_channels,
scale=1,
use_gt=False,
use_bg=False,
fetch_attention=False,
bn_type=None):
super(_ObjectAttentionBlock, self).__init__()
self.scale = scale
self.in_channels = in_channels
self.key_channels = key_channels
self.use_gt = use_gt
self.use_bg = use_bg
self.fetch_attention = fetch_attention
self.pool = nn.MaxPool2d(kernel_size=(scale, scale))
self.f_pixel = nn.Sequential(
nn.Conv2d(in_channels=self.in_channels,
out_channels=self.key_channels,
kernel_size=1,
stride=1,
padding=0),
ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type),
nn.Conv2d(in_channels=self.key_channels,
out_channels=self.key_channels,
kernel_size=1,
stride=1,
padding=0),
ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type),
)
self.f_object = nn.Sequential(
nn.Conv2d(in_channels=self.in_channels,
out_channels=self.key_channels,
kernel_size=1,
stride=1,
padding=0),
ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type),
nn.Conv2d(in_channels=self.key_channels,
out_channels=self.key_channels,
kernel_size=1,
stride=1,
padding=0),
ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type),
)
self.f_down = nn.Sequential(
nn.Conv2d(in_channels=self.in_channels,
out_channels=self.key_channels,
kernel_size=1,
stride=1,
padding=0),
ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type),
)
self.f_up = nn.Sequential(
nn.Conv2d(in_channels=self.key_channels,
out_channels=self.in_channels,
kernel_size=1,
stride=1,
padding=0),
ModuleHelper.BNReLU(self.in_channels, bn_type=bn_type),
)
def __init__(self,
features,
inner_features=512,
out_features=512,
dilations=(12, 24, 36),
bn_type=None):
super(ASPPModule, self).__init__()
self.conv1 = nn.Sequential(
nn.AdaptiveAvgPool2d((1, 1)),
nn.Conv2d(features,
inner_features,
kernel_size=1,
padding=0,
dilation=1,
bias=False),
ModuleHelper.BNReLU(inner_features, bn_type=bn_type))
self.conv2 = nn.Sequential(
nn.Conv2d(features,
inner_features,
kernel_size=1,
padding=0,
dilation=1,
bias=False),
ModuleHelper.BNReLU(inner_features, bn_type=bn_type))
self.conv3 = nn.Sequential(
nn.Conv2d(features,
inner_features,
kernel_size=3,
padding=dilations[0],
dilation=dilations[0],
bias=False),
ModuleHelper.BNReLU(inner_features, bn_type=bn_type))
self.conv4 = nn.Sequential(
nn.Conv2d(features,
inner_features,
kernel_size=3,
padding=dilations[1],
dilation=dilations[1],
bias=False),
ModuleHelper.BNReLU(inner_features, bn_type=bn_type))
self.conv5 = nn.Sequential(
nn.Conv2d(features,
inner_features,
kernel_size=3,
padding=dilations[2],
dilation=dilations[2],
bias=False),
ModuleHelper.BNReLU(inner_features, bn_type=bn_type))
self.bottleneck = nn.Sequential(
nn.Conv2d(inner_features * 5,
out_features,
kernel_size=1,
padding=0,
dilation=1,
bias=False),
ModuleHelper.BNReLU(out_features, bn_type=bn_type),
nn.Dropout2d(0.1))
def __init__(self, dim_in, proj_dim=256, proj='convmlp', bn_type='inplace_abn'):
super(ProjectionHead, self).__init__()
Log.info('proj_dim: {}'.format(proj_dim))
if proj == 'linear':
self.proj = nn.Conv2d(dim_in, proj_dim, kernel_size=1)
elif proj == 'convmlp':
self.proj = nn.Sequential(
nn.Conv2d(dim_in, dim_in, kernel_size=1),
ModuleHelper.BNReLU(dim_in, bn_type=bn_type),
nn.Conv2d(dim_in, proj_dim, kernel_size=1)
)
def __init__(self, configer):
self.inplanes = 128
super(FastAspOCNet, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
self.backbone = BackboneSelector(configer).get_backbone()
# extra added layers
self.down = nn.Sequential(
nn.Conv2d(2048, 512, kernel_size=3, stride=1, padding=1),
ModuleHelper.BNReLU(512, bn_type=self.configer.get('network', 'bn_type')),
)
# self.object_head = ObjectContext_Module(num_classes)
from lib.models.modules.fast_asp_oc_block import Fast_ASP_OC_Module
self.fast_oc_head = Fast_ASP_OC_Module(512, 512, num_classes=self.num_classes,
bn_type=self.configer.get('network', 'bn_type'))
self.head = nn.Conv2d(512, self.num_classes, kernel_size=1, stride=1, padding=0, bias=True)
self.dsn_head = nn.Sequential(
nn.Conv2d(1024, 512, kernel_size=3, stride=1, padding=1),
ModuleHelper.BNReLU(512, bn_type=self.configer.get('network', 'bn_type')),
nn.Conv2d(512, self.num_classes, kernel_size=1, stride=1, padding=0, bias=True)
)
def __init__(self, configer):
self.inplanes = 128
super(AspOCNet, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
self.backbone = BackboneSelector(configer).get_backbone()
# extra added layers
from lib.models.modules.asp_oc_block import ASP_OC_Module
self.context = nn.Sequential(
nn.Conv2d(2048, 512, kernel_size=3, stride=1, padding=1),
ModuleHelper.BNReLU(512,
bn_type=self.configer.get(
'network', 'bn_type')),
ASP_OC_Module(512,
512,
bn_type=self.configer.get('network', 'bn_type')),
)
self.cls = nn.Conv2d(512,
self.num_classes,
kernel_size=1,
stride=1,
padding=0,
bias=True)
self.dsn = nn.Sequential(
nn.Conv2d(1024, 512, kernel_size=3, stride=1, padding=1),
ModuleHelper.BNReLU(512,
bn_type=self.configer.get(
'network', 'bn_type')), nn.Dropout2d(0.10),
nn.Conv2d(512,
self.num_classes,
kernel_size=1,
stride=1,
padding=0,
bias=True))
self.context.apply(self.weights_init)
self.cls.apply(self.weights_init)
self.dsn.apply(self.weights_init)
def __init__(self,
in_channels,
key_channels,
value_channels,
out_channels,
down_factors=[[8, 8]],
dropout=0,
bn_type=None):
super(ISA_Module, self).__init__()
assert isinstance(down_factors, (tuple, list))
self.down_factors = down_factors
self.stages = nn.ModuleList([
ISA_Block(in_channels, key_channels, value_channels, out_channels,
d, bn_type) for d in down_factors
])
concat_channels = in_channels + out_channels
if len(self.down_factors) > 1:
self.up_conv = nn.Sequential(
nn.Conv2d(in_channels,
len(self.down_factors) * out_channels,
kernel_size=1,
padding=0,
bias=False),
ModuleHelper.BNReLU(len(self.down_factors) * out_channels,
bn_type=bn_type),
)
concat_channels = out_channels * len(self.down_factors) * 2
self.conv_bn = nn.Sequential(
nn.Conv2d(concat_channels, out_channels, kernel_size=1,
bias=False),
ModuleHelper.BNReLU(out_channels, bn_type=bn_type),
nn.Dropout2d(dropout),
)
def __init__(self, configer):
self.inplanes = 128
super(PyramidOCNet, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
self.backbone = BackboneSelector(configer).get_backbone()
self.layer5 = nn.Sequential(
nn.Conv2d(2048, 512, kernel_size=3, stride=1, padding=1),
ModuleHelper.BNReLU(512,
bn_type=self.configer.get(
'network', 'bn_type')))
# extra added layers
from lib.models.modules.pyramid_oc_block import Pyramid_OC_Module
self.context = Pyramid_OC_Module(in_channels=512,
out_channels=512,
dropout=0.05,
sizes=([1, 2, 3, 6]),
bn_type=self.configer.get(
'network', 'bn_type'))
self.cls = nn.Conv2d(512,
self.num_classes,
kernel_size=1,
stride=1,
padding=0,
bias=True)
self.dsn = nn.Sequential(
nn.Conv2d(1024, 512, kernel_size=3, stride=1, padding=1),
ModuleHelper.BNReLU(512,
bn_type=self.configer.get(
'network', 'bn_type')), nn.Dropout2d(0.05),
nn.Conv2d(512,
self.num_classes,
kernel_size=1,
stride=1,
padding=0,
bias=True))
def __init__(self, fc_dim=4096, bn_type=None):
super(PPMBilinearDeepsup, self).__init__()
self.bn_type = bn_type
pool_scales = (1, 2, 3, 6)
self.ppm = []
# assert bn_type == 'syncbn' or not self.training
# Torch BN can't handle feature map size with 1x1.
for scale in pool_scales:
self.ppm.append(nn.Sequential(
nn.AdaptiveAvgPool2d(scale),
nn.Conv2d(fc_dim, 512, kernel_size=1, bias=False),
ModuleHelper.BNReLU(512, bn_type=bn_type)
))
self.ppm = nn.ModuleList(self.ppm)
def __init__(self, configer):
self.inplanes = 128
super(ISANet, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
self.backbone = BackboneSelector(configer).get_backbone()
# extra added layers
bn_type = self.configer.get('network', 'bn_type')
factors = self.configer.get('network', 'factors')
from lib.models.modules.isa_block import ISA_Module
self.isa_head = nn.Sequential(
nn.Conv2d(2048, 512, kernel_size=3, stride=1, padding=1, bias=False),
ModuleHelper.BNReLU(512, bn_type=bn_type),
ISA_Module(in_channels=512, key_channels=256, value_channels=512,
out_channels=512, down_factors=factors, dropout=0.05, bn_type=bn_type),
)
self.cls_head = nn.Conv2d(512, self.num_classes, kernel_size=1, stride=1, padding=0, bias=True)
self.dsn_head = nn.Sequential(
nn.Conv2d(1024, 512, kernel_size=3, stride=1, padding=1, bias=False),
ModuleHelper.BNReLU(512, bn_type=bn_type),
nn.Dropout2d(0.05),
nn.Conv2d(512, self.num_classes, kernel_size=1, stride=1, padding=0, bias=True),
)
def __init__(self,
in_channels,
out_channels,
key_channels,
value_channels,
dropout,
sizes=([1]),
bn_type=None):
super(BaseOC_Module, self).__init__()
self.stages = []
self.stages = nn.ModuleList([self._make_stage(in_channels, in_channels,
key_channels, value_channels, size, bn_type) for size in sizes])
self.conv_bn_dropout = nn.Sequential(
nn.Conv2d(2*in_channels, out_channels, kernel_size=1, padding=0),
ModuleHelper.BNReLU(out_channels, bn_type=bn_type),
nn.Dropout2d(dropout)
)
def __init__(self, num_classes, bn_type=None):
super(DeepLabHead, self).__init__()
# auxiliary loss
self.layer_dsn = nn.Sequential(
nn.Conv2d(1024, 256, kernel_size=3, stride=1, padding=1),
ModuleHelper.BNReLU(256, bn_type=bn_type),
nn.Conv2d(256,
num_classes,
kernel_size=1,
stride=1,
padding=0,
bias=True))
# main pipeline
self.layer_aspp = ASPPModule(2048, 512, bn_type=bn_type)
self.refine = nn.Sequential(
nn.Conv2d(512, 512, kernel_size=3, padding=1, stride=1,
bias=False),
ModuleHelper.BatchNorm2d(bn_type=bn_type)(512),
nn.Conv2d(512, num_classes, kernel_size=1, stride=1, bias=True))