def vgg(cfg, progress, **kwargs):
pretrained = kwargs['pretrained']
batch_norm = kwargs['batch_norm']
model = VGG(
make_layers(cfgs[cfg], conv_kernel_size=3, batch_norm=batch_norm),
**kwargs)
if pretrained:
if batch_norm:
model_url = 'https://download.pytorch.org/models/vgg16_bn-6c64b313.pth'
else:
model_url = 'https://download.pytorch.org/models/vgg16-397923af.pth'
state_dict = load_state_dict_from_url(model_url, progress=progress)
state_dict['features.0.weight'] = state_dict[
'features.0.weight'][:, 0, :, :].unsqueeze(1)
state_dict = {
param_name: param
for param_name, param in state_dict.items()
if not ('classifier' in param_name)
}
model.load_state_dict(state_dict, strict=False)
return model
def __init__(self, subnets, rnns):
super().__init__()
assert len(subnets) == len(rnns)
self.blocks = len(subnets)
for index, (params, rnn) in enumerate(zip(subnets, rnns), 1):
setattr(self, 'stage' + str(index), make_layers(params))
setattr(self, 'rnn' + str(index), rnn)
def __init__(self, subnets, rnns, config=None):
super().__init__()
self.config = config
assert len(subnets) == len(rnns)
self.blocks = len(subnets)
for index, (params, rnn) in enumerate(zip(subnets, rnns)):
setattr(self, 'rnn' + str(self.blocks - index), rnn)
setattr(self, 'stage' + str(self.blocks - index),
make_layers(params))
def __init__(self, lock_features: bool = True):
super().__init__()
# This model uses a “two-branch multi-stage” CNN
# Two branch means that the CNN produces two different outputs.
# Multi-stage means that the network is stacked one on top of the other at every stage
# these layers have no relu layer
no_relu_layers = [
"conv5_5_CPM_L1",
"conv5_5_CPM_L2",
"Mconv7_stage2_L1",
"Mconv7_stage2_L2",
"Mconv7_stage3_L1",
"Mconv7_stage3_L2",
"Mconv7_stage4_L1",
"Mconv7_stage4_L2",
"Mconv7_stage5_L1",
"Mconv7_stage5_L2",
"Mconv7_stage6_L1",
"Mconv7_stage6_L1",
]
blocks = {}
# CNN layers of VGG-19
block0 = OrderedDict(
[
("conv1_1", [3, 64, 3, 1, 1]),
("conv1_2", [64, 64, 3, 1, 1]),
("pool1_stage1", [2, 2, 0]),
("conv2_1", [64, 128, 3, 1, 1]),
("conv2_2", [128, 128, 3, 1, 1]),
("pool2_stage1", [2, 2, 0]),
("conv3_1", [128, 256, 3, 1, 1]),
("conv3_2", [256, 256, 3, 1, 1]),
("conv3_3", [256, 256, 3, 1, 1]),
("conv3_4", [256, 256, 3, 1, 1]),
("pool3_stage1", [2, 2, 0]),
("conv4_1", [256, 512, 3, 1, 1]),
("conv4_2", [512, 512, 3, 1, 1]),
("conv4_3_CPM", [512, 256, 3, 1, 1]),
("conv4_4_CPM", [256, 128, 3, 1, 1]),
]
)
# Stage 1 - the network produces an initial set of
# - part affinity fields (PAFs) L
# - detection confidence maps S
# Below we use the notation blockX_1 to define blocks that belong to the first branch
# which predict the PAFs that represent a degree of association between different body parts
block1_1 = OrderedDict(
[
("conv5_1_CPM_L1", [128, 128, 3, 1, 1]),
("conv5_2_CPM_L1", [128, 128, 3, 1, 1]),
("conv5_3_CPM_L1", [128, 128, 3, 1, 1]),
("conv5_4_CPM_L1", [128, 512, 1, 1, 0]),
("conv5_5_CPM_L1", [512, 38, 1, 1, 0]),
]
)
# Below we use the notation blockX_2 to define blocks that belong to the second branch
# which predict the confidence maps of different body parts location
block1_2 = OrderedDict(
[
("conv5_1_CPM_L2", [128, 128, 3, 1, 1]),
("conv5_2_CPM_L2", [128, 128, 3, 1, 1]),
("conv5_3_CPM_L2", [128, 128, 3, 1, 1]),
("conv5_4_CPM_L2", [128, 512, 1, 1, 0]),
("conv5_5_CPM_L2", [512, 19, 1, 1, 0]),
]
)
blocks["block1_1"] = block1_1
blocks["block1_2"] = block1_2
self.model0 = make_layers(block0, no_relu_layers)
# Stages 2 - 6
# predictions from previous stages are iteratively refined
for i in range(2, 7):
blocks[f"block{i}_1"] = OrderedDict(
[
(f"Mconv1_stage{i}_L1", [185, 128, 7, 1, 3]),
(f"Mconv2_stage{i}_L1", [128, 128, 7, 1, 3]),
(f"Mconv3_stage{i}_L1", [128, 128, 7, 1, 3]),
(f"Mconv4_stage{i}_L1", [128, 128, 7, 1, 3]),
(f"Mconv5_stage{i}_L1", [128, 128, 7, 1, 3]),
(f"Mconv6_stage{i}_L1", [128, 128, 1, 1, 0]),
(f"Mconv7_stage{i}_L1", [128, 38, 1, 1, 0]),
]
)
blocks[f"block{i}_2"] = OrderedDict(
[
(f"Mconv1_stage{i}_L2", [185, 128, 7, 1, 3]),
(f"Mconv2_stage{i}_L2", [128, 128, 7, 1, 3]),
(f"Mconv3_stage{i}_L2", [128, 128, 7, 1, 3]),
(f"Mconv4_stage{i}_L2", [128, 128, 7, 1, 3]),
(f"Mconv5_stage{i}_L2", [128, 128, 7, 1, 3]),
(f"Mconv6_stage{i}_L2", [128, 128, 1, 1, 0]),
(f"Mconv7_stage{i}_L2", [128, 19, 1, 1, 0]),
]
)
for k in blocks.keys():
blocks[k] = make_layers(blocks[k], no_relu_layers)
# all the stages that belong to the first branch
self.model1_1 = blocks["block1_1"]
self.model2_1 = blocks["block2_1"]
self.model3_1 = blocks["block3_1"]
self.model4_1 = blocks["block4_1"]
self.model5_1 = blocks["block5_1"]
self.model6_1 = blocks["block6_1"]
# all the stages that belong to the second branch
self.model1_2 = blocks["block1_2"]
self.model2_2 = blocks["block2_2"]
self.model3_2 = blocks["block3_2"]
self.model4_2 = blocks["block4_2"]
self.model5_2 = blocks["block5_2"]
self.model6_2 = blocks["block6_2"]
if lock_features:
set_parameter_requires_grad(self, False)
# load pretrained model
model_dict = transfer(self, torch.load("../checkpoints/body_pose_model.pth"))
self.load_state_dict(model_dict)
def two_conv(cfg, **kwargs):
batch_norm = kwargs['batch_norm']
model = TwoConv(
make_layers(cfgs[cfg], conv_kernel_size=5, batch_norm=batch_norm),
**kwargs)
return model
def five_conv(cfg, **kwargs):
batch_norm = kwargs['batch_norm']
model = FiveConv(
make_layers(cfg=cfgs[cfg], conv_kernel_size=6, batch_norm=batch_norm),
**kwargs)
return model