def __init__(self,
channles,
down_resolution=None,
norm_type='BN',
act_type='prelu',
num_group=None):
super().__init__()
self.down_resolution = down_resolution
self.neck = nn.Sequential(*(optUnit(
norm_type=norm_type, act_type=act_type, out_ch=channles + 13) +
optUnit(opt_type='conv',
in_ch=channles + 13,
out_ch=channles,
ker_size=1,
stride=1)))
self.hourglass = HourglassBlock1(channles,
norm_type=norm_type,
act_type=act_type,
num_group=num_group)
self.foot = nn.Sequential(*optUnit(
norm_type=norm_type, out_ch=channles * 2, num_group=num_group))
def __init__(self, hourglass_channels, boundaries):
super(MPIBlock, self).__init__()
# B, 256, 64, 64 ->B, 13, 64, 64
self.up = nn.Sequential(*optUnit(opt_type='conv',
in_ch=hourglass_channels * boundaries,
out_ch=16 * boundaries,
ker_size=7,
stride=1,
conv_groups=boundaries))
# up 1
self.up1 = nn.Sequential(
*optUnit(norm_type='BN', act_type="prelu", out_ch=16))
# up 2
self.up1to2 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.up2 = nn.Sequential(
*optUnit(norm_type='BN', act_type='prelu', out_ch=16))
# up 3
self.up2to3 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.up3 = nn.Sequential(
*optUnit(norm_type='BN', act_type='prelu', out_ch=16))
# up 4
self.up1to4 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.up3to4 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.up4 = nn.Sequential(
*optUnit(norm_type='BN', act_type='prelu', out_ch=16))
# up 5
self.up2to5 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.up4to5 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.up5 = nn.Sequential(
*optUnit(norm_type='BN', act_type='prelu', out_ch=16))
# up 6
self.up1to6 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.up5to6 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.up6 = nn.Sequential(
*optUnit(norm_type='BN', act_type='prelu', out_ch=16))
# up 7
self.up1to7 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.up6to7 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.up7 = nn.Sequential(
*optUnit(norm_type='BN', act_type='prelu', out_ch=16))
# up 8 and 11
self.up1to8 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.up7to8 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.up8 = nn.Sequential(
*optUnit(norm_type='BN', act_type='prelu', out_ch=16))
self.up1to11 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.up7to11 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.up11 = nn.Sequential(
*optUnit(norm_type='BN', act_type='prelu', out_ch=16))
# up 9 and 12
self.up7to9 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.up8to9 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.up9 = nn.Sequential(
*optUnit(norm_type='BN', act_type='prelu', out_ch=16))
self.up7to12 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.up11to12 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.up12 = nn.Sequential(
*optUnit(norm_type='BN', act_type='prelu', out_ch=16))
# up 10 and 13
self.up7to10 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.up9to10 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.up10 = nn.Sequential(
*optUnit(norm_type='BN', act_type='prelu', out_ch=16))
self.up7to13 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.up12to13 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.up13 = nn.Sequential(
*optUnit(norm_type='BN', act_type='prelu', out_ch=16))
self.up_foot = nn.Sequential(*optUnit(opt_type='conv',
in_ch=16 * boundaries,
out_ch=hourglass_channels,
ker_size=1,
stride=1))
# down
self.down = nn.Sequential(*optUnit(opt_type='conv',
in_ch=hourglass_channels *
boundaries,
out_ch=16 * boundaries,
ker_size=7,
stride=1,
conv_groups=boundaries))
# down 13 and 10
self.down13 = nn.Sequential(
*optUnit(norm_type='BN', act_type="prelu", out_ch=16))
self.down10 = nn.Sequential(
*optUnit(norm_type='BN', act_type="prelu", out_ch=16))
# down 9 and 12
self.down10to9 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.down9 = nn.Sequential(
*optUnit(norm_type='BN', act_type='prelu', out_ch=16))
self.down13to12 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.down12 = nn.Sequential(
*optUnit(norm_type='BN', act_type='prelu', out_ch=16))
# down 8 and 11
self.down9to8 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.down8 = nn.Sequential(
*optUnit(norm_type='BN', act_type='prelu', out_ch=16))
self.down12to11 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.down11 = nn.Sequential(
*optUnit(norm_type='BN', act_type='prelu', out_ch=16))
# down 7
self.down9to7 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.down10to7 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.down13to7 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.down12to7 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.down8to7 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.down11to7 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.down7 = nn.Sequential(
*optUnit(norm_type='BN', act_type='prelu', out_ch=16))
# down 6
self.down7to6 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.down6 = nn.Sequential(
*optUnit(norm_type='BN', act_type='prelu', out_ch=16))
# down 5
self.down6to5 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.down5 = nn.Sequential(
*optUnit(norm_type='BN', act_type='prelu', out_ch=16))
# down 4
self.down5to4 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.down4 = nn.Sequential(
*optUnit(norm_type='BN', act_type='prelu', out_ch=16))
# down 3
self.down4to3 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.down3 = nn.Sequential(
*optUnit(norm_type='BN', act_type='prelu', out_ch=16))
# down 2
self.down3to2 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.down5to2 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.down2 = nn.Sequential(
*optUnit(norm_type='BN', act_type='prelu', out_ch=16))
# down 1
self.down2to1 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.down4to1 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.down6to1 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.down7to1 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.down8to1 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.down11to1 = nn.Sequential(*(
optUnit(opt_type='conv',
norm_type='BN',
act_type='prelu',
in_ch=16,
out_ch=8,
ker_size=7,
stride=1) +
optUnit(opt_type='conv', in_ch=8, out_ch=16, ker_size=7, stride=1)
))
self.down1 = nn.Sequential(
*optUnit(norm_type='BN', act_type='prelu', out_ch=16))
self.down_foot = nn.Sequential(*optUnit(opt_type='conv',
in_ch=16 * boundaries,
out_ch=hourglass_channels,
ker_size=1,
stride=1))
def __init__(self,
channels,
norm_type='BN',
act_type='prelu',
num_group=None):
super(HourglassBlock1, self).__init__()
# Lower 64
self.lower_res0_branch = ResidualBlock(channels,
channels,
1,
3,
norm_type=norm_type,
num_group=num_group,
act_type=act_type)
# pooling1 32
self.lower_res1_main = nn.Sequential(*[
nn.MaxPool2d(2, 2),
ResidualBlock(channels,
channels,
1,
3,
norm_type=norm_type,
num_group=num_group,
act_type=act_type)
])
self.lower_res1_branch = ResidualBlock(channels,
channels,
1,
3,
norm_type=norm_type,
num_group=num_group,
act_type=act_type)
# pooling 16
self.lower_res2_main = nn.Sequential(*[
nn.MaxPool2d(2, 2),
ResidualBlock(channels,
channels,
1,
3,
norm_type=norm_type,
num_group=num_group,
act_type=act_type)
])
self.lower_res2_branch = ResidualBlock(channels,
channels,
1,
3,
norm_type=norm_type,
num_group=num_group,
act_type=act_type)
# middle
self.middle = ResidualBlock(channels,
channels,
1,
3,
norm_type=norm_type,
num_group=num_group,
act_type=act_type)
# Upper 32
self.upper_res2_main = nn.Sequential(*[
ResidualBlock(channels,
channels,
1,
3,
norm_type=norm_type,
num_group=num_group,
act_type=act_type),
nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
])
# Upper 64
self.upper_res1_main = nn.Sequential(*[
ResidualBlock(channels,
channels,
1,
3,
norm_type=norm_type,
num_group=num_group,
act_type=act_type),
nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
])
# end_layer
self.output = nn.Sequential(
ResidualBlock(channels,
channels,
1,
3,
norm_type=norm_type,
num_group=num_group,
act_type=act_type),
*optUnit(norm_type=norm_type, act_type=act_type, out_ch=channels))
def __init__(self,
in_channels,
out_channels,
stride,
res_type=1,
channel_fac=4,
norm_type='BN',
num_group=None,
act_type='relu'):
super(ResidualBlock, self).__init__()
self.stride = stride
self.res_type = res_type
if res_type is 1:
layers = optUnit(norm_type=norm_type, num_group=num_group)
layers += optUnit(opt_type='conv',
norm_type=norm_type,
act_type=act_type,
in_ch=in_channels,
out_ch=out_channels,
ker_size=3,
stride=1,
num_group=num_group)
layers += optUnit(opt_type='conv',
in_ch=in_channels,
out_ch=out_channels,
ker_size=3,
stride=stride,
num_group=num_group)
self.right = nn.Sequential(*layers)
if self.stride is 2:
self.left = nn.Sequential(*optUnit(opt_type='conv',
in_ch=in_channels,
out_ch=out_channels,
ker_size=3,
stride=2,
num_group=num_group))
elif res_type is 2:
left_layers = optUnit(opt_type='conv',
in_ch=in_channels,
out_ch=out_channels,
ker_size=1,
stride=stride,
num_group=num_group)
self.left = nn.Sequential(*left_layers)
# bottleneck
right_layers = optUnit(opt_type='conv',
norm_type=norm_type,
act_type=act_type,
in_ch=in_channels,
out_ch=out_channels // channel_fac,
ker_size=1,
stride=1,
num_group=num_group)
right_layers += optUnit(opt_type='conv',
norm_type=norm_type,
act_type=act_type,
in_ch=out_channels // channel_fac,
out_ch=out_channels // channel_fac,
ker_size=1,
stride=stride,
num_group=num_group)
right_layers += optUnit(opt_type='conv',
in_ch=out_channels // channel_fac,
out_ch=out_channels,
ker_size=1,
stride=1,
num_group=num_group)
self.right = nn.Sequential(*right_layers)
elif res_type is 3:
right_layers = optUnit(norm_type=norm_type,
act_type=act_type,
out_ch=in_channels,
num_group=num_group)
right_layers += optUnit(opt_type='conv',
norm_type=norm_type,
act_type=act_type,
in_ch=in_channels,
out_ch=out_channels // channel_fac,
ker_size=1,
stride=1,
num_group=num_group)
right_layers += optUnit(opt_type='conv',
norm_type=norm_type,
act_type=act_type,
in_ch=out_channels // channel_fac,
out_ch=out_channels // channel_fac,
ker_size=3,
stride=1,
num_group=num_group)
right_layers += optUnit(opt_type='conv',
in_ch=out_channels // channel_fac,
out_ch=out_channels,
ker_size=1,
stride=1,
num_group=num_group)
self.right = nn.Sequential(*right_layers)
