def __init__(self, in_channels, out_channels, num_heatmaps, num_pafs):
super().__init__()
self.trunk = nn.Sequential(
RefinementStageBlock(in_channels, out_channels),
RefinementStageBlock(out_channels, out_channels),
RefinementStageBlock(out_channels, out_channels),
RefinementStageBlock(out_channels, out_channels),
RefinementStageBlock(out_channels, out_channels))
self.heatmaps = nn.Sequential(
conv(out_channels,
out_channels,
kernel_size=1,
padding=0,
bn=False),
conv(out_channels,
num_heatmaps,
kernel_size=1,
padding=0,
bn=False,
relu=False))
self.pafs = nn.Sequential(
conv(out_channels,
out_channels,
kernel_size=1,
padding=0,
bn=False),
conv(out_channels,
num_pafs,
kernel_size=1,
padding=0,
bn=False,
relu=False))
def __init__(self, in_channels, out_channels):
super().__init__()
self.initial = conv(in_channels,
out_channels,
kernel_size=1,
padding=0,
bn=False)
self.trunk = nn.Sequential(
conv(out_channels, out_channels),
conv(out_channels, out_channels, dilation=2, padding=2))
def __init__(self, in_channels, out_channels):
super().__init__()
self.align = conv(in_channels,
out_channels,
kernel_size=1,
padding=0,
bn=False)
self.trunk = nn.Sequential(conv_dw_no_bn(out_channels, out_channels),
conv_dw_no_bn(out_channels, out_channels),
conv_dw_no_bn(out_channels, out_channels))
self.conv = conv(out_channels, out_channels, bn=False)
def __init__(self,
num_refinement_stages=1,
num_channels=128,
num_heatmaps=19,
num_pafs=38):
super().__init__()
self.model = nn.Sequential(
conv(3, 32, stride=2, bias=False),
conv_dw(32, 64),
conv_dw(64, 128, stride=2),
conv_dw(128, 128),
conv_dw(128, 256, stride=2),
conv_dw(256, 256),
conv_dw(256, 512), # conv4_2
conv_dw(512, 512, dilation=2, padding=2),
conv_dw(512, 512),
conv_dw(512, 512),
conv_dw(512, 512),
conv_dw(512, 512) # conv5_5
)
self.cpm = Cpm(512, num_channels)
self.initial_stage = InitialStage(num_channels, num_heatmaps, num_pafs)
self.refinement_stages = nn.ModuleList()
for idx in range(num_refinement_stages):
self.refinement_stages.append(
RefinementStage(num_channels + num_heatmaps + num_pafs,
num_channels, num_heatmaps, num_pafs))
def __init__(self, num_channels, num_heatmaps, num_pafs):
super().__init__()
self.trunk = nn.Sequential(conv(num_channels, num_channels, bn=False),
conv(num_channels, num_channels, bn=False),
conv(num_channels, num_channels, bn=False))
self.heatmaps = nn.Sequential(
conv(num_channels, 512, kernel_size=1, padding=0, bn=False),
conv(512,
num_heatmaps,
kernel_size=1,
padding=0,
bn=False,
relu=False))
self.pafs = nn.Sequential(
conv(num_channels, 512, kernel_size=1, padding=0, bn=False),
conv(512, num_pafs, kernel_size=1, padding=0, bn=False,
relu=False))