Exemplo n.º 1
0
 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
         ),
     )
Exemplo n.º 2
0
 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)
Exemplo n.º 3
0
 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),
     )
Exemplo n.º 4
0
    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,
                )
            )
Exemplo n.º 5
0
 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),
     )