def __init__(self, configer):
super(OpenPose, self).__init__()
self.configer = configer
self.backbone = BackboneSelector(configer).get_backbone()
self.pose_model = PoseModel(configer, self.backbone.get_num_features())
self.valid_loss_dict = configer.get('loss', 'loss_weights',
configer.get('loss.loss_type'))
def __init__(self, configer):
super(PSPNet, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
self.backbone = BackboneSelector(configer).get_backbone()
num_features = self.backbone.get_num_features()
self.dsn = nn.Sequential(
_ConvBatchNormReluBlock(num_features // 2,
num_features // 4,
3,
1,
norm_type=self.configer.get(
'network', 'norm_type')),
nn.Dropout2d(0.1),
nn.Conv2d(num_features // 4, self.num_classes, 1, 1, 0))
self.ppm = PPMBilinearDeepsup(fc_dim=num_features,
norm_type=self.configer.get(
'network', 'norm_type'))
self.cls = nn.Sequential(
nn.Conv2d(num_features + 4 * 512,
512,
kernel_size=3,
padding=1,
bias=False),
ModuleHelper.BNReLU(512,
norm_type=self.configer.get(
'network', 'norm_type')),
nn.Dropout2d(0.1), nn.Conv2d(512, self.num_classes, kernel_size=1))
self.valid_loss_dict = configer.get('loss', 'loss_weights',
configer.get('loss.loss_type'))
def __init__(self, configer):
super(NonLocalNet_nowd, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
self.backbone = BackboneSelector(configer).get_backbone()
num_features = self.backbone.get_num_features()
self.dsn = nn.Sequential(
_ConvBatchNormReluBlock(num_features // 2, num_features // 4, 3, 1,
norm_type=self.configer.get('network', 'norm_type')),
nn.Dropout2d(0.1),
nn.Conv2d(num_features // 4, self.num_classes, 1, 1, 0)
)
self.nlm = NLModule_nowd(num_features, 512, self.num_classes, self.configer)
self.valid_loss_dict = configer.get('loss', 'loss_weights', configer.get('loss.loss_type'))
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(self.backbone.get_num_features(),
norm_type=self.configer.get('network', 'norm_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, norm_type=self.configer.get('network', 'norm_type')),
nn.Dropout2d(0.1),
nn.Conv2d(512, self.num_classes, kernel_size=1, stride=1, padding=0, bias=True)
)
self.valid_loss_dict = configer.get('loss', 'loss_weights', configer.get('loss.loss_type'))
class DeepLabV3(nn.Module):
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(self.backbone.get_num_features(),
norm_type=self.configer.get('network', 'norm_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, norm_type=self.configer.get('network', 'norm_type')),
nn.Dropout2d(0.1),
nn.Conv2d(512, self.num_classes, kernel_size=1, stride=1, padding=0, bias=True)
)
self.valid_loss_dict = configer.get('loss', 'loss_weights', configer.get('loss.loss_type'))
def forward(self, data_dict):
x = self.backbone(data_dict['img'])
x_dsn = self.dsn(x[-2])
x = self.head(x[-1])
x_dsn = F.interpolate(x_dsn, size=(data_dict['img'].size(2), data_dict['img'].size(3)),
mode="bilinear", align_corners=True)
x = F.interpolate(x, size=(data_dict['img'].size(2), data_dict['img'].size(3)),
mode="bilinear", align_corners=True)
out_dict = dict(dsn_out=x_dsn, out=x)
if self.configer.get('phase') == 'test':
return out_dict
loss_dict = dict()
if 'dsn_ce_loss' in self.valid_loss_dict:
loss_dict['dsn_ce_loss'] = dict(
params=[x_dsn, data_dict['labelmap']],
type=torch.cuda.LongTensor([BASE_LOSS_DICT['ce_loss']]),
weight=torch.cuda.FloatTensor([self.valid_loss_dict['dsn_ce_loss']])
)
if 'ce_loss' in self.valid_loss_dict:
loss_dict['ce_loss'] = dict(
params=[x, data_dict['labelmap']],
type=torch.cuda.LongTensor([BASE_LOSS_DICT['ce_loss']]),
weight=torch.cuda.FloatTensor([self.valid_loss_dict['ce_loss']])
)
if 'ohem_ce_loss' in self.valid_loss_dict:
loss_dict['ohem_ce_loss'] = dict(
params=[x, data_dict['labelmap']],
type=torch.cuda.LongTensor([BASE_LOSS_DICT['ohem_ce_loss']]),
weight=torch.cuda.FloatTensor([self.valid_loss_dict['ohem_ce_loss']])
)
return out_dict, loss_dict
class NonLocalNet_nowd(nn.Sequential):
def __init__(self, configer):
super(NonLocalNet_nowd, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
self.backbone = BackboneSelector(configer).get_backbone()
num_features = self.backbone.get_num_features()
self.dsn = nn.Sequential(
_ConvBatchNormReluBlock(num_features // 2, num_features // 4, 3, 1,
norm_type=self.configer.get('network', 'norm_type')),
nn.Dropout2d(0.1),
nn.Conv2d(num_features // 4, self.num_classes, 1, 1, 0)
)
self.nlm = NLModule_nowd(num_features, 512, self.num_classes, self.configer)
self.valid_loss_dict = configer.get('loss', 'loss_weights', configer.get('loss.loss_type'))
def forward(self, data_dict):
x = self.backbone(data_dict['img'])
aux_x = self.dsn(x[-2])
x = self.nlm(x[-1])
x_dsn = F.interpolate(aux_x, size=(data_dict['img'].size(2), data_dict['img'].size(3)),
mode="bilinear", align_corners=True)
x = F.interpolate(x, size=(data_dict['img'].size(2), data_dict['img'].size(3)),
mode="bilinear", align_corners=True)
out_dict = dict(dsn_out=x_dsn, out=x)
if self.configer.get('phase') == 'test':
return out_dict
loss_dict = dict()
if 'dsn_ce_loss' in self.valid_loss_dict:
loss_dict['dsn_ce_loss'] = dict(
params=[x_dsn, data_dict['labelmap']],
type=torch.cuda.LongTensor([BASE_LOSS_DICT['ce_loss']]),
weight=torch.cuda.FloatTensor([self.valid_loss_dict['dsn_ce_loss']])
)
if 'ce_loss' in self.valid_loss_dict:
loss_dict['ce_loss'] = dict(
params=[x, data_dict['labelmap']],
type=torch.cuda.LongTensor([BASE_LOSS_DICT['ce_loss']]),
weight=torch.cuda.FloatTensor([self.valid_loss_dict['ce_loss']])
)
if 'ohem_ce_loss' in self.valid_loss_dict:
loss_dict['ohem_ce_loss'] = dict(
params=[x, data_dict['labelmap']],
type=torch.cuda.LongTensor([BASE_LOSS_DICT['ohem_ce_loss']]),
weight=torch.cuda.FloatTensor([self.valid_loss_dict['ohem_ce_loss']])
)
return out_dict, loss_dict
def __init__(self, configer):
super(CPN, self).__init__()
self.configer = configer
self.backbone = BackboneSelector(configer).get_backbone()
input_size = self.configer.get('data', 'input_size')
stride = self.configer.get('network', 'stride')
output_shape = (input_size[0] // stride, input_size[1] // stride)
self.global_net = globalNet([2048, 1024, 512, 256],
output_shape,
self.configer.get('network', 'heatmap_out'))
self.refine_net = refineNet(256, output_shape, self.configer.get('network', 'heatmap_out'))
def __init__(self, configer):
super(asymmetric_non_local_network, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
self.backbone = BackboneSelector(configer).get_backbone()
# low_in_channels, high_in_channels, out_channels, key_channels, value_channels, dropout
self.fusion = AFNB(1024,
2048,
2048,
256,
256,
dropout=0.05,
sizes=([1]),
norm_type=self.configer.get('network', 'norm_type'))
# extra added layers
self.context = nn.Sequential(
nn.Conv2d(2048, 512, kernel_size=3, stride=1, padding=1),
ModuleHelper.BNReLU(512,
norm_type=self.configer.get(
'network', 'norm_type')),
APNB(in_channels=512,
out_channels=512,
key_channels=256,
value_channels=256,
dropout=0.05,
sizes=([1]),
norm_type=self.configer.get('network', 'norm_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,
norm_type=self.configer.get(
'network', 'norm_type')),
nn.Dropout2d(0.05),
nn.Conv2d(512,
self.num_classes,
kernel_size=1,
stride=1,
padding=0,
bias=True))
self.valid_loss_dict = configer.get('loss', 'loss_weights',
configer.get('loss.loss_type'))
class OpenPose(nn.Module):
def __init__(self, configer):
super(OpenPose, self).__init__()
self.configer = configer
self.backbone = BackboneSelector(configer).get_backbone()
self.pose_model = PoseModel(configer, self.backbone.get_num_features())
self.valid_loss_dict = configer.get('loss', 'loss_weights',
configer.get('loss.loss_type'))
def forward(self, data_dict):
x = self.backbone(data_dict['img'])
paf_out, heatmap_out = self.pose_model(x)
out_dict = dict(paf=paf_out[-1], heatmap=heatmap_out[-1])
if self.configer.get('phase') == 'test':
return out_dict
loss_dict = dict()
for i in range(len(paf_out)):
if 'paf_loss{}'.format(i) in self.valid_loss_dict:
loss_dict['paf_loss{}'.format(i)] = dict(
params=[
paf_out[i] * data_dict['maskmap'],
data_dict['vecmap'] * data_dict['maskmap']
],
type=torch.cuda.LongTensor([BASE_LOSS_DICT['mse_loss']]),
weight=torch.cuda.FloatTensor(
[self.valid_loss_dict['paf_loss{}'.format(i)]]))
for i in range(len(heatmap_out)):
if 'heatmap_loss{}'.format(i) in self.valid_loss_dict:
loss_dict['heatmap_loss{}'.format(i)] = dict(
params=[
heatmap_out[i] * data_dict['maskmap'],
data_dict['heatmap'] * data_dict['maskmap']
],
type=torch.cuda.LongTensor([BASE_LOSS_DICT['mse_loss']]),
weight=torch.cuda.FloatTensor(
[self.valid_loss_dict['heatmap_loss{}'.format(i)]]))
return out_dict, loss_dict
def __init__(self, configer, loss_dict=None, flag=""):
super(ClsModel, self).__init__()
self.configer = configer
self.flag = flag if len(flag) == 0 else "{}_".format(flag)
self.backbone = BackboneSelector(configer).get_backbone(
backbone=configer.get('network.{}backbone'.format(self.flag)),
pretrained=configer.get('network.{}pretrained'.format(self.flag)))
self.reduction = None
fc_dim_out = configer.get('network.{}fc_dim'.format(self.flag),
default=None)
fc_dim = self.backbone.num_features
if fc_dim_out is not None:
self.reduction = nn.Conv2d(self.backbone.num_features, fc_dim_out,
1)
fc_dim = fc_dim_out
self.linear_list = nn.ModuleList()
linear_type = configer.get('network',
'{}linear_type'.format(self.flag),
default='nobias_linear')
self.fc = LINEAR_DICT[linear_type](fc_dim,
configer.get('data.num_classes'))
self.embed = None
if configer.get('network.{}embed'.format(self.flag), default=False):
feat_dim = configer.get('network', '{}feat_dim'.format(self.flag))
self.embed = nn.Sequential(nn.Linear(fc_dim, feat_dim),
nn.BatchNorm1d(feat_dim))
self.bn = None
if configer.get('network.{}bn'.format(self.flag), default=False):
self.bn = nn.BatchNorm1d(fc_dim)
nn.init.zeros_(self.bn.bias)
self.bn.bias.requires_grad = False
self.valid_loss_dict = configer.get(
'loss.loss_weights',
configer.get('loss.loss_type')) if loss_dict is None else loss_dict
def __init__(self, configer):
super(FpnRCNN, self).__init__()
self.configer = configer
self.backbone = BackboneSelector(configer).get_backbone()
self.RCNN_layer0 = nn.Sequential(self.backbone.conv1,
self.backbone.bn1, self.backbone.relu,
self.backbone.maxpool)
self.RCNN_layer1 = nn.Sequential(self.backbone.layer1)
self.RCNN_layer2 = nn.Sequential(self.backbone.layer2)
self.RCNN_layer3 = nn.Sequential(self.backbone.layer3)
self.RCNN_layer4 = nn.Sequential(self.backbone.layer4)
# Top layer
self.RCNN_toplayer = nn.Conv2d(2048,
256,
kernel_size=1,
stride=1,
padding=0) # reduce channel
# Smooth layers
self.RCNN_smooth1 = nn.Conv2d(256,
256,
kernel_size=3,
stride=1,
padding=1)
self.RCNN_smooth2 = nn.Conv2d(256,
256,
kernel_size=3,
stride=1,
padding=1)
self.RCNN_smooth3 = nn.Conv2d(256,
256,
kernel_size=3,
stride=1,
padding=1)
# Lateral layers
self.RCNN_latlayer1 = nn.Conv2d(1024,
256,
kernel_size=1,
stride=1,
padding=0)
self.RCNN_latlayer2 = nn.Conv2d(512,
256,
kernel_size=1,
stride=1,
padding=0)
self.RCNN_latlayer3 = nn.Conv2d(256,
256,
kernel_size=1,
stride=1,
padding=0)
# ROI Pool feature downsampling
self.RCNN_roi_feat_ds = nn.Conv2d(256,
256,
kernel_size=3,
stride=2,
padding=1)
self.RCNN_top = nn.Sequential(
nn.Conv2d(256,
1024,
kernel_size=self.configer.get('roi', 'pooled_height'),
stride=self.configer.get('roi', 'pooled_height'),
padding=0), nn.ReLU(True),
nn.Conv2d(1024, 1024, kernel_size=1, stride=1, padding=0),
nn.ReLU(True))
self.rpn = NaiveRPN(configer)
self.roi = FRROIGenerator(configer)
self.roi_sampler = FRROISampler(configer)
self.head = RoIHead(configer)
def __init__(self, configer):
super(DenseASPP, self).__init__()
self.configer = configer
dropout0 = MODEL_CONFIG['dropout0']
dropout1 = MODEL_CONFIG['dropout1']
d_feature0 = MODEL_CONFIG['d_feature0']
d_feature1 = MODEL_CONFIG['d_feature1']
self.backbone = BackboneSelector(configer).get_backbone()
num_features = self.backbone.get_num_features()
self.trans = _Transition(num_input_features=self.num_features,
num_output_features=self.num_features // 2,
norm_type=self.configer.get(
'network', 'norm_type'))
self.num_features = self.num_features // 2
self.ASPP_3 = _DenseAsppBlock(input_num=num_features,
num1=d_feature0,
num2=d_feature1,
dilation_rate=3,
drop_out=dropout0,
norm_type=self.configer.get(
'network', 'norm_type'))
self.ASPP_6 = _DenseAsppBlock(input_num=num_features + d_feature1 * 1,
num1=d_feature0,
num2=d_feature1,
dilation_rate=6,
drop_out=dropout0,
norm_type=self.configer.get(
'network', 'norm_type'))
self.ASPP_12 = _DenseAsppBlock(input_num=num_features + d_feature1 * 2,
num1=d_feature0,
num2=d_feature1,
dilation_rate=12,
drop_out=dropout0,
norm_type=self.configer.get(
'network', 'norm_type'))
self.ASPP_18 = _DenseAsppBlock(input_num=num_features + d_feature1 * 3,
num1=d_feature0,
num2=d_feature1,
dilation_rate=18,
drop_out=dropout0,
norm_type=self.configer.get(
'network', 'norm_type'))
self.ASPP_24 = _DenseAsppBlock(input_num=num_features + d_feature1 * 4,
num1=d_feature0,
num2=d_feature1,
dilation_rate=24,
drop_out=dropout0,
norm_type=self.configer.get(
'network', 'norm_type'))
num_features = num_features + 5 * d_feature1
self.classification = nn.Sequential(
nn.Dropout2d(p=dropout1),
nn.Conv2d(num_features,
self.configer.get('data', 'num_classes'),
kernel_size=1,
padding=0))
self.valid_loss_dict = configer.get('loss', 'loss_weights',
configer.get('loss.loss_type'))
class DenseASPP(nn.Module):
"""
* output_scale can only set as 8 or 16
"""
def __init__(self, configer):
super(DenseASPP, self).__init__()
self.configer = configer
dropout0 = MODEL_CONFIG['dropout0']
dropout1 = MODEL_CONFIG['dropout1']
d_feature0 = MODEL_CONFIG['d_feature0']
d_feature1 = MODEL_CONFIG['d_feature1']
self.backbone = BackboneSelector(configer).get_backbone()
num_features = self.backbone.get_num_features()
self.trans = _Transition(num_input_features=self.num_features,
num_output_features=self.num_features // 2,
norm_type=self.configer.get(
'network', 'norm_type'))
self.num_features = self.num_features // 2
self.ASPP_3 = _DenseAsppBlock(input_num=num_features,
num1=d_feature0,
num2=d_feature1,
dilation_rate=3,
drop_out=dropout0,
norm_type=self.configer.get(
'network', 'norm_type'))
self.ASPP_6 = _DenseAsppBlock(input_num=num_features + d_feature1 * 1,
num1=d_feature0,
num2=d_feature1,
dilation_rate=6,
drop_out=dropout0,
norm_type=self.configer.get(
'network', 'norm_type'))
self.ASPP_12 = _DenseAsppBlock(input_num=num_features + d_feature1 * 2,
num1=d_feature0,
num2=d_feature1,
dilation_rate=12,
drop_out=dropout0,
norm_type=self.configer.get(
'network', 'norm_type'))
self.ASPP_18 = _DenseAsppBlock(input_num=num_features + d_feature1 * 3,
num1=d_feature0,
num2=d_feature1,
dilation_rate=18,
drop_out=dropout0,
norm_type=self.configer.get(
'network', 'norm_type'))
self.ASPP_24 = _DenseAsppBlock(input_num=num_features + d_feature1 * 4,
num1=d_feature0,
num2=d_feature1,
dilation_rate=24,
drop_out=dropout0,
norm_type=self.configer.get(
'network', 'norm_type'))
num_features = num_features + 5 * d_feature1
self.classification = nn.Sequential(
nn.Dropout2d(p=dropout1),
nn.Conv2d(num_features,
self.configer.get('data', 'num_classes'),
kernel_size=1,
padding=0))
self.valid_loss_dict = configer.get('loss', 'loss_weights',
configer.get('loss.loss_type'))
def forward(self, data_dict):
x = self.backbone(data_dict['img'])
feature = self.trans(x)
aspp3 = self.ASPP_3(feature)
feature = torch.cat((aspp3, feature), dim=1)
aspp6 = self.ASPP_6(feature)
feature = torch.cat((aspp6, feature), dim=1)
aspp12 = self.ASPP_12(feature)
feature = torch.cat((aspp12, feature), dim=1)
aspp18 = self.ASPP_18(feature)
feature = torch.cat((aspp18, feature), dim=1)
aspp24 = self.ASPP_24(feature)
feature = torch.cat((aspp24, feature), dim=1)
x = self.classification(feature)
x = F.interpolate(x,
size=(data_dict['img'].size(2),
data_dict['img'].size(3)),
mode="bilinear",
align_corners=True)
out_dict = dict(out=x)
if self.configer.get('phase') == 'test':
return out_dict
loss_dict = dict()
if 'ce_loss' in self.valid_loss_dict:
loss_dict['ce_loss'] = dict(params=[x, data_dict['labelmap']],
type=torch.cuda.LongTensor(
[BASE_LOSS_DICT['ce_loss']]),
weight=torch.cuda.FloatTensor(
[self.valid_loss_dict['ce_loss']]))
if 'ohem_ce_loss' in self.valid_loss_dict:
loss_dict['ohem_ce_loss'] = dict(
params=[x, data_dict['labelmap']],
type=torch.cuda.LongTensor([BASE_LOSS_DICT['ohem_ce_loss']]),
weight=torch.cuda.FloatTensor(
[self.valid_loss_dict['ohem_ce_loss']]))
return out_dict, loss_dict
def __init__(self, configer):
super(DarkNetYolov3, self).__init__()
self.configer = configer
self.backbone = BackboneSelector(configer).get_backbone()
self.yolov3_head = Yolov3Head(configer, out_filters=self.backbone.num_features)
self.yolo_detection_layer = YOLODetectionLayer(self.configer)