def _get_2d_pose_estimation_from_model(inp, model_pe, num_joints, num_blocks,
num_context_per_joint, full_trainable=False):
num_frames = K.int_shape(inp)[1]
stem = model_pe.get_layer('Stem')
stem.trainable = full_trainable
i = 1
recep_block = model_pe.get_layer('rBlock%d' % i)
recep_block.trainable = full_trainable
x1 = TimeDistributed(stem, name='td_%s' % stem.name)(inp)
xb1 = TimeDistributed(recep_block, name='td_%s' % recep_block.name)(x1)
inp_pe = Input(shape=K.int_shape(xb1)[2:])
sep_conv = model_pe.get_layer('SepConv%d' % i)
reg_map = model_pe.get_layer('RegMap%d' % i)
fre_map = model_pe.get_layer('fReMap%d' % i)
x2 = sep_conv(inp_pe)
x3 = fre_map(reg_map(x2))
x = add([inp_pe, x2, x3])
for i in range(2, num_blocks):
recep_block = model_pe.get_layer('rBlock%d' % i)
sep_conv = model_pe.get_layer('SepConv%d' % i)
reg_map = model_pe.get_layer('RegMap%d' % i)
fre_map = model_pe.get_layer('fReMap%d' % i)
x1 = recep_block(x)
x2 = sep_conv(x1)
x3 = fre_map(reg_map(x2))
x = add([x1, x2, x3])
recep_block = model_pe.get_layer('rBlock%d' % num_blocks)
sep_conv = model_pe.get_layer('SepConv%d' % num_blocks)
reg_map = model_pe.get_layer('RegMap%d' % num_blocks)
x = recep_block(x)
x = sep_conv(x)
x = reg_map(x)
model1 = Model(inp_pe, x, name='PoseReg')
model1.trainable = full_trainable
num_heatmaps = (num_context_per_joint + 1) * num_joints
num_rows = K.int_shape(model1.output)[1]
num_cols = K.int_shape(model1.output)[2]
sams_input_shape = (num_frames, num_rows, num_cols, num_joints)
samc_input_shape = \
(num_frames, num_rows, num_cols, num_heatmaps - num_joints)
# Build the time distributed models
model_pe.get_layer('sSAM').trainable = full_trainable
sam_s_model = TimeDistributed(model_pe.get_layer('sSAM'),
input_shape=sams_input_shape, name='sSAM')
if num_context_per_joint > 0:
model_pe.get_layer('cSAM').trainable = full_trainable
sam_c_model = TimeDistributed(model_pe.get_layer('cSAM'),
input_shape=samc_input_shape, name='cSAM')
model_pe.get_layer('sjProb').trainable = False
jprob_s_model = TimeDistributed(model_pe.get_layer('sjProb'),
input_shape=sams_input_shape, name='sjProb')
if num_context_per_joint > 0:
model_pe.get_layer('cjProb').trainable = False
jprob_c_model = TimeDistributed(model_pe.get_layer('cjProb'),
input_shape=samc_input_shape, name='cjProb')
agg_model = build_context_aggregation(num_joints,
num_context_per_joint, 0.8, num_frames=num_frames, name='Agg')
h = TimeDistributed(model1, name='td_Model1')(xb1)
if num_context_per_joint > 0:
hs = Lambda(lambda x: x[:,:,:,:, :num_joints])(h)
hc = Lambda(lambda x: x[:,:,:,:, num_joints:])(h)
else:
hs = h
ys = sam_s_model(hs)
if num_context_per_joint > 0:
yc = sam_c_model(hc)
pc = jprob_c_model(hc)
y = agg_model([ys, yc, pc])
else:
y = ys
p = jprob_s_model(Lambda(lambda x: 4*x)(hs))
hs = TimeDistributed(Activation(channel_softmax_2d()),
name='td_ChannelSoftmax')(hs)
return y, p, hs, xb1
def build(input_shape,
num_joints,
dim,
num_context_per_joint=None,
alpha=0.8,
num_blocks=4,
depth_maps=16,
ksize=(3, 3),
export_heatmaps=False,
export_vfeat_block=None,
old_model=False):
if dim == 2:
if num_context_per_joint is None:
num_context_per_joint = 2
num_heatmaps = (num_context_per_joint + 1) * num_joints
elif dim == 3:
assert num_context_per_joint == None, \
'For 3D pose estimation, contextual heat maps are not allowed.'
num_heatmaps = depth_maps * num_joints
else:
raise ValueError('"dim" must be 2 or 3 and not (%d)' % dim)
inp = Input(shape=input_shape)
outputs = []
vfeat = None
x = _stem(inp, old_model=old_model)
num_rows, num_cols, num_filters = K.int_shape(x)[1:]
# Build the soft-argmax models (no parameters) for specialized and
# contextual maps.
sams_input_shape = (num_rows, num_cols, num_joints)
sam_s_model = build_softargmax_2d(sams_input_shape, rho=0, name='sSAM')
jprob_s_model = build_joints_probability(sams_input_shape, name='sjProb')
# Build the aggregation model (no parameters)
if num_context_per_joint is not None:
samc_input_shape = (num_rows, num_cols, num_heatmaps - num_joints)
sam_c_model = build_softargmax_2d(samc_input_shape, rho=0, name='cSAM')
jprob_c_model = build_joints_probability(samc_input_shape,
name='cjProb')
agg_model = build_context_aggregation(num_joints,
num_context_per_joint,
alpha,
name='Agg')
if dim == 3:
samz_input_shape = (depth_maps, num_joints)
sam_z_model = build_softargmax_1d(samz_input_shape, name='zSAM')
for bidx in range(num_blocks):
block_shape = K.int_shape(x)[1:]
x = build_reception_block(x, name='rBlock%d' % (bidx + 1), ksize=ksize)
if export_vfeat_block == (bidx + 1):
vfeat = x
ident_map = x
x = build_sconv_block(x, name='SepConv%d' % (bidx + 1), ksize=ksize)
h = build_regmap_block(x, num_heatmaps, name='RegMap%d' % (bidx + 1))
if dim == 2:
if num_context_per_joint is not None:
pose, visible, hm = pose_regression_2d_context(
h, num_joints, sam_s_model, sam_c_model, jprob_c_model,
agg_model, jprob_s_model)
else:
pose, visible, hm = pose_regression_2d(h, sam_s_model,
jprob_s_model)
else:
pose, visible, hm = pose_regression_3d(h, num_joints, depth_maps,
sam_s_model, sam_z_model)
outputs.append(pose)
outputs.append(visible)
if export_heatmaps:
outputs.append(hm)
if bidx < num_blocks - 1:
h = build_fremap_block(h,
block_shape[-1],
name='fReMap%d' % (bidx + 1))
x = add([ident_map, x, h])
if vfeat is not None:
outputs.append(vfeat)
model = Model(inputs=inp, outputs=outputs)
return model