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