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 _get_3d_pose_estimation_from_model(inp,
model_pe,
num_joints,
num_blocks,
depth_maps,
full_trainable=False):
num_frames = K.int_shape(inp)[1]
model_pe.summary()
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_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)
samz_input_shape = (num_frames, depth_maps, 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')
model_pe.get_layer('zSAM').trainable = full_trainable
sam_z_model = TimeDistributed(model_pe.get_layer('zSAM'),
input_shape=samz_input_shape,
name='zSAM')
h = TimeDistributed(model1, name='td_Model1')(xb1)
assert K.int_shape(h)[-1] == depth_maps * num_joints
def _reshape_heatmaps(x):
x = K.expand_dims(x, axis=-1)
x = K.reshape(x, (-1, K.int_shape(x)[1], K.int_shape(x)[2],
K.int_shape(x)[3], depth_maps, num_joints))
return x
h = Lambda(_reshape_heatmaps)(h)
hxy = Lambda(lambda x: K.mean(x, axis=4))(h)
hz = Lambda(lambda x: K.mean(x, axis=(2, 3)))(h)
pxy = sam_s_model(hxy)
pz = sam_z_model(hz)
pose = concatenate([pxy, pz])
vxy = TimeDistributed(GlobalMaxPooling2D(),
name='td_GlobalMaxPooling2D',
input_shape=K.int_shape(hxy)[1:])(hxy)
vz = TimeDistributed(GlobalMaxPooling1D(),
name='td_GlobalMaxPooling1D',
input_shape=K.int_shape(hz)[1:])(hz)
v = add([vxy, vz])
v = Lambda(lambda x: 2 * K.expand_dims(x, axis=-1))(v)
visible = Activation('sigmoid')(v)
hxy = TimeDistributed(Activation(channel_softmax_2d()),
name='td_ChannelSoftmax')(hxy)
return pose, visible, hxy, xb1
def act_channel_softmax(x, name=None):
x = Activation(channel_softmax_2d(), name=name)(x)
return x