def GradDescArchitecture(param_trainable, init_wrapper, smpl_params, input_info, faces, emb_size=1000):
    """ Gradient descent optimisation network architecture """
    # An embedding layer is required to optimise the parameters
    optlearner_input = Input(shape=(1,), name="embedding_index")

    # Initialise the embedding layers
    emb_layers = init_emb_layers(optlearner_input, emb_size, param_trainable, init_wrapper)
    print("emb layer shape: " + str(emb_layers[0].shape))
    optlearner_params = Concatenate(axis=1, name="parameter_embedding")(emb_layers)
    print("Embedding shape: " + str(optlearner_params.shape))
    #exit(1)
    optlearner_params = Reshape(target_shape=(85,), name="learned_params")(optlearner_params)
    print("optlearner parameters shape: " +str(optlearner_params.shape))
    #exit(1)

    # Ground truth parameters and point cloud are inputs to the model as well
    gt_params = Input(shape=(85,), name="gt_params")
    gt_pc = Input(shape=(6890, 3), name="gt_pc")
    print("gt parameters shape: " + str(gt_params.shape))
    print("gt point cloud shape: " + str(gt_pc.shape))

    # Compute the true offset (i.e. difference) between the ground truth and learned parameters
    #pi = K.constant(np.pi)
    delta_d = Lambda(lambda x: x[0] - x[1], name="delta_d")([gt_params, optlearner_params])
    #delta_d = Lambda(lambda x: x[0] - x[1], name="delta_d_no_mod")([gt_params, optlearner_params])
    #delta_d = Lambda(lambda x: K.tf.math.floormod(x - pi, 2*pi) - pi, name="delta_d")(delta_d)  # custom modulo 2pi of delta_d
    #delta_d = custom_mod(delta_d, pi, name="delta_d")  # custom modulo 2pi of delta_d
    print("delta_d shape: " + str(delta_d.shape))
    #exit(1)

    # Calculate the (batched) MSE between the learned parameters and the ground truth parameters
    #false_loss_delta_d = Lambda(lambda x: K.mean(K.square(x), axis=1))(delta_d)
    false_loss_delta_d = Lambda(lambda x: K.mean(K.square(x), axis=1))(delta_d)
    print("delta_d loss shape: " + str(false_loss_delta_d.shape))
    #exit(1)
    false_loss_delta_d = Reshape(target_shape=(1,), name="delta_d_mse")(false_loss_delta_d)
    print("delta_d loss shape: " + str(false_loss_delta_d.shape))

    # Load SMPL model and get necessary parameters
    optlearner_pc = get_pc(optlearner_params, smpl_params, input_info, faces)

    # Get the (batched) Euclidean loss between the learned and ground truth point clouds
    pc_euclidean_diff = Lambda(lambda x: x[0] -  x[1])([gt_pc, optlearner_pc])
    pc_euclidean_dist = Lambda(lambda x: K.sum(K.square(x),axis=-1))(pc_euclidean_diff)
    #pc_euclidean_dist = Lambda(lambda x: K.square(x))(pc_euclidean_diff)
    print('pc euclidean dist '+str(pc_euclidean_dist.shape))
    #exit(1)
    false_loss_pc = Lambda(lambda x: K.mean(x, axis=1))(pc_euclidean_dist)
    false_loss_pc = Reshape(target_shape=(1,), name="pc_mean_euc_dist")(false_loss_pc)
    #false_loss_pc = Lambda(lambda x: x, name="pc_mean_euc_dist")(pc_euclidean_dist)
    print("point cloud loss shape: " + str(false_loss_pc.shape))
    #exit(1)

    return [optlearner_input, gt_params, gt_pc], [optlearner_params, false_loss_delta_d, optlearner_pc, false_loss_pc]
示例#2
0
def standard_input(emb_size, param_trainable, init_wrapper):
    # An embedding layer is required to optimise the parameters
    optlearner_input = Input(shape=(1, ), name="embedding_index")

    # Initialise the embedding layers
    emb_layers = init_emb_layers(optlearner_input, emb_size, param_trainable,
                                 init_wrapper)
    optlearner_params = Concatenate(name="parameter_embedding")(emb_layers)
    optlearner_params = Reshape(target_shape=(85, ),
                                name="learned_params")(optlearner_params)
    print("optlearner parameters shape: " + str(optlearner_params.shape))
    #exit(1)

    # Ground truth parameters and point cloud are inputs to the model as well
    gt_params = Input(shape=(85, ), name="gt_params")
    gt_pc = Input(shape=(6890, 3), name="gt_pc")
    print("gt parameters shape: " + str(gt_params.shape))
    print("gt point cloud shape: " + str(gt_pc.shape))

    return optlearner_params, gt_params, gt_pc
def DeepConv1DOptLearnerStaticArchitecture(param_trainable,
                                           init_wrapper,
                                           smpl_params,
                                           input_info,
                                           faces,
                                           emb_size=1000,
                                           input_type="3D_POINTS"):
    """ Optimised learner network architecture """
    # An embedding layer is required to optimise the parameters
    optlearner_input = Input(shape=(1, ), name="embedding_index")

    # Initialise the embedding layers
    emb_layers = init_emb_layers(optlearner_input, emb_size, param_trainable,
                                 init_wrapper)
    optlearner_params = Concatenate(name="parameter_embedding")(emb_layers)
    optlearner_params = Reshape(target_shape=(85, ),
                                name="learned_params")(optlearner_params)
    print("optlearner parameters shape: " + str(optlearner_params.shape))
    #exit(1)

    # Ground truth parameters and point cloud are inputs to the model as well
    gt_params = Input(shape=(85, ), name="gt_params")
    gt_pc = Input(shape=(6890, 3), name="gt_pc")
    print("gt parameters shape: " + str(gt_params.shape))
    print("gt point cloud shape: " + str(gt_pc.shape))

    # Compute the true offset (i.e. difference) between the ground truth and learned parameters
    pi = K.constant(np.pi)
    delta_d = Lambda(lambda x: x[0] - x[1],
                     name="delta_d")([gt_params, optlearner_params])
    #delta_d = Lambda(lambda x: x[0] - x[1], name="delta_d_no_mod")([gt_params, optlearner_params])
    #delta_d = Lambda(lambda x: K.tf.math.floormod(x - pi, 2*pi) - pi, name="delta_d")(delta_d)  # custom modulo 2pi of delta_d
    #delta_d = custom_mod(delta_d, pi, name="delta_d")  # custom modulo 2pi of delta_d
    print("delta_d shape: " + str(delta_d.shape))
    #exit(1)

    # Calculate the (batched) MSE between the learned parameters and the ground truth parameters
    false_loss_delta_d = Lambda(lambda x: K.mean(K.square(x), axis=1))(delta_d)
    print("delta_d loss shape: " + str(false_loss_delta_d.shape))
    #exit(1)
    false_loss_delta_d = Reshape(target_shape=(1, ),
                                 name="delta_d_mse")(false_loss_delta_d)
    print("delta_d loss shape: " + str(false_loss_delta_d.shape))

    # Load SMPL model and get necessary parameters
    optlearner_pc = get_pc(optlearner_params, smpl_params, input_info,
                           faces)  # UNCOMMENT
    print("optlearner_pc shape: " + str(optlearner_pc.shape))
    #exit(1)
    #optlearner_pc = Dense(6890*3)(delta_d)
    #optlearner_pc = Reshape((6890, 3))(optlearner_pc)

    # Get the (batched) Euclidean loss between the learned and ground truth point clouds
    pc_euclidean_diff = Lambda(lambda x: x[0] - x[1])([gt_pc, optlearner_pc])
    pc_euclidean_dist = Lambda(lambda x: K.sum(K.square(x), axis=-1))(
        pc_euclidean_diff)
    print('pc euclidean dist ' + str(pc_euclidean_dist.shape))
    #exit(1)
    false_loss_pc = Lambda(lambda x: K.mean(x, axis=1))(pc_euclidean_dist)
    false_loss_pc = Reshape(target_shape=(1, ),
                            name="pc_mean_euc_dist")(false_loss_pc)
    print("point cloud loss shape: " + str(false_loss_pc.shape))
    #exit(1)

    # Gather sets of points and compute their cross product to get mesh normals
    # In order of: right hand, right wrist, right forearm, right bicep end, right bicep, right shoulder, top of cranium, left shoulder, left bicep, left bicep end, left forearm, left wrist, left hand,
    # chest, belly/belly button, back of neck, upper back, central back, lower back/tailbone,
    # left foot, left over-ankle, left shin, left over-knee, left quadricep, left hip, right, hip, right, quadricep, right over-knee, right shin, right, over-ankle, right foot
    vertex_list = [
        5674, 5705, 5039, 5151, 4977, 4198, 411, 606, 1506, 1682, 1571, 2244,
        2212, 3074, 3500, 460, 2878, 3014, 3021, 3365, 4606, 4588, 4671, 6877,
        1799, 5262, 3479, 1187, 1102, 1120, 6740
    ]
    #face_array = np.array([11396, 8620, 7866, 5431, 6460, 1732, 4507])
    pc_euclidean_diff_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(
        pc_euclidean_diff
    )  # This is added to avoid influencing embedding layer parameters by a "bad" gradient network
    vertex_diff_NOGRAD = Lambda(lambda x: K.tf.gather(
        x, np.array(vertex_list).astype(np.int32), axis=-2))(
            pc_euclidean_diff_NOGRAD)
    print("vertex_diff_NOGRAD shape: " + str(vertex_diff_NOGRAD.shape))
    vertex_diff_NOGRAD = Flatten()(vertex_diff_NOGRAD)
    #exit(1)
    face_array = np.array([[face for face in faces if vertex in face][0]
                           for vertex in vertex_list
                           ])  # only take a single face for each vertex
    print("face_array shape: " + str(face_array.shape))
    gt_normals = get_mesh_normals(gt_pc,
                                  face_array,
                                  layer_name="gt_cross_product")
    print("gt_normals shape: " + str(gt_normals.shape))
    opt_normals = get_mesh_normals(optlearner_pc,
                                   face_array,
                                   layer_name="opt_cross_product")
    print("opt_normals shape: " + str(opt_normals.shape))
    #exit(1)

    # Learn the offset in parameters from the difference between the ground truth and learned mesh normals
    diff_normals = Lambda(lambda x: K.tf.cross(x[0], x[1]),
                          name="diff_cross_product")([gt_normals, opt_normals])
    diff_normals_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(
        diff_normals
    )  # This is added to avoid influencing embedding layer parameters by a "bad" gradient network
    diff_angles = Lambda(lambda x: K.tf.subtract(x[0], x[1]),
                         name="diff_angle")([gt_normals, opt_normals])
    diff_angles_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(diff_angles)
    diff_angles_norm_NOGRAD = Lambda(
        lambda x: K.tf.norm(x, axis=-1),
        name="diff_angle_norm")(diff_angles_NOGRAD)
    dist_angles = Lambda(lambda x: K.mean(K.square(x), axis=-1),
                         name="diff_angle_mse")(diff_angles)
    dist_angles_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(dist_angles)
    print("diff_angles shape: " + str(diff_angles.shape))
    print("dist_angles shape: " + str(dist_angles.shape))
    #pc_euclidean_diff_NOGRAD =  Lambda(lambda x: K.stop_gradient(x))(pc_euclidean_diff) # This is added to avoid influencing embedding layer parameters by a "bad" gradient network
    #print("diff_normals_NOGRAD shape: " + str(diff_normals_NOGRAD.shape))
    diff_normals_NOGRAD = Flatten()(diff_normals_NOGRAD)
    diff_angles_NOGRAD = Flatten()(diff_angles_NOGRAD)
    mesh_diff_NOGRAD = Concatenate()([diff_normals_NOGRAD, dist_angles_NOGRAD])

    if input_type == "3D_POINTS":
        #optlearner_architecture = Dense(2**9, activation="relu")(vertex_diff_NOGRAD)
        optlearner_architecture = Dense(2**7,
                                        activation="relu")(vertex_diff_NOGRAD)
    if input_type == "MESH_NORMALS":
        #optlearner_architecture = Dense(2**11, activation="relu")(diff_angles_norm_NOGRAD)
        #optlearner_architecture = Dense(2**11, activation="relu")(diff_angles_NOGRAD)
        #optlearner_architecture = Dense(2**9, activation="relu")(mesh_diff_NOGRAD)
        optlearner_architecture = Dense(2**7,
                                        activation="relu")(mesh_diff_NOGRAD)
    #optlearner_architecture = BatchNormalization()(optlearner_architecture)
    #optlearner_architecture = Dropout(0.5)(optlearner_architecture)
    print('optlearner_architecture shape: ' +
          str(optlearner_architecture.shape))
    optlearner_architecture = Reshape(
        (optlearner_architecture.shape[1].value, 1))(optlearner_architecture)
    print('optlearner_architecture shape: ' +
          str(optlearner_architecture.shape))
    optlearner_architecture = Conv1D(
        64, 3, activation="relu")(optlearner_architecture)
    optlearner_architecture = Conv1D(
        64, 3, activation="relu")(optlearner_architecture)
    optlearner_architecture = BatchNormalization()(optlearner_architecture)
    optlearner_architecture = MaxPooling1D(2)(optlearner_architecture)
    #optlearner_architecture = MaxPooling1D(3)(optlearner_architecture)
    optlearner_architecture = Conv1D(
        128, 3, activation="relu")(optlearner_architecture)
    optlearner_architecture = Conv1D(
        128, 3, activation="relu")(optlearner_architecture)
    optlearner_architecture = BatchNormalization()(optlearner_architecture)
    optlearner_architecture = MaxPooling1D(3)(optlearner_architecture)
    optlearner_architecture = Conv1D(
        256, 3, activation="relu")(optlearner_architecture)
    optlearner_architecture = Conv1D(
        256, 3, activation="relu")(optlearner_architecture)
    optlearner_architecture = BatchNormalization()(optlearner_architecture)
    optlearner_architecture = MaxPooling1D(3)(optlearner_architecture)
    optlearner_architecture = Conv1D(
        512, 3, activation="relu")(optlearner_architecture)
    optlearner_architecture = Conv1D(
        512, 3, activation="relu")(optlearner_architecture)
    optlearner_architecture = BatchNormalization()(optlearner_architecture)
    #optlearner_architecture = AveragePooling1D(2)(optlearner_architecture)
    optlearner_architecture = GlobalAveragePooling1D()(optlearner_architecture)
    #print('optlearner_architecture shape: '+str(optlearner_architecture.shape))
    #optlearner_architecture = Flatten()(optlearner_architecture)
    print('optlearner_architecture shape: ' +
          str(optlearner_architecture.shape))
    #optlearner_architecture = Dropout(0.5)(optlearner_architecture)
    optlearner_architecture = Dense(2**7,
                                    activation="relu")(optlearner_architecture)
    print('optlearner_architecture shape: ' +
          str(optlearner_architecture.shape))
    #delta_d_hat = Dense(85, activation=pos_scaled_tanh, name="delta_d_hat")(optlearner_architecture)
    delta_d_hat = Dense(85, activation="linear",
                        name="delta_d_hat")(optlearner_architecture)
    #delta_d_hat = Dense(85, activation=centred_linear, name="delta_d_hat")(optlearner_architecture)
    print('delta_d_hat shape: ' + str(delta_d_hat.shape))
    #exit(1)

    # Calculate the (batched) MSE between the learned and ground truth offset in the parameters
    delta_d_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(delta_d)
    false_loss_delta_d_hat = Lambda(
        lambda x: K.mean(K.square(x[0] - x[1]), axis=1))(
            [delta_d_NOGRAD, delta_d_hat])
    #false_loss_delta_d_hat = Lambda(lambda x: K.sum(K.square(x[0] - x[1]), axis=1))([delta_d_NOGRAD, delta_d_hat])
    #false_loss_delta_d_hat = Lambda(lambda x: mape(x[0], x[1]))([delta_d_NOGRAD, delta_d_hat])
    false_loss_delta_d_hat = Reshape(
        target_shape=(1, ), name="delta_d_hat_mse")(false_loss_delta_d_hat)
    print("delta_d_hat loss shape: " + str(false_loss_delta_d_hat.shape))
    #false_sin_loss_delta_d_hat = get_sin_metric(delta_d_NOGRAD, delta_d_hat)
    false_sin_loss_delta_d_hat = get_sin_metric(delta_d_NOGRAD,
                                                delta_d_hat,
                                                average=False)
    false_sin_loss_delta_d_hat = Lambda(
        lambda x: x, name="delta_d_hat_sin_output")(false_sin_loss_delta_d_hat)
    print("delta_d_hat sin loss shape: " +
          str(false_sin_loss_delta_d_hat.shape))

    # Prevent model from using the delta_d_hat gradient in final loss
    delta_d_hat_NOGRAD = Lambda(lambda x: K.stop_gradient(x),
                                name='optlearner_output_NOGRAD')(delta_d_hat)

    # False loss designed to pass the learned offset as a gradient to the embedding layer
    false_loss_smpl = Multiply(name="smpl_diff")(
        [optlearner_params, delta_d_hat_NOGRAD])
    print("smpl loss shape: " + str(false_loss_smpl.shape))

    #return [optlearner_input, gt_params, gt_pc], [optlearner_params, false_loss_delta_d, optlearner_pc, false_loss_pc, false_loss_delta_d_hat, false_sin_loss_delta_d_hat,  false_loss_smpl, delta_d, delta_d_hat, delta_d_hat_NOGRAD]
    return [optlearner_input, gt_params, gt_pc], [
        optlearner_params, false_loss_delta_d, optlearner_pc, false_loss_pc,
        false_loss_delta_d_hat, false_sin_loss_delta_d_hat, false_loss_smpl,
        delta_d, delta_d_hat, dist_angles
    ]
def OptLearnerMeshNormalStaticModArchitecture(param_trainable,
                                              init_wrapper,
                                              smpl_params,
                                              input_info,
                                              faces,
                                              emb_size=1000):
    """ Optimised learner network architecture """
    # An embedding layer is required to optimise the parameters
    optlearner_input = Input(shape=(1, ), name="embedding_index")

    # Initialise the embedding layers
    emb_layers = init_emb_layers(optlearner_input, emb_size, param_trainable,
                                 init_wrapper)
    optlearner_params = Concatenate(name="parameter_embedding")(emb_layers)
    optlearner_params = Reshape(target_shape=(85, ),
                                name="learned_params")(optlearner_params)
    print("optlearner parameters shape: " + str(optlearner_params.shape))
    #exit(1)

    # Ground truth parameters and point cloud are inputs to the model as well
    gt_params = Input(shape=(85, ), name="gt_params")
    gt_pc = Input(shape=(6890, 3), name="gt_pc")
    print("gt parameters shape: " + str(gt_params.shape))
    print("gt point cloud shape: " + str(gt_pc.shape))

    # Compute the true offset (i.e. difference) between the ground truth and learned parameters
    pi = K.constant(np.pi)
    delta_d = Lambda(lambda x: x[0] - x[1],
                     name="delta_d")([gt_params, optlearner_params])
    #delta_d = Lambda(lambda x: x[0] - x[1], name="delta_d_no_mod")([gt_params, optlearner_params])
    #delta_d = Lambda(lambda x: K.tf.math.floormod(x - pi, 2*pi) - pi, name="delta_d")(delta_d)  # custom modulo 2pi of delta_d
    #delta_d = custom_mod(delta_d, pi, name="delta_d")  # custom modulo 2pi of delta_d
    print("delta_d shape: " + str(delta_d.shape))
    #exit(1)

    # Calculate the (batched) MSE between the learned parameters and the ground truth parameters
    false_loss_delta_d = Lambda(lambda x: K.mean(K.square(x), axis=1))(delta_d)
    print("delta_d loss shape: " + str(false_loss_delta_d.shape))
    #exit(1)
    false_loss_delta_d = Reshape(target_shape=(1, ),
                                 name="delta_d_mse")(false_loss_delta_d)
    print("delta_d loss shape: " + str(false_loss_delta_d.shape))

    # Load SMPL model and get necessary parameters
    optlearner_pc = get_pc(optlearner_params, smpl_params, input_info, faces)

    # Get the (batched) Euclidean loss between the learned and ground truth point clouds
    pc_euclidean_diff = Lambda(lambda x: x[0] - x[1])([gt_pc, optlearner_pc])
    pc_euclidean_dist = Lambda(lambda x: K.sum(K.square(x), axis=-1))(
        pc_euclidean_diff)
    print('pc euclidean dist ' + str(pc_euclidean_dist.shape))
    #exit(1)
    false_loss_pc = Lambda(lambda x: K.mean(x, axis=1))(pc_euclidean_dist)
    false_loss_pc = Reshape(target_shape=(1, ),
                            name="pc_mean_euc_dist")(false_loss_pc)
    print("point cloud loss shape: " + str(false_loss_pc.shape))
    #exit(1)

    # Gather sets of points and compute their cross product to get mesh normals
    #vertex_list=[1850, 1600, 2050, 1300, 5350, 5050, 5500]
    #vertex_list=[1850, 1600, 2050, 5350, 5050, 5500]
    # In order of right wrist, right, bicep, right shoulder, cranium top, left shoulder, left bicep, left wrist
    vertex_list = [5705, 4977, 4198, 411, 606, 1506, 2244]
    #face_array = np.array([11396, 8620, 7866, 5431, 6460, 1732, 4507])
    pc_euclidean_diff_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(
        pc_euclidean_diff
    )  # This is added to avoid influencing embedding layer parameters by a "bad" gradient network
    vertex_diff_NOGRAD = Lambda(lambda x: K.tf.gather(
        x, np.array(vertex_list).astype(np.int32), axis=-2))(
            pc_euclidean_diff_NOGRAD)
    print("vertex_diff_NOGRAD shape: " + str(vertex_diff_NOGRAD.shape))
    vertex_diff_NOGRAD = Flatten()(vertex_diff_NOGRAD)
    #exit(1)
    face_array = np.array([[face for face in faces if vertex in face][0]
                           for vertex in vertex_list
                           ])  # only take a single face for each vertex
    print("face_array shape: " + str(face_array.shape))
    gt_normals = get_mesh_normals(gt_pc,
                                  face_array,
                                  layer_name="gt_cross_product")
    print("gt_normals shape: " + str(gt_normals.shape))
    opt_normals = get_mesh_normals(optlearner_pc,
                                   face_array,
                                   layer_name="opt_cross_product")
    print("opt_normals shape: " + str(opt_normals.shape))
    #exit(1)

    # Learn the offset in parameters from the difference between the ground truth and learned mesh normals
    diff_normals = Lambda(lambda x: K.tf.cross(x[0], x[1]),
                          name="diff_cross_product")([gt_normals, opt_normals])
    diff_normals_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(
        diff_normals
    )  # This is added to avoid influencing embedding layer parameters by a "bad" gradient network
    diff_angles = Lambda(lambda x: K.tf.subtract(x[0], x[1]),
                         name="diff_angle")([gt_normals, opt_normals])
    diff_angles_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(diff_angles)
    diff_angles_norm_NOGRAD = Lambda(
        lambda x: K.tf.norm(x, axis=-1),
        name="diff_angle_norm")(diff_angles_NOGRAD)
    dist_angles = Lambda(lambda x: K.mean(K.square(x), axis=-1),
                         name="diff_angle_mse")(diff_angles)
    dist_angles_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(dist_angles)
    print("diff_angles shape: " + str(diff_angles.shape))
    print("dist_angles shape: " + str(dist_angles.shape))
    #pc_euclidean_diff_NOGRAD =  Lambda(lambda x: K.stop_gradient(x))(pc_euclidean_diff) # This is added to avoid influencing embedding layer parameters by a "bad" gradient network
    #print("diff_normals_NOGRAD shape: " + str(diff_normals_NOGRAD.shape))
    diff_normals_NOGRAD = Flatten()(diff_normals_NOGRAD)
    diff_angles_NOGRAD = Flatten()(diff_angles_NOGRAD)
    mesh_diff_NOGRAD = Concatenate()([diff_normals_NOGRAD, dist_angles_NOGRAD])

    optlearner_architecture = Dense(2**11,
                                    activation="relu")(vertex_diff_NOGRAD)
    #optlearner_architecture = Dense(2**11, activation="relu")(diff_angles_norm_NOGRAD)
    #optlearner_architecture = Dense(2**11, activation="relu")(diff_angles_NOGRAD)
    #optlearner_architecture = Dense(2**11, activation="relu")(mesh_diff_NOGRAD)
    optlearner_architecture = BatchNormalization()(optlearner_architecture)
    optlearner_architecture = Dropout(0.5)(optlearner_architecture)
    #optlearner_architecture = Dropout(0.2)(optlearner_architecture)
    print('optlearner_architecture shape: ' +
          str(optlearner_architecture.shape))
    #delta_d_hat = Dense(85, activation=pos_scaled_tanh, name="delta_d_hat")(optlearner_architecture)
    delta_d_hat = Dense(85, activation="linear",
                        name="delta_d_hat")(optlearner_architecture)
    #delta_d_hat = Dense(85, activation=centred_linear, name="delta_d_hat")(optlearner_architecture)
    print('delta_d_hat shape: ' + str(delta_d_hat.shape))
    #exit(1)

    # Calculate the (batched) MSE between the learned and ground truth offset in the parameters
    delta_d_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(delta_d)
    false_loss_delta_d_hat = Lambda(
        lambda x: K.mean(K.square(x[0] - x[1]), axis=1))(
            [delta_d_NOGRAD, delta_d_hat])
    #false_loss_delta_d_hat = Lambda(lambda x: K.sum(K.square(x[0] - x[1]), axis=1))([delta_d_NOGRAD, delta_d_hat])
    #false_loss_delta_d_hat = Lambda(lambda x: mape(x[0], x[1]))([delta_d_NOGRAD, delta_d_hat])
    false_loss_delta_d_hat = Reshape(
        target_shape=(1, ), name="delta_d_hat_mse")(false_loss_delta_d_hat)
    print("delta_d_hat loss shape: " + str(false_loss_delta_d_hat.shape))
    false_sin_loss_delta_d_hat = get_sin_metric(delta_d_NOGRAD, delta_d_hat)
    false_sin_loss_delta_d_hat = Lambda(
        lambda x: x, name="delta_d_hat_sin_output")(false_sin_loss_delta_d_hat)
    print("delta_d_hat sin loss shape: " +
          str(false_sin_loss_delta_d_hat.shape))

    # Prevent model from using the delta_d_hat gradient in final loss
    delta_d_hat_NOGRAD = Lambda(lambda x: K.stop_gradient(x),
                                name='optlearner_output_NOGRAD')(delta_d_hat)

    # False loss designed to pass the learned offset as a gradient to the embedding layer
    false_loss_smpl = Multiply(name="smpl_diff")(
        [optlearner_params, delta_d_hat_NOGRAD])
    print("smpl loss shape: " + str(false_loss_smpl.shape))

    #return [optlearner_input, gt_params, gt_pc], [optlearner_params, false_loss_delta_d, optlearner_pc, false_loss_pc, false_loss_delta_d_hat, false_sin_loss_delta_d_hat,  false_loss_smpl, delta_d, delta_d_hat, delta_d_hat_NOGRAD]
    return [optlearner_input, gt_params, gt_pc], [
        optlearner_params, false_loss_delta_d, optlearner_pc, false_loss_pc,
        false_loss_delta_d_hat, false_sin_loss_delta_d_hat, false_loss_smpl,
        delta_d, delta_d_hat, dist_angles
    ]
def OptLearnerMeshNormalStaticArchitecture(param_trainable,
                                           init_wrapper,
                                           emb_size=1000):
    """ Optimised learner network architecture """
    # An embedding layer is required to optimise the parameters
    #print('parameter initializer pose '+str(parameter_initializer([1000,85])[0,:15]))
    #print('parameter initializer shape '+str(parameter_initializer([1000,85])[0,72:82]))
    #print('parameter initializer T '+str(parameter_initializer([1000,85])[0,82:85]))
    #exit(1)

    optlearner_input = Input(shape=(1, ), name="embedding_index")

    def init_emb_layers(index, param_trainable, init_wrapper):
        """ Initialise the parameter embedding layers """
        emb_layers = []
        num_params = 85
        for i in range(num_params):
            layer_name = "param_{:02d}".format(i)
            initialiser = init_wrapper(param=i,
                                       offset=param_trainable[layer_name])

            emb_layer = Embedding(emb_size,
                                  1,
                                  name=layer_name,
                                  trainable=param_trainable[layer_name],
                                  embeddings_initializer=initialiser)(index)
            emb_layers.append(emb_layer)

        return emb_layers

    # Initialise the embedding layers
    emb_layers = init_emb_layers(optlearner_input, param_trainable,
                                 init_wrapper)
    optlearner_params = Concatenate(name="parameter_embedding")(emb_layers)

    #optlearner_params = Embedding(1000, 85, embeddings_initializer=parameter_initializer, name="parameter_embedding")(optlearner_input)
    #print("optlearner parameters shape: " +str(optlearner_params.shape))
    optlearner_params = Reshape(target_shape=(85, ),
                                name="learned_params")(optlearner_params)
    print("optlearner parameters shape: " + str(optlearner_params.shape))
    #exit(1)

    # Ground truth parameters and point cloud are inputs to the model as well
    gt_params = Input(shape=(85, ), name="gt_params")
    gt_pc = Input(shape=(6890, 3), name="gt_pc")
    print("gt parameters shape: " + str(gt_params.shape))
    print("gt point cloud shape: " + str(gt_pc.shape))

    # Compute the true offset (i.e. difference) between the ground truth and learned parameters
    #pi = K.constant(np.pi)
    #delta_d = Lambda(lambda x: x[0] - x[1])([gt_params, optlearner_params])
    #delta_d = Lambda(lambda x: K.tf.math.floormod(x - pi, 2*pi) - pi, name="delta_d")(delta_d)  # custom modulo 2pi of delta_d
    #print("delta_d shape: " + str(delta_d.shape))
    delta_d = Lambda(lambda x: x[0] - x[1],
                     name="delta_d")([gt_params, optlearner_params])
    print("delta_d shape: " + str(delta_d.shape))
    #exit(1)
    # Calculate the (batched) MSE between the learned parameters and the ground truth parameters

    false_loss_delta_d = Lambda(lambda x: K.mean(K.square(x), axis=1))(delta_d)
    print("delta_d loss shape: " + str(false_loss_delta_d.shape))
    #exit(1)
    false_loss_delta_d = Reshape(target_shape=(1, ),
                                 name="delta_d_mse")(false_loss_delta_d)
    print("delta_d loss shape: " + str(false_loss_delta_d.shape))

    # Load SMPL model and get necessary parameters
    smpl_params = load_params(
        './keras_rotationnet_v2_demo_for_hidde/basicModel_f_lbs_10_207_0_v1.0.0.pkl'
    )
    _, _, input_info = get_parameters()
    faces = smpl_params['f']  # canonical mesh faces
    print("faces shape: " + str(faces.shape))
    #exit(1)

    input_betas = Lambda(lambda x: x[:, 72:82])(optlearner_params)
    input_pose_rodrigues = Lambda(lambda x: x[:, 0:72])(optlearner_params)
    input_trans = Lambda(lambda x: x[:, 82:85])(optlearner_params)

    # Get the point cloud corresponding to these parameters
    optlearner_pc = Points3DFromSMPLParams(input_betas, input_pose_rodrigues,
                                           input_trans, smpl_params,
                                           input_info)
    print("optlearner point cloud shape: " + str(optlearner_pc.shape))
    #optlearner_pc = Lambda(lambda x: x * 0.0)[optlearner_pc]

    # Get the (batched) Euclidean loss between the learned and ground truth point clouds
    pc_euclidean_diff = Lambda(lambda x: x[0] - x[1])([gt_pc, optlearner_pc])
    pc_euclidean_dist = Lambda(lambda x: K.sum(K.square(x), axis=-1))(
        pc_euclidean_diff)
    print('pc euclidean dist ' + str(pc_euclidean_dist.shape))
    #exit(1)
    false_loss_pc = Lambda(lambda x: K.mean(x, axis=1))(pc_euclidean_dist)
    false_loss_pc = Reshape(target_shape=(1, ),
                            name="pc_mean_euc_dist")(false_loss_pc)
    print("point cloud loss shape: " + str(false_loss_pc.shape))
    #exit(1)

    #    # Get the (batched) MSE between the learned and ground truth point clouds
    #    false_loss_pc = Lambda(lambda x: tf.reduce_mean(tf.square(tf.subtract(x[0], x[1])), axis=[1,2]))([gt_pc, optlearner_pc])
    #    false_loss_pc = Reshape(target_shape=(1,), name="pointcloud_mse")(false_loss_pc)
    #    print("point cloud loss shape: " + str(false_loss_pc.shape))

    # Gather sets of points and compute their cross product to get mesh normals
    vertex_list = [1850, 1600, 2050, 5350, 5050, 5500]
    face_array = np.array(
        [face for face in faces for vertex in vertex_list if vertex in face])
    #gt_p0 = Lambda(lambda x: K.tf.gather(x, np.array(faces[:,0]).astype(np.int32), axis=-2))(gt_pc)
    #gt_p1 = Lambda(lambda x: K.tf.gather(x, np.array(faces[:,1]).astype(np.int32), axis=-2))(gt_pc)
    #gt_p2 = Lambda(lambda x: K.tf.gather(x, np.array(faces[:,2]).astype(np.int32), axis=-2))(gt_pc)
    #print("gt_p0 shape: " + str(gt_p0.shape))
    #print("gt_p1 shape: " + str(gt_p1.shape))
    #print("gt_p2 shape: " + str(gt_p2.shape))
    #gt_vec1 = Lambda(lambda x: x[1] - x[0])([gt_p0, gt_p1])
    #gt_vec2 = Lambda(lambda x: x[1] - x[0])([gt_p0, gt_p2])
    #print("gt_vec1 shape: " + str(gt_vec1.shape))
    #print("gt_vec2 shape: " + str(gt_vec2.shape))
    #gt_normals = Lambda(lambda x: K.l2_normalize(K.tf.cross(x[0], x[1]), axis=-1), name="gt_cross_product")([gt_vec1, gt_vec2])
    #gt_normals = get_mesh_normals(gt_pc, faces, layer_name="gt_cross_product")
    gt_normals = get_mesh_normals(gt_pc,
                                  face_array,
                                  layer_name="gt_cross_product")
    print("gt_normals shape: " + str(gt_normals.shape))

    #opt_p0 = Lambda(lambda x: K.tf.gather(x, np.array(faces[:,0]).astype(np.int32), axis=-2))(optlearner_pc)
    #opt_p1 = Lambda(lambda x: K.tf.gather(x, np.array(faces[:,1]).astype(np.int32), axis=-2))(optlearner_pc)
    #opt_p2 = Lambda(lambda x: K.tf.gather(x, np.array(faces[:,2]).astype(np.int32), axis=-2))(optlearner_pc)
    #print("opt_p0 shape: " + str(opt_p0.shape))
    #print("opt_p1 shape: " + str(opt_p1.shape))
    #print("opt_p2 shape: " + str(opt_p2.shape))
    #opt_vec1 = Lambda(lambda x: x[1] - x[0])([opt_p0, opt_p1])
    #opt_vec2 = Lambda(lambda x: x[1] - x[0])([opt_p0, opt_p2])
    #print("opt_vec1 shape: " + str(opt_vec1.shape))
    #print("opt_vec2 shape: " + str(opt_vec2.shape))
    #opt_normals = Lambda(lambda x: K.l2_normalize(K.tf.cross(x[0], x[1]), axis=-1), name="opt_cross_product")([opt_vec1, opt_vec2])
    #opt_normals = get_mesh_normals(optlearner_pc, faces, layer_name="opt_cross_product")
    opt_normals = get_mesh_normals(optlearner_pc,
                                   face_array,
                                   layer_name="opt_cross_product")
    print("opt_normals shape: " + str(opt_normals.shape))
    #exit(1)

    # Learn the offset in parameters from the difference between the ground truth and learned mesh normals
    diff_normals = Lambda(lambda x: K.tf.cross(x[0], x[1]),
                          name="diff_cross_product")([gt_normals, opt_normals])
    #pc_euclidean_diff_NOGRAD =  Lambda(lambda x: K.stop_gradient(x))(pc_euclidean_diff) # This is added to avoid influencing embedding layer parameters by a "bad" gradient network
    diff_normals_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(
        diff_normals
    )  # This is added to avoid influencing embedding layer parameters by a "bad" gradient network
    print("diff_normals_NOGRAD shape: " + str(diff_normals_NOGRAD.shape))

    # Get the mesh normal angle magnitudes (for evaluation)
    gt_angles = Lambda(lambda x: K.tf.norm(x, axis=-1),
                       name="gt_angles")(gt_normals)
    print("gt_angles shape: " + str(gt_angles.shape))
    opt_angles = Lambda(lambda x: K.tf.norm(x, axis=-1),
                        name="opt_angles")(opt_normals)
    print("opt_angles shape: " + str(opt_angles.shape))
    diff_angles = Lambda(lambda x: K.tf.norm(x, axis=-1),
                         name="diff_angles")(diff_normals_NOGRAD)
    print("diff_angles shape: " + str(diff_angles.shape))
    #exit(1)

    # Keep every 5xth normal entry
    #indices = np.array([i for i in enumerate()])
    #diff_normals_NOGRAD = Lambda(lambda x: x[:, ::10], name="reduce_num_normals")(diff_normals_NOGRAD)
    diff_normals_NOGRAD = Flatten()(diff_normals_NOGRAD)

    optlearner_architecture = Dense(2**11,
                                    activation="relu")(diff_normals_NOGRAD)
    #optlearner_architecture = Dense(2**12, activation="relu")(diff_normals_NOGRAD)
    optlearner_architecture = BatchNormalization()(optlearner_architecture)
    optlearner_architecture = Dropout(0.5)(optlearner_architecture)
    #optlearner_architecture = Dense(2**10, activation="relu")(optlearner_architecture)
    #optlearner_architecture = BatchNormalization()(optlearner_architecture)
    print('optlearner_architecture shape: ' +
          str(optlearner_architecture.shape))
    #exit(1)
    #optlearner_architecture = Dropout(0.1)(optlearner_architecture)
    #optlearner_architecture = Dense(1024, activation="relu")(optlearner_architecture)
    #optlearner_architecture = Dropout(0.1)(optlearner_architecture)
    delta_d_hat = Dense(85, activation="linear",
                        name="delta_d_hat")(optlearner_architecture)
    print('delta_d_hat shape: ' + str(delta_d_hat.shape))
    #exit(1)

    # Calculate the (batched) MSE between the learned and ground truth offset in the parameters
    #false_loss_delta_d_hat = Lambda(lambda x: K.mean(K.square(x[0] - x[1]), axis=1))([delta_d, delta_d_hat])
    delta_d_NOGRAD = Lambda(lambda x: K.stop_gradient(x),
                            name="delta_d_NOGRAD")(delta_d)
    #false_loss_delta_d_hat = Lambda(lambda x: K.sum(K.square(x[0] - x[1]), axis=1))([delta_d_NOGRAD, delta_d_hat])
    false_loss_delta_d_hat = Lambda(
        lambda x: K.mean(K.square(x[0] - x[1]), axis=1))(
            [delta_d_NOGRAD, delta_d_hat])
    false_loss_delta_d_hat = Reshape(
        target_shape=(1, ), name="delta_d_hat_mse")(false_loss_delta_d_hat)
    print("delta_d_hat loss shape: " + str(false_loss_delta_d_hat.shape))

    # Prevent model from using the delta_d_hat gradient in final loss
    delta_d_hat_NOGRAD = Lambda(lambda x: K.stop_gradient(x),
                                name='optlearner_output_NOGRAD')(delta_d_hat)

    # False loss designed to pass the learned offset as a gradient to the embedding layer
    false_loss_smpl = Lambda(lambda x: x[1] * x[0], name="smpl_diff")(
        [optlearner_params, delta_d_hat_NOGRAD])
    print("smpl loss shape: " + str(false_loss_smpl.shape))

    #return [optlearner_input, gt_params, gt_pc], [optlearner_params, false_loss_delta_d, optlearner_pc, false_loss_pc, false_loss_delta_d_hat, false_loss_smpl, delta_d, delta_d_hat,delta_d_hat_NOGRAD]
    return [optlearner_input, gt_params, gt_pc], [
        optlearner_params, false_loss_delta_d, optlearner_pc, false_loss_pc,
        false_loss_delta_d_hat, false_loss_smpl, delta_d, delta_d_hat,
        delta_d_hat_NOGRAD, gt_angles, opt_angles, diff_angles
    ]
示例#6
0
def GroupedConv1DOptLearnerArchitecture(param_trainable,
                                        init_wrapper,
                                        smpl_params,
                                        input_info,
                                        faces,
                                        emb_size=1000,
                                        input_type="3D_POINTS",
                                        groups=[]):
    """ Optimised learner network architecture """
    # An embedding layer is required to optimise the parameters
    optlearner_input = Input(shape=(1, ), name="embedding_index")

    # Initialise the embedding layers
    emb_layers = init_emb_layers(optlearner_input, emb_size, param_trainable,
                                 init_wrapper)
    optlearner_params = Concatenate(name="parameter_embedding")(emb_layers)
    optlearner_params = Reshape(target_shape=(85, ),
                                name="learned_params")(optlearner_params)
    print("optlearner parameters shape: " + str(optlearner_params.shape))
    #exit(1)

    # Ground truth parameters and point cloud are inputs to the model as well
    gt_params = Input(shape=(85, ), name="gt_params")
    gt_pc = Input(shape=(6890, 3), name="gt_pc")
    print("gt parameters shape: " + str(gt_params.shape))
    print("gt point cloud shape: " + str(gt_pc.shape))

    # Get trainable parameters
    trainable_params = sorted([
        int(param.replace("param_", ""))
        for param, trainable in param_trainable.items() if trainable
    ])
    print("trainable_params: " + str(trainable_params))
    non_trainable_params = [
        param for param in range(85) if param not in trainable_params
    ]
    print("non_trainable_params: " + str(non_trainable_params))
    joints_with_trainable = sorted(
        np.unique(
            [int((param - (param % 3)) / 3) for param in trainable_params]))
    print("joints_with_trainable: " + str(joints_with_trainable))
    params_from_joints = np.concatenate([[3 * i, 3 * i + 1, 3 * i + 2]
                                         for i in joints_with_trainable])
    print("params_from_joints: " + str(params_from_joints))

    #group1 = [0,1,2,3]
    #group2 = [4,5,6,7]
    #group3 = [8,9,10,11]
    #group4 = [12,13,14,15]
    #group5 = [16,17,18,19]
    #group6 = [20,21,22,23]

    #groups = [group1, group2, group3, group4, group5, group6]
    #groups = [group1 + group2, group3 + group4, group5 + group6]
    #groups = [[i] for i in range(24)]
    print("groups: " + str(groups))
    flattened_groups = [i for sublist in groups for i in sublist]
    print(flattened_groups)
    assert np.all([i in flattened_groups for i in range(24)])

    # Compute the true offset (i.e. difference) between the ground truth and learned parameters
    pi = K.constant(np.pi)
    delta_d = Lambda(lambda x: x[0] - x[1],
                     name="delta_d")([gt_params, optlearner_params])
    #delta_d = Lambda(lambda x: x[0] - x[1], name="delta_d_no_mod")([gt_params, optlearner_params])
    #delta_d = Lambda(lambda x: K.tf.math.floormod(x - pi, 2*pi) - pi, name="delta_d")(delta_d)  # custom modulo 2pi of delta_d
    #delta_d = custom_mod(delta_d, pi, name="delta_d")  # custom modulo 2pi of delta_d
    print("delta_d shape: " + str(delta_d.shape))
    #exit(1)

    # Calculate the (batched) MSE between the learned parameters and the ground truth parameters
    false_loss_delta_d = Lambda(lambda x: K.mean(K.square(x), axis=1))(delta_d)
    print("delta_d loss shape: " + str(false_loss_delta_d.shape))
    #exit(1)
    false_loss_delta_d = Reshape(target_shape=(1, ),
                                 name="delta_d_mse")(false_loss_delta_d)
    print("delta_d loss shape: " + str(false_loss_delta_d.shape))

    # Load SMPL model and get necessary parameters
    optlearner_pc = get_pc(optlearner_params, smpl_params, input_info,
                           faces)  # UNCOMMENT
    print("optlearner_pc shape: " + str(optlearner_pc.shape))
    #exit(1)
    #optlearner_pc = Dense(6890*3)(delta_d)
    #optlearner_pc = Reshape((6890, 3))(optlearner_pc)

    # Get the (batched) Euclidean loss between the learned and ground truth point clouds
    pc_euclidean_diff = Lambda(lambda x: x[0] - x[1])([gt_pc, optlearner_pc])
    pc_euclidean_dist = Lambda(lambda x: K.sum(K.square(x), axis=-1))(
        pc_euclidean_diff)
    print('pc euclidean dist ' + str(pc_euclidean_dist.shape))
    #exit(1)
    false_loss_pc = Lambda(lambda x: K.mean(x, axis=1))(pc_euclidean_dist)
    false_loss_pc = Reshape(target_shape=(1, ),
                            name="pc_mean_euc_dist")(false_loss_pc)
    print("point cloud loss shape: " + str(false_loss_pc.shape))
    #exit(1)

    # Gather sets of points and compute their cross product to get mesh normals
    # In order of: right hand, right wrist, right forearm, right bicep end, right bicep, right shoulder, top of cranium, left shoulder, left bicep, left bicep end, left forearm, left wrist, left hand,
    # chest, belly/belly button, back of neck, upper back, central back, lower back/tailbone,
    # left foot, left over-ankle, left shin, left over-knee, left quadricep, left hip, right, hip, right, quadricep, right over-knee, right shin, right, over-ankle, right foot
    vertex_list = [
        5674, 5705, 5039, 5151, 4977, 4198, 411, 606, 1506, 1682, 1571, 2244,
        2212, 3074, 3500, 460, 2878, 3014, 3021, 3365, 4606, 4588, 4671, 6877,
        1799, 5262, 3479, 1187, 1102, 1120, 6740
    ]
    # with added vertices in the feet
    #vertex_list = [5674, 5705, 5039, 5151, 4977, 4198, 411, 606, 1506, 1682, 1571, 2244, 2212,
    #        3074, 3500, 460, 2878, 3014, 3021,
    #        3365, 4606, 4588, 4671, 6877, 1799, 5262, 3479, 1187, 1102, 1120, 6740, 3392, 3545, 3438, 6838, 6781, 6792]
    #face_array = np.array([11396, 8620, 7866, 5431, 6460, 1732, 4507])
    pc_euclidean_diff_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(
        pc_euclidean_diff
    )  # This is added to avoid influencing embedding layer parameters by a "bad" gradient network
    vertex_diff_NOGRAD = Lambda(lambda x: K.tf.gather(
        x, np.array(vertex_list).astype(np.int32), axis=-2))(
            pc_euclidean_diff_NOGRAD)
    print("vertex_diff_NOGRAD shape: " + str(vertex_diff_NOGRAD.shape))
    vertex_diff_NOGRAD = Flatten()(vertex_diff_NOGRAD)
    #exit(1)
    face_array = np.array([[face for face in faces if vertex in face][0]
                           for vertex in vertex_list
                           ])  # only take a single face for each vertex
    print("face_array shape: " + str(face_array.shape))
    gt_normals = get_mesh_normals(gt_pc,
                                  face_array,
                                  layer_name="gt_cross_product")
    print("gt_normals shape: " + str(gt_normals.shape))
    opt_normals = get_mesh_normals(optlearner_pc,
                                   face_array,
                                   layer_name="opt_cross_product")
    print("opt_normals shape: " + str(opt_normals.shape))
    #exit(1)

    # Learn the offset in parameters from the difference between the ground truth and learned mesh normals
    diff_normals = Lambda(lambda x: K.tf.cross(x[0], x[1]),
                          name="diff_cross_product")([gt_normals, opt_normals])
    diff_normals_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(
        diff_normals
    )  # This is added to avoid influencing embedding layer parameters by a "bad" gradient network
    diff_angles = Lambda(lambda x: K.tf.subtract(x[0], x[1]),
                         name="diff_angle")([gt_normals, opt_normals])
    diff_angles_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(diff_angles)
    diff_angles_norm_NOGRAD = Lambda(
        lambda x: K.tf.norm(x, axis=-1),
        name="diff_angle_norm")(diff_angles_NOGRAD)
    dist_angles = Lambda(lambda x: K.mean(K.square(x), axis=-1),
                         name="diff_angle_mse")(diff_angles)
    #dist_angles = Lambda(lambda x: K.mean(K.abs(x), axis=-1), name="diff_angle_mse")(diff_angles)
    dist_angles_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(dist_angles)
    print("diff_angles shape: " + str(diff_angles.shape))
    print("dist_angles shape: " + str(dist_angles.shape))
    #pc_euclidean_diff_NOGRAD =  Lambda(lambda x: K.stop_gradient(x))(pc_euclidean_diff) # This is added to avoid influencing embedding layer parameters by a "bad" gradient network
    #print("diff_normals_NOGRAD shape: " + str(diff_normals_NOGRAD.shape))
    diff_normals_NOGRAD = Flatten()(diff_normals_NOGRAD)
    diff_angles_NOGRAD = Flatten()(diff_angles_NOGRAD)
    mesh_diff_NOGRAD = Concatenate()([diff_normals_NOGRAD, dist_angles_NOGRAD])

    if input_type == "3D_POINTS":
        deep_opt_input = Dense(2**9, activation="relu")(vertex_diff_NOGRAD)
    if input_type == "MESH_NORMALS":
        deep_opt_input = Dense(2**9, activation="relu")(mesh_diff_NOGRAD)
    print('deep_opt_input shape: ' + str(deep_opt_input.shape))
    deep_opt_input = Reshape((-1, 1))(deep_opt_input)
    print('deep_opt_input shape: ' + str(deep_opt_input.shape))
    #DROPOUT = 0.1
    DROPOUT = 0.0
    indices_ordering = []
    group_outputs = []
    group_losses = []
    group_sin_losses = []
    group_param_mse = []
    for group in groups:
        optlearner_architecture = Conv1D(64, 5, strides=2,
                                         activation="relu")(deep_opt_input)
        optlearner_architecture = Dropout(DROPOUT)(optlearner_architecture)
        optlearner_architecture = Conv1D(
            128, 5, strides=2, activation="relu")(optlearner_architecture)
        optlearner_architecture = Dropout(DROPOUT)(optlearner_architecture)
        optlearner_architecture = Conv1D(
            256, 3, strides=2, activation="relu")(optlearner_architecture)
        optlearner_architecture = Dropout(DROPOUT)(optlearner_architecture)
        #        optlearner_architecture = Conv1D(512, 3, strides=2, activation="relu")(optlearner_architecture)
        #        optlearner_architecture = Dropout(DROPOUT)(optlearner_architecture)
        #print('optlearner_architecture shape: '+str(optlearner_architecture.shape))
        #optlearner_architecture = Flatten()(optlearner_architecture)
        print('optlearner_architecture shape: ' +
              str(optlearner_architecture.shape))
        optlearner_architecture = Reshape((-1, ))(optlearner_architecture)
        print('optlearner_architecture shape: ' +
              str(optlearner_architecture.shape))
        #optlearner_architecture = Dropout(0.5)(optlearner_architecture)
        #optlearner_architecture = Dense(2**7, activation="relu")(optlearner_architecture)
        #print('optlearner_architecture shape: '+str(optlearner_architecture.shape))
        #delta_d_hat = Dense(85, activation="linear", name="delta_d_hat")(optlearner_architecture)
        delta_d_hat = Dense(3 * len(group),
                            activation="linear")(optlearner_architecture)
        print('delta_d_hat shape: ' + str(delta_d_hat.shape))
        group_outputs.append(delta_d_hat)
        #exit(1)

        indices = []
        for joint in group:
            j_base = 3 * joint
            j1 = j_base
            j2 = j_base + 1
            j3 = j_base + 2
            indices += [j1, j2, j3]
        indices_ordering += indices

        #indices = K.constant(indices)
        ## Filter parameters such that the model is only evaluated on trainable parameters
        #delta_d_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(delta_d)
        #delta_d_NOGRAD_FILTERED = Lambda(lambda x: K.tf.gather(x, indices, axis=-1))(delta_d_NOGRAD)
        #delta_d_hat_FILTERED = Lambda(lambda x: K.tf.gather(x, indices, axis=-1))(delta_d_hat)

    # predict shape and translation parameters
    optlearner_architecture = Dense(2**9, activation="relu")(deep_opt_input)
    deep_opt_input = Reshape((-1, 1))(deep_opt_input)
    print('deep_opt_input shape: ' + str(deep_opt_input.shape))
    optlearner_architecture = Conv1D(
        64, 5, strides=2, activation="relu")(optlearner_architecture)
    optlearner_architecture = Dropout(DROPOUT)(optlearner_architecture)
    optlearner_architecture = Conv1D(
        128, 5, strides=2, activation="relu")(optlearner_architecture)
    optlearner_architecture = Dropout(DROPOUT)(optlearner_architecture)
    optlearner_architecture = Conv1D(
        256, 3, strides=2, activation="relu")(optlearner_architecture)
    optlearner_architecture = Dropout(DROPOUT)(optlearner_architecture)
    optlearner_architecture = Conv1D(
        512, 3, strides=2, activation="relu")(optlearner_architecture)
    optlearner_architecture = Dropout(DROPOUT)(optlearner_architecture)
    #print('optlearner_architecture shape: '+str(optlearner_architecture.shape))
    #optlearner_architecture = Flatten()(optlearner_architecture)
    print('optlearner_architecture shape: ' +
          str(optlearner_architecture.shape))
    optlearner_architecture = Reshape((-1, ))(optlearner_architecture)
    print('optlearner_architecture shape: ' +
          str(optlearner_architecture.shape))
    #optlearner_architecture = Dropout(0.5)(optlearner_architecture)
    #optlearner_architecture = Dense(2**7, activation="relu")(optlearner_architecture)
    #print('optlearner_architecture shape: '+str(optlearner_architecture.shape))
    #delta_d_hat = Dense(85, activation="linear", name="delta_d_hat")(optlearner_architecture)
    shape_params = Dense(13, activation="linear",
                         name="shape_params")(optlearner_architecture)
    print('delta_d_hat shape: ' + str(delta_d_hat.shape))
    group_outputs.append(shape_params)
    #exit(1)

    indices_ordering += [i for i in range(72, 85)]
    print(indices_ordering)

    # process indices ordering to re-order array
    reordered_indices = []
    for i in sorted(indices_ordering):
        reordered_indices.append(indices_ordering.index(i))
    reordered_indices = K.constant(reordered_indices, dtype=K.tf.int32)
    #print(reordered_indices)
    #exit(1)

    delta_d_hat = Concatenate(axis=-1)(group_outputs)
    print("delta_d_hat shape: " + str(delta_d_hat.shape))
    delta_d_hat = Lambda(lambda x: K.tf.gather(x, reordered_indices, axis=-1),
                         name="delta_d_hat_collected")(delta_d_hat)
    print("delta_d_hat shape: " + str(delta_d_hat.shape))

    # Stop gradients for non-trainable parameters
    delta_d_hat_non_trainable = Lambda(
        lambda x: K.tf.gather(x, non_trainable_params, axis=-1))(delta_d_hat)
    print("delta_d_hat_non_trainable shape: " +
          str(delta_d_hat_non_trainable.shape))
    delta_d_hat_non_trainable_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(
        delta_d_hat_non_trainable)
    delta_d_hat_trainable = Lambda(
        lambda x: K.tf.gather(x, trainable_params, axis=-1))(delta_d_hat)
    print("delta_d_hat_trainable shape: " + str(delta_d_hat_trainable.shape))
    delta_d_hat_unordered = Concatenate()(
        [delta_d_hat_non_trainable_NOGRAD, delta_d_hat_trainable])
    print("delta_d_hat_unordered shape: " + str(delta_d_hat_unordered.shape))

    # Re-order indices
    params_ordering = non_trainable_params + trainable_params
    delta_d_hat_ordering = []
    for i in sorted(params_ordering):
        delta_d_hat_ordering.append(params_ordering.index(i))
    delta_d_hat_ordering = K.constant(delta_d_hat_ordering, dtype=K.tf.int32)

    delta_d_hat = Lambda(
        lambda x: K.tf.gather(x, delta_d_hat_ordering, axis=-1),
        name="delta_d_hat")(delta_d_hat_unordered)
    print("delta_d_hat shape: " + str(delta_d_hat.shape))

    # Get trainable true differences
    delta_d_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(delta_d)
    delta_d_NOGRAD_trainable = Lambda(
        lambda x: K.tf.gather(x, trainable_params, axis=-1))(delta_d_NOGRAD)

    # Rotation loss for trainable parameters
    delta_d_hat_vec, _, _ = split_and_reshape_euler_angles(delta_d_hat)
    delta_d_NOGRAD_vec, _, _ = split_and_reshape_euler_angles(delta_d_NOGRAD)
    delta_d_hat_SO3 = rot3d_from_rodrigues(delta_d_hat_vec)
    delta_d_NOGRAD_SO3 = rot3d_from_rodrigues(delta_d_NOGRAD_vec)
    rotational_loss = rotation_loss(delta_d_NOGRAD_SO3, delta_d_hat_SO3)
    print("rotational_loss shape: " + str(rotational_loss.shape))
    rotational_loss = Lambda(lambda x: K.tf.gather(
        x, joints_with_trainable, axis=-1))(rotational_loss)
    print("rotational_loss shape: " + str(rotational_loss.shape))
    rotational_loss = Lambda(lambda x: K.mean(x, axis=-1, keepdims=True),
                             name="rotational_loss")(rotational_loss)
    print("rotational_loss shape: " + str(rotational_loss.shape))

    #false_loss_delta_d_hat = Lambda(lambda x: K.mean(K.square(x[0] - x[1]), axis=1))([delta_d_NOGRAD, delta_d_hat])
    false_loss_delta_d_hat = Lambda(
        lambda x: K.mean(K.square(x[0] - x[1]), axis=1))(
            [delta_d_NOGRAD_trainable, delta_d_hat_trainable])
    #false_loss_delta_d_hat = geodesic_loss(delta_d_NOGRAD_SO3, delta_d_hat_SO3)
    #false_loss_delta_d_hat = Lambda(lambda x: x)(rotational_loss)
    #print("false_loss_delta_d_hat shape: " + str(false_loss_delta_d_hat.shape))
    #false_loss_delta_d_hat = Add()([false_loss_delta_d_hat, rotational_loss])
    print("false_loss_delta_d_hat shape: " + str(false_loss_delta_d_hat.shape))
    false_loss_delta_d_hat = Reshape(
        target_shape=(1, ), name="delta_d_hat_mse")(false_loss_delta_d_hat)
    print("delta_d_hat loss shape: " + str(false_loss_delta_d_hat.shape))

    # Metrics
    false_sin_loss_delta_d_hat = get_angular_distance_metric(
        delta_d_NOGRAD, delta_d_hat)
    #false_sin_loss_delta_d_hat = get_sin_metric(delta_d_NOGRAD, delta_d_hat)
    #false_sin_loss_delta_d_hat = get_sin_metric(delta_d_NOGRAD, delta_d_hat, average=False)
    false_sin_loss_delta_d_hat = Lambda(
        lambda x: x, name="delta_d_hat_sin_output")(false_sin_loss_delta_d_hat)
    print("delta_d_hat sin loss shape: " +
          str(false_sin_loss_delta_d_hat.shape))
    per_param_mse = Lambda(lambda x: K.square(K.sin(x[0] - x[1])))(
        [delta_d_NOGRAD, delta_d_hat])
    #per_param_mse = Lambda(lambda x: K.square(x[0] - x[1]))([delta_d_NOGRAD, delta_d_hat])
    per_param_mse = Reshape((85, ), name="params_mse")(per_param_mse)

    # Prevent model from using the delta_d_hat gradient in final loss
    delta_d_hat_NOGRAD = Lambda(lambda x: K.stop_gradient(x),
                                name='optlearner_output_NOGRAD')(delta_d_hat)

    # False loss designed to pass the learned offset as a gradient to the embedding layer
    false_loss_smpl = Multiply(name="smpl_diff")(
        [optlearner_params, delta_d_hat_NOGRAD])
    print("smpl loss shape: " + str(false_loss_smpl.shape))

    #return [optlearner_input, gt_params, gt_pc], [optlearner_params, false_loss_delta_d, optlearner_pc, false_loss_pc, false_loss_delta_d_hat, false_sin_loss_delta_d_hat,  false_loss_smpl, delta_d, delta_d_hat, dist_angles, per_param_mse]
    return [optlearner_input, gt_params, gt_pc], [
        optlearner_params, false_loss_delta_d, optlearner_pc, false_loss_pc,
        false_loss_delta_d_hat, false_sin_loss_delta_d_hat, false_loss_smpl,
        delta_d, delta_d_hat, dist_angles, per_param_mse, gt_normals,
        opt_normals
    ]
def RotConv1DOptLearnerArchitecture(param_trainable,
                                    init_wrapper,
                                    smpl_params,
                                    input_info,
                                    faces,
                                    emb_size=1000,
                                    input_type="3D_POINTS"):
    """ Optimised learner network architecture """
    # An embedding layer is required to optimise the parameters
    optlearner_input = Input(shape=(1, ), name="embedding_index")

    # Initialise the embedding layers
    emb_layers = init_emb_layers(optlearner_input, emb_size, param_trainable,
                                 init_wrapper)
    optlearner_params = Concatenate(name="parameter_embedding")(emb_layers)
    optlearner_params = Reshape(target_shape=(85, ),
                                name="learned_params")(optlearner_params)
    print("optlearner parameters shape: " + str(optlearner_params.shape))
    #exit(1)

    # Ground truth parameters and point cloud are inputs to the model as well
    gt_params = Input(shape=(85, ), name="gt_params")
    gt_pc = Input(shape=(6890, 3), name="gt_pc")
    print("gt parameters shape: " + str(gt_params.shape))
    print("gt point cloud shape: " + str(gt_pc.shape))

    # Compute the true offset (i.e. difference) between the ground truth and learned parameters
    pi = K.constant(np.pi)
    delta_d = Lambda(lambda x: x[0] - x[1],
                     name="delta_d")([gt_params, optlearner_params])
    #delta_d = Lambda(lambda x: x[0] - x[1], name="delta_d_no_mod")([gt_params, optlearner_params])
    #delta_d = Lambda(lambda x: K.tf.math.floormod(x - pi, 2*pi) - pi, name="delta_d")(delta_d)  # custom modulo 2pi of delta_d
    #delta_d = custom_mod(delta_d, pi, name="delta_d")  # custom modulo 2pi of delta_d
    print("delta_d shape: " + str(delta_d.shape))
    #exit(1)

    # Calculate the (batched) MSE between the learned parameters and the ground truth parameters
    false_loss_delta_d = Lambda(lambda x: K.mean(K.square(x), axis=1))(delta_d)
    print("delta_d loss shape: " + str(false_loss_delta_d.shape))
    #exit(1)
    false_loss_delta_d = Reshape(target_shape=(1, ),
                                 name="delta_d_mse")(false_loss_delta_d)
    print("delta_d loss shape: " + str(false_loss_delta_d.shape))

    # Load SMPL model and get necessary parameters
    optlearner_pc = get_pc(optlearner_params, smpl_params, input_info,
                           faces)  # UNCOMMENT
    print("optlearner_pc shape: " + str(optlearner_pc.shape))
    #exit(1)
    #optlearner_pc = Dense(6890*3)(delta_d)
    #optlearner_pc = Reshape((6890, 3))(optlearner_pc)

    # Get the (batched) Euclidean loss between the learned and ground truth point clouds
    pc_euclidean_diff = Lambda(lambda x: x[0] - x[1])([gt_pc, optlearner_pc])
    pc_euclidean_dist = Lambda(lambda x: K.sum(K.square(x), axis=-1))(
        pc_euclidean_diff)
    print('pc euclidean dist ' + str(pc_euclidean_dist.shape))
    #exit(1)
    false_loss_pc = Lambda(lambda x: K.mean(x, axis=1))(pc_euclidean_dist)
    false_loss_pc = Reshape(target_shape=(1, ),
                            name="pc_mean_euc_dist")(false_loss_pc)
    print("point cloud loss shape: " + str(false_loss_pc.shape))
    #exit(1)

    # Gather sets of points and compute their cross product to get mesh normals
    # In order of: right hand, right wrist, right forearm, right bicep end, right bicep, right shoulder, top of cranium, left shoulder, left bicep, left bicep end, left forearm, left wrist, left hand,
    # chest, belly/belly button, back of neck, upper back, central back, lower back/tailbone,
    # left foot, left over-ankle, left shin, left over-knee, left quadricep, left hip, right, hip, right, quadricep, right over-knee, right shin, right, over-ankle, right foot
    vertex_list = [
        5674, 5705, 5039, 5151, 4977, 4198, 411, 606, 1506, 1682, 1571, 2244,
        2212, 3074, 3500, 460, 2878, 3014, 3021, 3365, 4606, 4588, 4671, 6877,
        1799, 5262, 3479, 1187, 1102, 1120, 6740
    ]
    #vertex_list = [5674, 5512, 5474, 5705, 6335, 5438, 5039, 5047, 5074, 5151, 5147, 5188, 4977, 4870, 4744, 4198, 4195, 5324, 411, 164, 3676, 606, 1826, 1863, 1506, 1382, 1389, 1682, 1675, 1714, 1571, 1579, 1603, 2244, 1923, 1936, 2212, 2007, 2171, 3074, 539, 4081, 3500, 6875, 3477, 460, 426, 3921, 2878, 2969, 6428, 3014, 892, 4380, 3021, 1188, 4675, 3365, 3336, 3338, 1120, 1136, 1158, 1102, 1110, 1114, 1187, 1144, 977, 3479, 872, 1161, 1799, 3150, 1804, 5262, 6567, 5268, 6877, 4359, 4647, 4671, 4630, 4462, 4588, 4641, 4600, 4606, 4622, 4644, 6740, 6744, 6737]
    #face_array = np.array([11396, 8620, 7866, 5431, 6460, 1732, 4507])
    pc_euclidean_diff_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(
        pc_euclidean_diff
    )  # This is added to avoid influencing embedding layer parameters by a "bad" gradient network
    vertex_diff_NOGRAD = Lambda(lambda x: K.tf.gather(
        x, np.array(vertex_list).astype(np.int32), axis=-2))(
            pc_euclidean_diff_NOGRAD)
    print("vertex_diff_NOGRAD shape: " + str(vertex_diff_NOGRAD.shape))
    vertex_diff_NOGRAD = Flatten()(vertex_diff_NOGRAD)
    #exit(1)
    face_array = np.array([[face for face in faces if vertex in face][0]
                           for vertex in vertex_list
                           ])  # only take a single face for each vertex
    print("face_array shape: " + str(face_array.shape))
    gt_normals = get_mesh_normals(gt_pc,
                                  face_array,
                                  layer_name="gt_cross_product")
    print("gt_normals shape: " + str(gt_normals.shape))
    opt_normals = get_mesh_normals(optlearner_pc,
                                   face_array,
                                   layer_name="opt_cross_product")
    print("opt_normals shape: " + str(opt_normals.shape))
    #exit(1)

    # Learn the offset in parameters from the difference between the ground truth and learned mesh normals
    diff_normals = Lambda(lambda x: K.tf.cross(x[0], x[1]),
                          name="diff_cross_product")([gt_normals, opt_normals])
    diff_normals_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(
        diff_normals
    )  # This is added to avoid influencing embedding layer parameters by a "bad" gradient network
    diff_angles = Lambda(lambda x: K.tf.subtract(x[0], x[1]),
                         name="diff_angle")([gt_normals, opt_normals])
    diff_angles_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(diff_angles)
    diff_angles_norm_NOGRAD = Lambda(
        lambda x: K.tf.norm(x, axis=-1),
        name="diff_angle_norm")(diff_angles_NOGRAD)
    dist_angles = Lambda(lambda x: K.mean(K.square(x), axis=-1),
                         name="diff_angle_mse")(diff_angles)
    dist_angles_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(dist_angles)
    print("diff_angles shape: " + str(diff_angles.shape))
    print("dist_angles shape: " + str(dist_angles.shape))
    #pc_euclidean_diff_NOGRAD =  Lambda(lambda x: K.stop_gradient(x))(pc_euclidean_diff) # This is added to avoid influencing embedding layer parameters by a "bad" gradient network
    #print("diff_normals_NOGRAD shape: " + str(diff_normals_NOGRAD.shape))
    diff_normals_NOGRAD = Flatten()(diff_normals_NOGRAD)
    diff_angles_NOGRAD = Flatten()(diff_angles_NOGRAD)
    mesh_diff_NOGRAD = Concatenate()([diff_normals_NOGRAD, dist_angles_NOGRAD])

    if input_type == "3D_POINTS":
        optlearner_architecture = Dense(2**9,
                                        activation="relu")(vertex_diff_NOGRAD)
    if input_type == "MESH_NORMALS":
        #optlearner_architecture = Dense(2**11, activation="relu")(diff_angles_norm_NOGRAD)
        #optlearner_architecture = Dense(2**11, activation="relu")(diff_angles_NOGRAD)
        optlearner_architecture = Dense(2**9,
                                        activation="relu")(mesh_diff_NOGRAD)
    #optlearner_architecture = BatchNormalization()(optlearner_architecture)
    #optlearner_architecture = Dropout(0.5)(optlearner_architecture)
    print('optlearner_architecture shape: ' +
          str(optlearner_architecture.shape))
    optlearner_architecture = Reshape(
        (optlearner_architecture.shape[1].value, 1))(optlearner_architecture)
    print('optlearner_architecture shape: ' +
          str(optlearner_architecture.shape))
    optlearner_architecture = Conv1D(
        64, 5, strides=2, activation="relu")(optlearner_architecture)
    optlearner_architecture = Conv1D(
        128, 5, strides=2, activation="relu")(optlearner_architecture)
    optlearner_architecture = Conv1D(
        256, 3, strides=2, activation="relu")(optlearner_architecture)
    optlearner_architecture = Conv1D(
        512, 3, strides=2, activation="relu")(optlearner_architecture)
    #print('optlearner_architecture shape: '+str(optlearner_architecture.shape))
    #optlearner_architecture = Flatten()(optlearner_architecture)
    print('optlearner_architecture shape: ' +
          str(optlearner_architecture.shape))
    optlearner_architecture = Reshape((-1, ))(optlearner_architecture)
    print('optlearner_architecture shape: ' +
          str(optlearner_architecture.shape))

    # Learn a 6D representation of the parameters
    mapped_pose = Dense(24 * 6, activation="linear",
                        name="mapped_pose")(optlearner_architecture)
    shape_params = Dense(10, activation="linear",
                         name="shape_params")(optlearner_architecture)
    mapped_trans = Dense(6, activation="linear",
                         name="mapped_trans")(optlearner_architecture)

    # Reshape mapped predictions into vectors
    mapped_pose_vec = Reshape((24, 3, 2), name="mapped_pose_mat")(mapped_pose)
    mapped_trans_vec = Reshape((1, 3, 2),
                               name="mapped_trans_mat")(mapped_trans)
    print("mapped_pose shape: " + str(mapped_pose_vec.shape))
    print("mapped_trans shape: " + str(mapped_trans_vec.shape))

    # Convert 6D representation to rotation matrix in SO(3)
    rot3d_pose = rot3d_from_ortho6d(mapped_pose_vec)
    rot3d_trans = rot3d_from_ortho6d(mapped_trans_vec)
    rot3d_pose = Lambda(lambda x: x, name="rot3d_pose")(rot3d_pose)
    rot3d_trans = Lambda(lambda x: x, name="rot3d_trans")(rot3d_trans)
    print("rot3d_pose shape: " + str(rot3d_pose.shape))
    print("rot3d_trans shape: " + str(rot3d_trans.shape))

    # Cast GT difference SO(3) representation to 6D for loss calculation
    delta_d_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(delta_d)
    delta_d_pose = Lambda(lambda x: x[:, 0:72])(delta_d_NOGRAD)
    delta_d_shape = Lambda(lambda x: x[:, 72:82])(delta_d_NOGRAD)
    delta_d_trans = Lambda(lambda x: x[:, 82:85])(delta_d_NOGRAD)
    print("delta_d_pose shape: " + str(delta_d_pose.shape))
    print("delta_d_shape shape: " + str(delta_d_shape.shape))
    print("delta_d_trans shape: " + str(delta_d_trans.shape))

    delta_d_pose_vec = Reshape((24, 3), name="delta_d_pose_vec")(delta_d_pose)
    delta_d_trans_vec = Reshape((1, 3),
                                name="delta_d_trans_vec")(delta_d_trans)
    print("delta_d_pose_vec shape: " + str(delta_d_pose_vec.shape))
    print("delta_d_trans_vec shape: " + str(delta_d_trans_vec.shape))

    #rot3d_delta_d_pose = rot3d_from_euler(delta_d_pose_vec)
    #rot3d_delta_d_trans = rot3d_from_euler(delta_d_trans_vec)
    rot3d_delta_d_pose = rot3d_from_rodrigues(delta_d_pose_vec)
    rot3d_delta_d_trans = rot3d_from_rodrigues(delta_d_trans_vec)
    rot3d_delta_d_pose = Lambda(lambda x: x,
                                name="rot3d_delta_d_pose")(rot3d_delta_d_pose)
    rot3d_delta_d_trans = Lambda(
        lambda x: x, name="rot3d_delta_d_trans")(rot3d_delta_d_trans)
    print("rot3d_delta_d_pose shape: " + str(rot3d_delta_d_pose.shape))
    print("rot3d_delta_d_trans shape: " + str(rot3d_delta_d_trans.shape))

    mapped_delta_d_pose_vec = ortho6d_from_rot3d(rot3d_delta_d_pose)
    mapped_delta_d_trans_vec = ortho6d_from_rot3d(rot3d_delta_d_trans)
    print("mapped_delta_d_pose_vec shape: " +
          str(mapped_delta_d_pose_vec.shape))
    print("mapped_delta_d_trans_vec shape: " +
          str(mapped_delta_d_trans_vec.shape))

    mapped_delta_d_pose = Reshape(
        (24 * 3 * 2, ), name="mapped_delta_d_pose")(mapped_delta_d_pose_vec)
    mapped_delta_d_trans = Reshape(
        (1 * 3 * 2, ), name="mapped_delta_d_trans")(mapped_delta_d_trans_vec)
    print("mapped_delta_d_pose shape: " + str(mapped_delta_d_pose.shape))
    print("mapped_delta_d_trans shape: " + str(mapped_delta_d_trans.shape))
    mapped_delta_d = Concatenate(name="mapped_delta_d")(
        [mapped_delta_d_pose, delta_d_shape, mapped_delta_d_trans])
    print("mapped_delta_d shape: " + str(mapped_delta_d.shape))
    #exit(1)

    # Calculate L2-norm on 6D orthogonal representation or geodesic loss on SO(3) representation
    mapped_delta_d_hat = Concatenate(name="mapped_delta_d_hat")(
        [mapped_pose, shape_params, mapped_trans])
    print("mapped_delta_d_hat shape: " + str(mapped_delta_d_hat.shape))
    #exit(1)
    mapped_delta_d_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(
        mapped_delta_d)

    # L2-norm loss on 6D representation
    false_loss_delta_d_hat = Lambda(
        lambda x: K.mean(K.square(x[0] - x[1]), axis=-1))(
            [mapped_delta_d_NOGRAD, mapped_delta_d_hat])
    # 'Angular' loss between vectors in 6D representation
    mapped_pose_vec1 = Lambda(lambda x: x[:, :, :, 0])(mapped_pose_vec)
    mapped_pose_vec2 = Lambda(lambda x: x[:, :, :, 1])(mapped_pose_vec)
    mapped_delta_d_pose_vec1 = Lambda(lambda x: x[:, :, :, 0])(
        mapped_delta_d_pose_vec)
    mapped_delta_d_pose_vec2 = Lambda(lambda x: x[:, :, :, 1])(
        mapped_delta_d_pose_vec)
    mapped_trans_vec1 = Lambda(lambda x: x[:, :, :, 0])(mapped_trans_vec)
    mapped_trans_vec2 = Lambda(lambda x: x[:, :, :, 1])(mapped_trans_vec)
    mapped_delta_d_trans_vec1 = Lambda(lambda x: x[:, :, :, 0])(
        mapped_delta_d_trans_vec)
    mapped_delta_d_trans_vec2 = Lambda(lambda x: x[:, :, :, 1])(
        mapped_delta_d_trans_vec)
    pose_dot1 = angle_between_vectors(mapped_pose_vec1,
                                      mapped_delta_d_pose_vec1)
    pose_dot2 = angle_between_vectors(mapped_pose_vec2,
                                      mapped_delta_d_pose_vec2)
    trans_dot1 = angle_between_vectors(mapped_trans_vec1,
                                       mapped_delta_d_trans_vec1)
    trans_dot2 = angle_between_vectors(mapped_trans_vec2,
                                       mapped_delta_d_trans_vec2)
    #false_loss_delta_d_hat_shape = Lambda(lambda x: K.mean(K.square(x[0] - x[1]), axis=-1))([delta_d_shape, shape_params])
    #false_loss_delta_d_hat = Lambda(lambda x: K.mean(x[0] + x[1] + x[2] + x[3], axis=[-2, -1]) + x[-1])([pose_dot1, pose_dot2, trans_dot1, trans_dot2, false_loss_delta_d_hat_shape])
    # L2-norm loss on rotation matrix in SO(3)
    #false_loss_delta_d_hat_pose = Lambda(lambda x: K.mean(K.square(x[0] - x[1]), axis=[1, 2, 3]))([rot3d_delta_d_pose, rot3d_pose])
    #false_loss_delta_d_hat_shape = Lambda(lambda x: K.mean(K.square(x[0] - x[1]), axis=-1))([delta_d_shape, shape_params])
    #false_loss_delta_d_hat_trans = Lambda(lambda x: K.mean(K.square(x[0] - x[1]), axis=[1, 2, 3]))([rot3d_delta_d_trans, rot3d_trans])
    #false_loss_delta_d_hat = Lambda(lambda x: K.mean(K.tf.stack(x, axis=-1), axis=-1))([false_loss_delta_d_hat_pose, false_loss_delta_d_hat_shape, false_loss_delta_d_hat_trans])
    # Geodesic loss on rotation matrix in SO(3)
    #false_loss_delta_d_hat_pose = geodesic_loss(rot3d_delta_d_pose, rot3d_pose)
    #false_loss_delta_d_hat_shape = Lambda(lambda x: K.mean(K.square(x[0] - x[1]), axis=-1))([delta_d_shape, shape_params])
    #false_loss_delta_d_hat_trans = geodesic_loss(rot3d_delta_d_trans, rot3d_trans)
    #false_loss_delta_d_hat = Lambda(lambda x: K.mean(K.tf.stack(x, axis=-1), axis=-1))([false_loss_delta_d_hat_pose, false_loss_delta_d_hat_shape, false_loss_delta_d_hat_trans])

    # Calculate the loss for direction and magnitude separately
    #sign_loss = Lambda(lambda x: 0.5*(1. - softsign(10*x[0]*x[1])))([mapped_delta_d_NOGRAD, mapped_delta_d_hat])
    ##magnitude_loss = Lambda(lambda x: K.abs(K.abs(x[0]) - K.abs(x[1])))([delta_d_NOGRAD, delta_d_hat])
    #magnitude_loss = Lambda(lambda x: K.exp(K.abs(x[1]) - K.abs(x[0])) )([mapped_delta_d_NOGRAD, mapped_delta_d_hat])
    #weighting = 0.1
    #false_loss_delta_d_hat = Lambda(lambda x: K.mean(x[0] + weighting*x[1], axis=-1))([sign_loss, magnitude_loss])
    false_loss_delta_d_hat = Reshape(
        target_shape=(1, ), name="delta_d_hat_mse")(false_loss_delta_d_hat)
    print("delta_d_hat loss shape: " + str(false_loss_delta_d_hat.shape))
    #exit(1)

    #    # Apply rotation in Euler angles
    #    pos_euler_poses = euler_from_rot3d(rot3d_pose)
    #    pos_euler_trans = euler_from_rot3d(rot3d_trans)
    #    #pos_euler_poses = euler_from_rot3d(rot3d_delta_d_pose)     # DEBUG ONLY
    #    #pos_euler_trans = euler_from_rot3d(rot3d_delta_d_trans)    # DEBUG ONLY
    #    print("pos_euler_poses shape: " + str(pos_euler_poses.shape))
    #    print("pos_euler_trans shape: " + str(pos_euler_trans.shape))
    #
    #    # Only consider first possibility for simplicity - this needs to be taken into account when evaluating predictions
    #    euler_pose = Lambda(lambda x: x[:, :, :, 0], name="euler_pose")(pos_euler_poses)
    #    euler_trans = Lambda(lambda x: x[:, :, :, 0], name="euler_trans")(pos_euler_trans)
    #    #euler_pose = Lambda(lambda x: x[:, :, :, 1], name="euler_pose")(pos_euler_poses)
    #    #euler_trans = Lambda(lambda x: x[:, :, :, 1], name="euler_trans")(pos_euler_trans)
    #    print("euler_pose shape: " + str(euler_pose.shape))
    #    print("euler_trans shape: " + str(euler_trans.shape))
    #
    #    # Reshape
    #    #delta_d_hat_pose = Reshape((72,), name="delta_d_hat_pose_reshaped")(euler_pose)
    #    #delta_d_hat_trans = Reshape((3,), name="delta_d_hat_trans_reshaped")(euler_trans)

    # Apply rotation in Rodrigues angles
    rodrigues_pose = rodrigues_from_rot3d(rot3d_pose)
    rodrigues_trans = rodrigues_from_rot3d(rot3d_trans)
    print("rodrigues_pose shape: " + str(rodrigues_pose.shape))
    print("rodrigues_trans shape: " + str(rodrigues_trans.shape))
    delta_d_hat_pose = Reshape(
        (72, ), name="delta_d_hat_pose_reshaped")(rodrigues_pose)
    delta_d_hat_trans = Reshape(
        (3, ), name="delta_d_hat_trans_reshaped")(rodrigues_trans)
    print("delta_d_hat_pose shape: " + str(delta_d_hat_pose.shape))
    print("delta_d_hat_trans shape: " + str(delta_d_hat_trans.shape))
    #exit(1)

    # Concatenate to form final update vector
    delta_d_hat = Concatenate(name="delta_d_hat")(
        [delta_d_hat_pose, shape_params, delta_d_hat_trans])
    #delta_d_hat = Concatenate(name="delta_d_hat")([euler_pose, shape_params, euler_trans])
    print('delta_d_hat shape: ' + str(delta_d_hat.shape))
    #exit(1)

    false_sin_loss_delta_d_hat = get_angular_distance_metric(
        delta_d_NOGRAD, delta_d_hat)
    #false_sin_loss_delta_d_hat = get_sin_metric(delta_d_NOGRAD, delta_d_hat)
    #false_sin_loss_delta_d_hat = get_sin_metric(delta_d_NOGRAD, delta_d_hat, average=False)
    false_sin_loss_delta_d_hat = Lambda(
        lambda x: x, name="delta_d_hat_sin_output")(false_sin_loss_delta_d_hat)
    #false_sin_loss_delta_d_hat = Lambda(lambda x: x, name="delta_d_hat_sin_output")(false_loss_new_pc)
    print("delta_d_hat sin loss shape: " +
          str(false_sin_loss_delta_d_hat.shape))

    # Prevent model from using the delta_d_hat gradient in final loss
    delta_d_hat_NOGRAD = Lambda(lambda x: K.stop_gradient(x),
                                name='optlearner_output_NOGRAD')(delta_d_hat)

    # False loss designed to pass the learned offset as a gradient to the embedding layer
    false_loss_smpl = Multiply(name="smpl_diff")(
        [optlearner_params, delta_d_hat_NOGRAD])
    print("smpl loss shape: " + str(false_loss_smpl.shape))

    # FOR DEBUGGING ONLY!!!
    #mapped_delta_d_pose_vec = Reshape((24, 3, 2))(mapped_delta_d_pose)
    #reverse_mapped_delta_d_pose = rot3d_from_ortho6d(mapped_delta_d_pose_vec)
    #rot3d_pose = Lambda(lambda x: x, name="rot3d_pose")(reverse_mapped_delta_d_pose)
    #test_rot3d_pose = Lambda(lambda x: x, name="test_rot3d_pose")(reverse_mapped_delta_d_pose)
    #rodrigues_delta_d_pose = rodrigues_from_rot3d(test_rot3d_pose)
    #rodrigues_delta_d_pose = rodrigues_from_rot3d(rot3d_delta_d_pose)
    #rodrigues_delta_d_pose = Lambda(lambda x: x, name="rodrigues_delta_d_pose")(rodrigues_delta_d_pose)

    return [optlearner_input, gt_params, gt_pc], [
        optlearner_params, false_loss_delta_d, optlearner_pc, false_loss_pc,
        false_loss_delta_d_hat, false_sin_loss_delta_d_hat, false_loss_smpl,
        delta_d, delta_d_hat, dist_angles, rot3d_delta_d_pose, rot3d_pose,
        mapped_pose, mapped_delta_d_pose
    ]
示例#8
0
def ConditionalOptLearnerArchitecture(param_trainable,
                                      init_wrapper,
                                      smpl_params,
                                      input_info,
                                      faces,
                                      emb_size=1000,
                                      input_type="3D_POINTS"):
    """ Optimised learner network architecture """
    # An embedding layer is required to optimise the parameters
    optlearner_input = Input(shape=(1, ), name="embedding_index")

    # Initialise the embedding layers
    emb_layers = init_emb_layers(optlearner_input, emb_size, param_trainable,
                                 init_wrapper)
    optlearner_params = Concatenate(name="parameter_embedding")(emb_layers)
    optlearner_params = Reshape(target_shape=(85, ),
                                name="learned_params")(optlearner_params)
    print("optlearner parameters shape: " + str(optlearner_params.shape))
    #exit(1)

    # Ground truth parameters and point cloud are inputs to the model as well
    gt_params = Input(shape=(85, ), name="gt_params")
    gt_pc = Input(shape=(6890, 3), name="gt_pc")
    print("gt parameters shape: " + str(gt_params.shape))
    print("gt point cloud shape: " + str(gt_pc.shape))

    # Compute the true offset (i.e. difference) between the ground truth and learned parameters
    pi = K.constant(np.pi)
    delta_d = Lambda(lambda x: x[0] - x[1],
                     name="delta_d")([gt_params, optlearner_params])
    print("delta_d shape: " + str(delta_d.shape))
    #exit(1)

    # Calculate the (batched) MSE between the learned parameters and the ground truth parameters
    false_loss_delta_d = Lambda(lambda x: K.mean(K.square(x), axis=1))(delta_d)
    print("delta_d loss shape: " + str(false_loss_delta_d.shape))
    #exit(1)
    false_loss_delta_d = Reshape(target_shape=(1, ),
                                 name="delta_d_mse")(false_loss_delta_d)
    print("delta_d loss shape: " + str(false_loss_delta_d.shape))

    # Load SMPL model and get necessary parameters
    optlearner_pc = get_pc(optlearner_params, smpl_params, input_info,
                           faces)  # UNCOMMENT
    print("optlearner_pc shape: " + str(optlearner_pc.shape))
    #exit(1)

    # Get the (batched) Euclidean loss between the learned and ground truth point clouds
    pc_euclidean_diff = Lambda(lambda x: x[0] - x[1])([gt_pc, optlearner_pc])
    pc_euclidean_dist = Lambda(lambda x: K.sum(K.square(x), axis=-1))(
        pc_euclidean_diff)
    print('pc euclidean dist ' + str(pc_euclidean_dist.shape))
    #exit(1)
    false_loss_pc = Lambda(lambda x: K.mean(x, axis=1))(pc_euclidean_dist)
    false_loss_pc = Reshape(target_shape=(1, ),
                            name="pc_mean_euc_dist")(false_loss_pc)
    print("point cloud loss shape: " + str(false_loss_pc.shape))
    #exit(1)

    # Gather sets of points and compute their cross product to get mesh normals
    # In order of: right hand, right wrist, right forearm, right bicep end, right bicep, right shoulder, top of cranium, left shoulder, left bicep, left bicep end, left forearm, left wrist, left hand,
    # chest, belly/belly button, back of neck, upper back, central back, lower back/tailbone,
    # left foot, left over-ankle, left shin, left over-knee, left quadricep, left hip, right, hip, right, quadricep, right over-knee, right shin, right, over-ankle, right foot
    vertex_list = [
        5674, 5705, 5039, 5151, 4977, 4198, 411, 606, 1506, 1682, 1571, 2244,
        2212, 3074, 3500, 460, 2878, 3014, 3021, 3365, 4606, 4588, 4671, 6877,
        1799, 5262, 3479, 1187, 1102, 1120, 6740
    ]
    #vertex_list = [5674, 5512, 5474, 5705, 6335, 5438, 5039, 5047, 5074, 5151, 5147, 5188, 4977, 4870, 4744, 4198, 4195, 5324, 411, 164, 3676, 606, 1826, 1863, 1506, 1382, 1389, 1682, 1675, 1714, 1571, 1579, 1603, 2244, 1923, 1936, 2212, 2007, 2171, 3074, 539, 4081, 3500, 6875, 3477, 460, 426, 3921, 2878, 2969, 6428, 3014, 892, 4380, 3021, 1188, 4675, 3365, 3336, 3338, 1120, 1136, 1158, 1102, 1110, 1114, 1187, 1144, 977, 3479, 872, 1161, 1799, 3150, 1804, 5262, 6567, 5268, 6877, 4359, 4647, 4671, 4630, 4462, 4588, 4641, 4600, 4606, 4622, 4644, 6740, 6744, 6737]
    #face_array = np.array([11396, 8620, 7866, 5431, 6460, 1732, 4507])
    pc_euclidean_diff_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(
        pc_euclidean_diff
    )  # This is added to avoid influencing embedding layer parameters by a "bad" gradient network
    vertex_diff_NOGRAD = Lambda(lambda x: K.tf.gather(
        x, np.array(vertex_list).astype(np.int32), axis=-2))(
            pc_euclidean_diff_NOGRAD)
    print("vertex_diff_NOGRAD shape: " + str(vertex_diff_NOGRAD.shape))
    vertex_diff_NOGRAD = Flatten()(vertex_diff_NOGRAD)
    #exit(1)
    face_array = np.array([[face for face in faces if vertex in face][0]
                           for vertex in vertex_list
                           ])  # only take a single face for each vertex
    print("face_array shape: " + str(face_array.shape))
    gt_normals = get_mesh_normals(gt_pc,
                                  face_array,
                                  layer_name="gt_cross_product")
    print("gt_normals shape: " + str(gt_normals.shape))
    opt_normals = get_mesh_normals(optlearner_pc,
                                   face_array,
                                   layer_name="opt_cross_product")
    print("opt_normals shape: " + str(opt_normals.shape))
    #exit(1)

    # Learn the offset in parameters from the difference between the ground truth and learned mesh normals
    diff_normals = Lambda(lambda x: K.tf.cross(x[0], x[1]),
                          name="diff_cross_product")([gt_normals, opt_normals])
    diff_normals_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(
        diff_normals
    )  # This is added to avoid influencing embedding layer parameters by a "bad" gradient network
    diff_angles = Lambda(lambda x: K.tf.subtract(x[0], x[1]),
                         name="diff_angle")([gt_normals, opt_normals])
    diff_angles_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(diff_angles)
    diff_angles_norm_NOGRAD = Lambda(
        lambda x: K.tf.norm(x, axis=-1),
        name="diff_angle_norm")(diff_angles_NOGRAD)
    dist_angles = Lambda(lambda x: K.mean(K.square(x), axis=-1),
                         name="diff_angle_mse")(diff_angles)
    dist_angles_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(dist_angles)
    print("diff_angles shape: " + str(diff_angles.shape))
    print("dist_angles shape: " + str(dist_angles.shape))
    diff_normals_NOGRAD = Flatten()(diff_normals_NOGRAD)
    diff_angles_NOGRAD = Flatten()(diff_angles_NOGRAD)
    mesh_diff_NOGRAD = Concatenate()([diff_normals_NOGRAD, dist_angles_NOGRAD])

    if input_type == "3D_POINTS":
        optlearner_architecture = Dense(2**9,
                                        activation="relu")(vertex_diff_NOGRAD)
    if input_type == "MESH_NORMALS":
        optlearner_architecture = Dense(2**9,
                                        activation="relu")(mesh_diff_NOGRAD)
    print('optlearner_architecture shape: ' +
          str(optlearner_architecture.shape))
    optlearner_architecture = Reshape(
        (optlearner_architecture.shape[1].value, 1))(optlearner_architecture)
    print('optlearner_architecture shape: ' +
          str(optlearner_architecture.shape))
    optlearner_architecture = Conv1D(
        64, 5, strides=2, activation="relu")(optlearner_architecture)
    optlearner_architecture = Conv1D(
        128, 5, strides=2, activation="relu")(optlearner_architecture)
    optlearner_architecture = Conv1D(
        256, 3, strides=2, activation="relu")(optlearner_architecture)
    optlearner_architecture = Conv1D(
        512, 3, strides=2, activation="relu")(optlearner_architecture)
    print('optlearner_architecture shape: ' +
          str(optlearner_architecture.shape))
    optlearner_architecture = Reshape((-1, ))(optlearner_architecture)
    print('optlearner_architecture shape: ' +
          str(optlearner_architecture.shape))

    # Learn a 6D representation of the parameters
    trainable_params = [
        int(param[6:8]) for param, trainable in param_trainable.items()
        if trainable
    ]
    trainable_joints = 3 * np.unique(
        [param // 3 for param in trainable_params])
    trainable_params = np.ravel([[param, param + 1, param + 2]
                                 for param in trainable_joints])
    print("trainable_params in conditional network: " + str(trainable_params))
    num_trainable_joints = len(trainable_joints)
    num_trainable_params = len(trainable_params)
    mapped_pose = Dense(num_trainable_joints * 6,
                        activation="linear",
                        name="mapped_pose")(optlearner_architecture)

    # Reshape mapped predictions into vectors
    mapped_pose_vec = Reshape((num_trainable_joints, 3, 2),
                              name="mapped_pose_mat")(mapped_pose)
    print("mapped_pose shape: " + str(mapped_pose_vec.shape))

    # Convert 6D representation to rotation matrix in SO(3)
    rot3d_pose = rot3d_from_ortho6d(mapped_pose_vec)
    rot3d_pose = Lambda(lambda x: x, name="rot3d_pose")(rot3d_pose)
    print("rot3d_pose shape: " + str(rot3d_pose.shape))

    # Cast GT difference SO(3) representation to 6D for loss calculation
    delta_d_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(delta_d)
    delta_d_pose = Lambda(lambda x: K.tf.gather(
        x, np.array(trainable_params), axis=1))(delta_d_NOGRAD)
    print("delta_d_pose shape: " + str(delta_d_pose.shape))

    delta_d_pose_vec = Reshape((num_trainable_joints, 3),
                               name="delta_d_pose_vec")(delta_d_pose)
    print("delta_d_pose_vec shape: " + str(delta_d_pose_vec.shape))

    #rot3d_delta_d_pose = rot3d_from_euler(delta_d_pose_vec)
    #rot3d_delta_d_trans = rot3d_from_euler(delta_d_trans_vec)
    rot3d_delta_d_pose = rot3d_from_rodrigues(delta_d_pose_vec)
    rot3d_delta_d_pose = Lambda(lambda x: x,
                                name="rot3d_delta_d_pose")(rot3d_delta_d_pose)
    print("rot3d_delta_d_pose shape: " + str(rot3d_delta_d_pose.shape))

    mapped_delta_d_pose_vec = ortho6d_from_rot3d(rot3d_delta_d_pose)
    print("mapped_delta_d_pose_vec shape: " +
          str(mapped_delta_d_pose_vec.shape))

    mapped_delta_d_pose = Reshape(
        (num_trainable_joints * 3 * 2, ),
        name="mapped_delta_d_pose")(mapped_delta_d_pose_vec)
    print("mapped_delta_d_pose shape: " + str(mapped_delta_d_pose.shape))
    mapped_delta_d = Lambda(lambda x: x,
                            name="mapped_delta_d")(mapped_delta_d_pose)
    print("mapped_delta_d shape: " + str(mapped_delta_d.shape))
    #exit(1)

    # Calculate L2-norm on 6D orthogonal representation or geodesic loss on SO(3) representation
    mapped_delta_d_hat = Lambda(lambda x: x,
                                name="mapped_delta_d_hat")(mapped_pose)
    print("mapped_delta_d_hat shape: " + str(mapped_delta_d_hat.shape))
    #exit(1)
    mapped_delta_d_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(
        mapped_delta_d)

    # L2-norm loss on 6D representation
    false_loss_delta_d_hat = Lambda(
        lambda x: K.mean(K.square(x[0] - x[1]), axis=-1))(
            [mapped_delta_d_NOGRAD, mapped_delta_d_hat])
    false_loss_delta_d_hat = Reshape(
        (1, ), name="delta_d_hat_mse")(false_loss_delta_d_hat)

    # Apply rotation in Rodrigues angles
    rodrigues_pose = rodrigues_from_rot3d(rot3d_pose)
    print("rodrigues_pose shape: " + str(rodrigues_pose.shape))
    delta_d_hat_pose = Reshape(
        (num_trainable_params, ),
        name="delta_d_hat_pose_reshaped")(rodrigues_pose)
    print("delta_d_hat_pose shape: " + str(delta_d_hat_pose.shape))
    #exit(1)

    # Concatenate to form final update vector
    # indices = K.constant(trainable_params, shape=(1, num_trainable_params), dtype="int32")
    # print("indices shape: " + str(indices.shape))
    # indices = Lambda(lambda x: K.tf.placeholder_with_default(x, [None, num_trainable_params]))(indices)
    # print("indices shape: " + str(indices.shape))
    # zeros = Lambda(lambda x: K.zeros_like(x))(delta_d_NOGRAD)
    # scatter = Lambda(lambda x: tf.scatter_nd_add(x[0], x[1], x[2]))([zeros, indices, delta_d_hat_pose])
    # print("scatter shape: " +str(scatter.shape))
    # exit(1)
    delta_d_hat = Lambda(lambda x: x, name="delta_d_hat")(delta_d_hat_pose)
    print('delta_d_hat shape: ' + str(delta_d_hat.shape))
    #exit(1)

    #false_sin_loss_delta_d_hat = get_angular_distance_metric(delta_d_NOGRAD, delta_d_hat)
    false_sin_loss_delta_d_hat = get_sin_metric(delta_d_pose, delta_d_hat)
    #false_sin_loss_delta_d_hat = get_sin_metric(delta_d_NOGRAD, delta_d_hat, average=False)
    false_sin_loss_delta_d_hat = Lambda(
        lambda x: x, name="delta_d_hat_sin_output")(false_sin_loss_delta_d_hat)
    #false_sin_loss_delta_d_hat = Lambda(lambda x: x, name="delta_d_hat_sin_output")(false_loss_new_pc)
    print("delta_d_hat sin loss shape: " +
          str(false_sin_loss_delta_d_hat.shape))

    # Prevent model from using the delta_d_hat gradient in final loss
    delta_d_hat_NOGRAD = Lambda(lambda x: K.stop_gradient(x),
                                name='optlearner_output_NOGRAD')(delta_d_hat)

    # False loss designed to pass the learned offset as a gradient to the embedding layer
    #false_loss_smpl = Multiply(name="smpl_diff")([optlearner_params, delta_d_hat_NOGRAD])
    false_loss_smpl = Lambda(lambda x: x, name="smpl_diff")(optlearner_params)
    print("smpl loss shape: " + str(false_loss_smpl.shape))

    return [optlearner_input, gt_params, gt_pc], [
        optlearner_params, false_loss_delta_d, optlearner_pc, false_loss_pc,
        false_loss_delta_d_hat, false_sin_loss_delta_d_hat, false_loss_smpl,
        delta_d, delta_d_hat, dist_angles, rot3d_delta_d_pose, rot3d_pose,
        mapped_pose, mapped_delta_d_pose
    ]
示例#9
0
def PeriodicOptLearnerArchitecture(param_trainable,
                                   init_wrapper,
                                   smpl_params,
                                   input_info,
                                   faces,
                                   emb_size=1000,
                                   input_type="3D_POINTS",
                                   groups=[],
                                   update_weight=1.0):
    """ Optimised learner network architecture """
    # An embedding layer is required to optimise the parameters
    optlearner_input = Input(shape=(1, ), name="embedding_index")

    # Initialise the embedding layers
    emb_layers = init_emb_layers(optlearner_input, emb_size, param_trainable,
                                 init_wrapper)
    optlearner_params = Concatenate(name="parameter_embedding")(emb_layers)
    optlearner_params = Reshape(target_shape=(85, ),
                                name="learned_params")(optlearner_params)
    print("optlearner parameters shape: " + str(optlearner_params.shape))
    #exit(1)

    # Ground truth parameters and point cloud are inputs to the model as well
    gt_params = Input(shape=(85, ), name="gt_params")
    gt_pc = Input(shape=(6890, 3), name="gt_pc")
    print("gt parameters shape: " + str(gt_params.shape))
    print("gt point cloud shape: " + str(gt_pc.shape))

    #group1 = [0,1,2,3]
    #group2 = [4,5,6,7]
    #group3 = [8,9,10,11]
    #group4 = [12,13,14,15]
    #group5 = [16,17,18,19]
    #group6 = [20,21,22,23]

    #groups = [group1, group2, group3, group4, group5, group6]
    #groups = [group1 + group2, group3 + group4, group5 + group6]
    #groups = [[i] for i in range(24)]
    print("groups: " + str(groups))
    flattened_groups = [i for sublist in groups for i in sublist]
    print(flattened_groups)
    assert np.all([i in flattened_groups for i in range(24)])

    # Params to train this epoch
    params_to_train = Input(shape=(85, ), dtype="bool", name="params_to_train")
    print("params_to_train shape: " + str(params_to_train.shape))
    params_to_train_indices = Lambda(lambda x: x[0])(params_to_train)
    print("params_to_train_indices shape: " +
          str(params_to_train_indices.shape))
    params_to_train_indices = Lambda(
        lambda x: K.tf.where(K.tf.cast(x, "bool")))(params_to_train_indices)
    print("params_to_train_indices shape: " +
          str(params_to_train_indices.shape))
    params_to_train_indices = Lambda(lambda x: K.squeeze(x, 1))(
        params_to_train_indices)
    print("params_to_train_indices shape: " +
          str(params_to_train_indices.shape))
    #exit(1)

    #DROPOUT = 0.1
    DROPOUT = 0.0

    # Compute the true offset (i.e. difference) between the ground truth and learned parameters
    pi = K.constant(np.pi)
    delta_d = Lambda(lambda x: x[0] - x[1],
                     name="delta_d")([gt_params, optlearner_params])
    #delta_d = Lambda(lambda x: x[0] - x[1], name="delta_d_no_mod")([gt_params, optlearner_params])
    #delta_d = Lambda(lambda x: K.tf.math.floormod(x - pi, 2*pi) - pi, name="delta_d")(delta_d)  # custom modulo 2pi of delta_d
    #delta_d = custom_mod(delta_d, pi, name="delta_d")  # custom modulo 2pi of delta_d
    print("delta_d shape: " + str(delta_d.shape))
    delta_d_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(delta_d)
    #exit(1)

    # Calculate the (batched) MSE between the learned parameters and the ground truth parameters
    false_loss_delta_d = Lambda(lambda x: K.mean(K.square(x), axis=1))(delta_d)
    print("delta_d loss shape: " + str(false_loss_delta_d.shape))
    #exit(1)
    false_loss_delta_d = Reshape(target_shape=(1, ),
                                 name="delta_d_mse")(false_loss_delta_d)
    print("delta_d loss shape: " + str(false_loss_delta_d.shape))

    # Load SMPL model and get necessary parameters
    optlearner_pc = get_pc(optlearner_params, smpl_params, input_info,
                           faces)  # UNCOMMENT
    print("optlearner_pc shape: " + str(optlearner_pc.shape))
    #exit(1)
    #optlearner_pc = Dense(6890*3)(delta_d)
    #optlearner_pc = Reshape((6890, 3))(optlearner_pc)

    # Get the (batched) Euclidean loss between the learned and ground truth point clouds
    pc_euclidean_diff = Lambda(lambda x: x[0] - x[1])([gt_pc, optlearner_pc])
    pc_euclidean_dist = Lambda(lambda x: K.sum(K.square(x), axis=-1))(
        pc_euclidean_diff)
    print('pc euclidean dist ' + str(pc_euclidean_dist.shape))
    #exit(1)
    false_loss_pc = Lambda(lambda x: K.mean(x, axis=1))(pc_euclidean_dist)
    false_loss_pc = Reshape(target_shape=(1, ),
                            name="pc_mean_euc_dist")(false_loss_pc)
    print("point cloud loss shape: " + str(false_loss_pc.shape))
    #exit(1)

    # Gather sets of points and compute their cross product to get mesh normals
    # In order of: right hand, right wrist, right forearm, right bicep end, right bicep, right shoulder, top of cranium, left shoulder, left bicep, left bicep end, left forearm, left wrist, left hand,
    # chest, belly/belly button, back of neck, upper back, central back, lower back/tailbone,
    # left foot, left over-ankle, left shin, left over-knee, left quadricep, left hip, right, hip, right, quadricep, right over-knee, right shin, right, over-ankle, right foot
    vertex_list = [
        5674, 5705, 5039, 5151, 4977, 4198, 411, 606, 1506, 1682, 1571, 2244,
        2212, 3074, 3500, 460, 2878, 3014, 3021, 3365, 4606, 4588, 4671, 6877,
        1799, 5262, 3479, 1187, 1102, 1120, 6740
    ]
    # with added vertices in the feet
    #vertex_list = [5674, 5705, 5039, 5151, 4977, 4198, 411, 606, 1506, 1682, 1571, 2244, 2212,
    #        3074, 3500, 460, 2878, 3014, 3021,
    #        3365, 4606, 4588, 4671, 6877, 1799, 5262, 3479, 1187, 1102, 1120, 6740, 3392, 3545, 3438, 6838, 6781, 6792]
    #face_array = np.array([11396, 8620, 7866, 5431, 6460, 1732, 4507])
    pc_euclidean_diff_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(
        pc_euclidean_diff
    )  # This is added to avoid influencing embedding layer parameters by a "bad" gradient network
    vertex_diff_NOGRAD = Lambda(lambda x: K.tf.gather(
        x, np.array(vertex_list).astype(np.int32), axis=-2))(
            pc_euclidean_diff_NOGRAD)
    print("vertex_diff_NOGRAD shape: " + str(vertex_diff_NOGRAD.shape))
    vertex_diff_NOGRAD = Flatten()(vertex_diff_NOGRAD)
    #exit(1)
    face_array = np.array([[face for face in faces if vertex in face][0]
                           for vertex in vertex_list
                           ])  # only take a single face for each vertex
    print("face_array shape: " + str(face_array.shape))
    gt_normals = get_mesh_normals(gt_pc,
                                  face_array,
                                  layer_name="gt_cross_product")
    print("gt_normals shape: " + str(gt_normals.shape))
    opt_normals = get_mesh_normals(optlearner_pc,
                                   face_array,
                                   layer_name="opt_cross_product")
    print("opt_normals shape: " + str(opt_normals.shape))
    #exit(1)

    # Learn the offset in parameters from the difference between the ground truth and learned mesh normals
    diff_normals = Lambda(lambda x: K.tf.cross(x[0], x[1]),
                          name="diff_cross_product")([gt_normals, opt_normals])
    diff_normals_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(
        diff_normals
    )  # This is added to avoid influencing embedding layer parameters by a "bad" gradient network
    diff_angles = Lambda(lambda x: K.tf.subtract(x[0], x[1]),
                         name="diff_angle")([gt_normals, opt_normals])
    diff_angles_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(diff_angles)
    diff_angles_norm_NOGRAD = Lambda(
        lambda x: K.tf.norm(x, axis=-1),
        name="diff_angle_norm")(diff_angles_NOGRAD)
    dist_angles = Lambda(lambda x: K.mean(K.square(x), axis=-1),
                         name="diff_angle_mse")(diff_angles)
    dist_angles_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(dist_angles)
    print("diff_angles shape: " + str(diff_angles.shape))
    print("dist_angles shape: " + str(dist_angles.shape))
    diff_normals_NOGRAD = Flatten()(diff_normals_NOGRAD)
    diff_angles_NOGRAD = Flatten()(diff_angles_NOGRAD)
    mesh_diff_NOGRAD = Concatenate()([diff_normals_NOGRAD, dist_angles_NOGRAD])

    # Learn to predict in conditional manner
    indices_ordering = []
    group_outputs = []
    new_params = Lambda(lambda x: x)(optlearner_params)
    group_vertex_diff_NOGRAD = Lambda(lambda x: x)(vertex_diff_NOGRAD)
    group_mesh_diff_NOGRAD = Lambda(lambda x: x)(mesh_diff_NOGRAD)
    for group in groups:
        old_params = Lambda(lambda x: x)(new_params)

        if input_type == "3D_POINTS":
            deep_opt_input = Dense(2**9,
                                   activation="relu")(group_vertex_diff_NOGRAD)
            #deep_opt_input = Dense(2**7, activation="relu")(group_vertex_diff_NOGRAD)
        if input_type == "MESH_NORMALS":
            deep_opt_input = Dense(2**9,
                                   activation="relu")(group_mesh_diff_NOGRAD)
            #deep_opt_input = Dense(2**7, activation="relu")(group_mesh_diff_NOGRAD)
        if input_type == "ONLY_NORMALS":
            deep_opt_input = Dense(2**9,
                                   activation="relu")(group_mesh_diff_NOGRAD)
            #deep_opt_input = Dense(2**7, activation="relu")(group_diff_normals_NOGRAD)

        print('deep_opt_input shape: ' + str(deep_opt_input.shape))
        deep_opt_input = Reshape((-1, 1))(deep_opt_input)
        print('deep_opt_input shape: ' + str(deep_opt_input.shape))

        optlearner_architecture = Conv1D(64, 5, strides=2,
                                         activation="relu")(deep_opt_input)
        optlearner_architecture = Dropout(DROPOUT)(optlearner_architecture)
        optlearner_architecture = Conv1D(
            128, 5, strides=2, activation="relu")(optlearner_architecture)
        optlearner_architecture = Dropout(DROPOUT)(optlearner_architecture)
        optlearner_architecture = Conv1D(
            256, 3, strides=2, activation="relu")(optlearner_architecture)
        optlearner_architecture = Dropout(DROPOUT)(optlearner_architecture)
        #optlearner_architecture = Conv1D(512, 3, strides=2, activation="relu")(optlearner_architecture)
        #optlearner_architecture = Dropout(DROPOUT)(optlearner_architecture)
        print('optlearner_architecture shape: ' +
              str(optlearner_architecture.shape))
        optlearner_architecture = Reshape((-1, ))(optlearner_architecture)
        print('optlearner_architecture shape: ' +
              str(optlearner_architecture.shape))
        delta_d_hat = Dense(3 * len(group),
                            activation="linear")(optlearner_architecture)
        print('delta_d_hat shape: ' + str(delta_d_hat.shape))
        group_outputs.append(delta_d_hat)
        #exit(1)

        indices = []
        for joint in group:
            j_base = 3 * joint
            j1 = j_base
            j2 = j_base + 1
            j3 = j_base + 2
            indices += [j1, j2, j3]
        indices_ordering += indices

        # Get remaining indices
        rem_indices = [i for i in range(85) if i not in indices]
        group_indices_ord = indices + rem_indices
        group_ind_reord = []
        for i in sorted(group_indices_ord):
            group_ind_reord.append(group_indices_ord.index(i))

        # Apply prediction and re-render to prepare input for next group
        # Update parameters
        old_params_group = Lambda(lambda x: K.tf.gather(x, indices, axis=-1))(
            old_params)
        print("old_params_group shape: " + str(old_params_group.shape))
        new_params_group = Lambda(lambda x: x[0] + update_weight * x[1])(
            [old_params_group, delta_d_hat])
        print("new_params_group shape: " + str(new_params_group.shape))
        # Complete the parameter vector
        rem_params = Lambda(lambda x: K.tf.gather(x, rem_indices, axis=-1))(
            old_params)
        new_params = Concatenate()([new_params_group, rem_params])
        print("new_params shape: " + str(new_params.shape))
        # Re-order the parameter vector
        new_params = Lambda(
            lambda x: K.tf.gather(x, group_ind_reord, axis=-1))(new_params)
        print("new_params shape: " + str(new_params.shape))

        # Render new parameters
        new_pc = Lambda(lambda x: get_pc(x, smpl_params, input_info, faces))(
            new_params)
        print("new_pc shape: " + str(new_pc.shape))

        # Gather input vertices
        group_pc_euclidean_diff = Lambda(lambda x: x[0] - x[1])(
            [gt_pc, new_pc])
        group_pc_euclidean_diff_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(
            group_pc_euclidean_diff)
        group_vertex_diff_NOGRAD = Lambda(lambda x: K.tf.gather(
            x, np.array(vertex_list).astype(np.int32), axis=-2))(
                group_pc_euclidean_diff_NOGRAD)
        print("group_vertex_diff_NOGRAD shape: " +
              str(group_vertex_diff_NOGRAD.shape))
        group_vertex_diff_NOGRAD = Flatten()(group_vertex_diff_NOGRAD)
        #exit(1)
        group_gt_normals = Lambda(lambda x: get_mesh_normals(x, face_array))(
            gt_pc)
        print("group_gt_normals shape: " + str(group_gt_normals.shape))
        group_opt_normals = Lambda(lambda x: get_mesh_normals(x, face_array))(
            new_pc)
        print("group_opt_normals shape: " + str(group_opt_normals.shape))
        #exit(1)

        # Learn the offset in parameters from the difference between the ground truth and learned mesh normals
        group_diff_normals = Lambda(lambda x: K.tf.cross(x[0], x[1]))(
            [group_gt_normals, group_opt_normals])
        group_diff_normals_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(
            group_diff_normals)
        group_diff_angles = Lambda(lambda x: K.tf.subtract(x[0], x[1]))(
            [group_gt_normals, group_opt_normals])
        group_diff_angles_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(
            group_diff_angles)
        group_dist_angles = Lambda(lambda x: K.mean(K.square(x), axis=-1))(
            group_diff_angles)
        group_dist_angles_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(
            group_dist_angles)
        print("group_diff_angles shape: " + str(group_diff_angles.shape))
        print("group_dist_angles shape: " + str(group_dist_angles.shape))
        group_diff_normals_NOGRAD = Flatten()(group_diff_normals_NOGRAD)
        group_diff_angles_NOGRAD = Flatten()(group_diff_angles_NOGRAD)
        group_mesh_diff_NOGRAD = Concatenate()(
            [group_diff_normals_NOGRAD, group_dist_angles_NOGRAD])

    if input_type == "3D_POINTS":
        deep_opt_input = Dense(2**9, activation="relu")(vertex_diff_NOGRAD)
        #deep_opt_input = Dense(2**7, activation="relu")(vertex_diff_NOGRAD)
    if input_type == "MESH_NORMALS":
        deep_opt_input = Dense(2**9, activation="relu")(mesh_diff_NOGRAD)
        #deep_opt_input = Dense(2**7, activation="relu")(mesh_diff_NOGRAD)
    if input_type == "ONLY_NORMALS":
        deep_opt_input = Dense(2**9, activation="relu")(mesh_diff_NOGRAD)
        #deep_opt_input = Dense(2**7, activation="relu")(diff_normals_NOGRAD)

    # predict shape and translation parameters
    deep_opt_input = Reshape((-1, 1))(deep_opt_input)
    print('deep_opt_input shape: ' + str(deep_opt_input.shape))
    optlearner_architecture = Conv1D(64, 5, strides=2,
                                     activation="relu")(deep_opt_input)
    optlearner_architecture = Dropout(DROPOUT)(optlearner_architecture)
    optlearner_architecture = Conv1D(
        128, 5, strides=2, activation="relu")(optlearner_architecture)
    optlearner_architecture = Dropout(DROPOUT)(optlearner_architecture)
    optlearner_architecture = Conv1D(
        256, 3, strides=2, activation="relu")(optlearner_architecture)
    optlearner_architecture = Dropout(DROPOUT)(optlearner_architecture)
    optlearner_architecture = Conv1D(
        512, 3, strides=2, activation="relu")(optlearner_architecture)
    optlearner_architecture = Dropout(DROPOUT)(optlearner_architecture)
    #print('optlearner_architecture shape: '+str(optlearner_architecture.shape))
    #optlearner_architecture = Flatten()(optlearner_architecture)
    print('optlearner_architecture shape: ' +
          str(optlearner_architecture.shape))
    optlearner_architecture = Reshape((-1, ))(optlearner_architecture)
    print('optlearner_architecture shape: ' +
          str(optlearner_architecture.shape))
    #optlearner_architecture = Dropout(0.5)(optlearner_architecture)
    #optlearner_architecture = Dense(2**7, activation="relu")(optlearner_architecture)
    #print('optlearner_architecture shape: '+str(optlearner_architecture.shape))
    #delta_d_hat = Dense(85, activation="linear", name="delta_d_hat")(optlearner_architecture)
    shape_params = Dense(13, activation="linear",
                         name="shape_params")(optlearner_architecture)
    print('delta_d_hat shape: ' + str(delta_d_hat.shape))
    group_outputs.append(shape_params)
    #exit(1)

    indices_ordering += [i for i in range(72, 85)]
    print(indices_ordering)

    # process indices ordering to re-order array
    reordered_indices = []
    for i in sorted(indices_ordering):
        reordered_indices.append(indices_ordering.index(i))
    reordered_indices = K.constant(reordered_indices, dtype=K.tf.int32)
    #print(reordered_indices)
    #exit(1)

    delta_d_hat = Concatenate(axis=-1)(group_outputs)
    print("delta_d_hat shape: " + str(delta_d_hat.shape))
    delta_d_hat = Lambda(lambda x: K.tf.gather(x, reordered_indices, axis=-1),
                         name="delta_d_hat")(delta_d_hat)
    print("delta_d_hat shape: " + str(delta_d_hat.shape))

    #delta_d_hat_trainable = Lambda(lambda x: K.tf.boolean_mask(x[0], K.tf.cast(x[1], "bool")))([delta_d_hat, params_to_train])
    #delta_d_NOGRAD_trainable = Lambda(lambda x: K.tf.boolean_mask(x[0], K.tf.cast(x[1], "bool")))([delta_d_NOGRAD, params_to_train])
    delta_d_hat_trainable = Lambda(
        lambda x: K.tf.gather(x[0], K.tf.cast(x[1], "int32"), axis=-1))(
            [delta_d_hat, params_to_train_indices])
    print("delta_d_hat_trainable shape: " + str(delta_d_hat_trainable.shape))
    delta_d_NOGRAD_trainable = Lambda(
        lambda x: K.tf.gather(x[0], K.tf.cast(x[1], "int32"), axis=-1))(
            [delta_d_NOGRAD, params_to_train_indices])
    print("delta_d_NOGRAD_trainable shape: " +
          str(delta_d_NOGRAD_trainable.shape))

    #false_loss_delta_d_hat = Lambda(lambda x: K.mean(K.square(x[0] - x[1]), axis=1))([delta_d_NOGRAD, delta_d_hat])
    false_loss_delta_d_hat = Lambda(
        lambda x: K.mean(K.square(x[0] - x[1]), axis=1))(
            [delta_d_NOGRAD_trainable, delta_d_hat_trainable])
    print("delta_d_hat loss shape: " + str(false_loss_delta_d_hat.shape))
    false_loss_delta_d_hat = Reshape(
        target_shape=(1, ), name="delta_d_hat_mse")(false_loss_delta_d_hat)
    print("delta_d_hat loss shape: " + str(false_loss_delta_d_hat.shape))

    # Metrics
    false_sin_loss_delta_d_hat = get_angular_distance_metric(
        delta_d_NOGRAD, delta_d_hat)
    #false_sin_loss_delta_d_hat = get_sin_metric(delta_d_NOGRAD, delta_d_hat)
    #false_sin_loss_delta_d_hat = get_sin_metric(delta_d_NOGRAD, delta_d_hat, average=False)
    false_sin_loss_delta_d_hat = Lambda(
        lambda x: x, name="delta_d_hat_sin_output")(false_sin_loss_delta_d_hat)
    print("delta_d_hat sin loss shape: " +
          str(false_sin_loss_delta_d_hat.shape))
    per_param_mse = Lambda(lambda x: K.square(K.sin(x[0] - x[1])))(
        [delta_d_NOGRAD, delta_d_hat])
    #per_param_mse = Lambda(lambda x: K.square(x[0] - x[1]))([delta_d_NOGRAD, delta_d_hat])
    per_param_mse = Reshape((85, ), name="params_mse")(per_param_mse)

    # Prevent model from using the delta_d_hat gradient in final loss
    delta_d_hat_NOGRAD = Lambda(lambda x: K.stop_gradient(x),
                                name='optlearner_output_NOGRAD')(delta_d_hat)

    # False loss designed to pass the learned offset as a gradient to the embedding layer
    false_loss_smpl = Multiply(name="smpl_diff")(
        [optlearner_params, delta_d_hat_NOGRAD])
    print("smpl loss shape: " + str(false_loss_smpl.shape))

    return [optlearner_input, gt_params, gt_pc, params_to_train], [
        optlearner_params, false_loss_delta_d, optlearner_pc, false_loss_pc,
        false_loss_delta_d_hat, false_sin_loss_delta_d_hat, false_loss_smpl,
        delta_d, delta_d_hat, dist_angles, per_param_mse
    ]
示例#10
0
def ProbCNNOptLearnerStaticArchitecture(param_trainable,
                                        init_wrapper,
                                        smpl_params,
                                        input_info,
                                        faces,
                                        emb_size=1000,
                                        input_type="3D_POINTS"):
    """ Optimised learner network architecture that outputs samples from a distribution instead of deterministic estimates """
    # An embedding layer is required to optimise the parameters
    optlearner_input = Input(shape=(1, ), name="embedding_index")

    # Initialise the embedding layers
    emb_layers = init_emb_layers(optlearner_input, emb_size, param_trainable,
                                 init_wrapper)
    optlearner_params = Concatenate(name="parameter_embedding")(emb_layers)
    optlearner_params = Reshape(target_shape=(85, ),
                                name="learned_params")(optlearner_params)
    print("optlearner parameters shape: " + str(optlearner_params.shape))
    #exit(1)

    # Ground truth parameters and point cloud are inputs to the model as well
    gt_params = Input(shape=(85, ), name="gt_params")
    gt_pc = Input(shape=(6890, 3), name="gt_pc")
    print("gt parameters shape: " + str(gt_params.shape))
    print("gt point cloud shape: " + str(gt_pc.shape))

    # Compute the true offset (i.e. difference) between the ground truth and learned parameters
    pi = K.constant(np.pi)
    delta_d = Lambda(lambda x: x[0] - x[1],
                     name="delta_d")([gt_params, optlearner_params])
    #delta_d = Lambda(lambda x: x[0] - x[1], name="delta_d_no_mod")([gt_params, optlearner_params])
    #delta_d = Lambda(lambda x: K.tf.math.floormod(x - pi, 2*pi) - pi, name="delta_d")(delta_d)  # custom modulo 2pi of delta_d
    #delta_d = custom_mod(delta_d, pi, name="delta_d")  # custom modulo 2pi of delta_d
    print("delta_d shape: " + str(delta_d.shape))
    #exit(1)

    # Calculate the (batched) MSE between the learned parameters and the ground truth parameters
    false_loss_delta_d = Lambda(lambda x: K.mean(K.square(x), axis=1))(delta_d)
    print("delta_d loss shape: " + str(false_loss_delta_d.shape))
    #exit(1)
    false_loss_delta_d = Reshape(target_shape=(1, ),
                                 name="delta_d_mse")(false_loss_delta_d)
    print("delta_d loss shape: " + str(false_loss_delta_d.shape))

    # Load SMPL model and get necessary parameters
    optlearner_pc = get_pc(optlearner_params, smpl_params, input_info,
                           faces)  # UNCOMMENT
    print("optlearner_pc shape: " + str(optlearner_pc.shape))
    #exit(1)
    #optlearner_pc = Dense(6890*3)(delta_d)
    #optlearner_pc = Reshape((6890, 3))(optlearner_pc)

    # Get the (batched) Euclidean loss between the learned and ground truth point clouds
    pc_euclidean_diff = Lambda(lambda x: x[0] - x[1])([gt_pc, optlearner_pc])
    pc_euclidean_dist = Lambda(lambda x: K.sum(K.square(x), axis=-1))(
        pc_euclidean_diff)
    print('pc euclidean dist ' + str(pc_euclidean_dist.shape))
    #exit(1)
    false_loss_pc = Lambda(lambda x: K.mean(x, axis=1))(pc_euclidean_dist)
    false_loss_pc = Reshape(target_shape=(1, ),
                            name="pc_mean_euc_dist")(false_loss_pc)
    print("point cloud loss shape: " + str(false_loss_pc.shape))
    #exit(1)

    # Gather sets of points and compute their cross product to get mesh normals
    # In order of: right hand, right wrist, right forearm, right bicep end, right bicep, right shoulder, top of cranium, left shoulder, left bicep, left bicep end, left forearm, left wrist, left hand,
    # chest, belly/belly button, back of neck, upper back, central back, lower back/tailbone,
    # left foot, left over-ankle, left shin, left over-knee, left quadricep, left hip, right, hip, right, quadricep, right over-knee, right shin, right, over-ankle, right foot
    vertex_list = [
        5674, 5705, 5039, 5151, 4977, 4198, 411, 606, 1506, 1682, 1571, 2244,
        2212, 3074, 3500, 460, 2878, 3014, 3021, 3365, 4606, 4588, 4671, 6877,
        1799, 5262, 3479, 1187, 1102, 1120, 6740
    ]
    #face_array = np.array([11396, 8620, 7866, 5431, 6460, 1732, 4507])
    pc_euclidean_diff_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(
        pc_euclidean_diff
    )  # This is added to avoid influencing embedding layer parameters by a "bad" gradient network
    vertex_diff_NOGRAD = Lambda(lambda x: K.tf.gather(
        x, np.array(vertex_list).astype(np.int32), axis=-2))(
            pc_euclidean_diff_NOGRAD)
    print("vertex_diff_NOGRAD shape: " + str(vertex_diff_NOGRAD.shape))
    vertex_diff_NOGRAD = Flatten()(vertex_diff_NOGRAD)
    #exit(1)
    face_array = np.array([[face for face in faces if vertex in face][0]
                           for vertex in vertex_list
                           ])  # only take a single face for each vertex
    print("face_array shape: " + str(face_array.shape))
    gt_normals = get_mesh_normals(gt_pc,
                                  face_array,
                                  layer_name="gt_cross_product")
    print("gt_normals shape: " + str(gt_normals.shape))
    opt_normals = get_mesh_normals(optlearner_pc,
                                   face_array,
                                   layer_name="opt_cross_product")
    print("opt_normals shape: " + str(opt_normals.shape))
    #exit(1)

    # Learn the offset in parameters from the difference between the ground truth and learned mesh normals
    diff_normals = Lambda(lambda x: K.tf.cross(x[0], x[1]),
                          name="diff_cross_product")([gt_normals, opt_normals])
    diff_normals_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(
        diff_normals
    )  # This is added to avoid influencing embedding layer parameters by a "bad" gradient network
    diff_angles = Lambda(lambda x: K.tf.subtract(x[0], x[1]),
                         name="diff_angle")([gt_normals, opt_normals])
    diff_angles_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(diff_angles)
    diff_angles_norm_NOGRAD = Lambda(
        lambda x: K.tf.norm(x, axis=-1),
        name="diff_angle_norm")(diff_angles_NOGRAD)
    dist_angles = Lambda(lambda x: K.mean(K.square(x), axis=-1),
                         name="diff_angle_mse")(diff_angles)
    dist_angles_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(dist_angles)
    print("diff_angles shape: " + str(diff_angles.shape))
    print("dist_angles shape: " + str(dist_angles.shape))
    #pc_euclidean_diff_NOGRAD =  Lambda(lambda x: K.stop_gradient(x))(pc_euclidean_diff) # This is added to avoid influencing embedding layer parameters by a "bad" gradient network
    #print("diff_normals_NOGRAD shape: " + str(diff_normals_NOGRAD.shape))
    diff_normals_NOGRAD = Flatten()(diff_normals_NOGRAD)
    diff_angles_NOGRAD = Flatten()(diff_angles_NOGRAD)
    mesh_diff_NOGRAD = Concatenate()([diff_normals_NOGRAD, dist_angles_NOGRAD])

    if input_type == "3D_POINTS":
        optlearner_architecture = Dense(2**9,
                                        activation="relu")(vertex_diff_NOGRAD)
    if input_type == "MESH_NORMALS":
        optlearner_architecture = Dense(2**9,
                                        activation="relu")(mesh_diff_NOGRAD)
    print('optlearner_architecture shape: ' +
          str(optlearner_architecture.shape))
    optlearner_architecture = Reshape(
        (optlearner_architecture.shape[1].value, 1))(optlearner_architecture)
    print('optlearner_architecture shape: ' +
          str(optlearner_architecture.shape))
    optlearner_architecture = Conv1D(
        64, 5, activation="relu")(optlearner_architecture)
    optlearner_architecture = MaxPooling1D(3)(optlearner_architecture)
    print('optlearner_architecture shape: ' +
          str(optlearner_architecture.shape))
    optlearner_architecture = Conv1D(
        128, 3, activation="relu")(optlearner_architecture)
    optlearner_architecture = AveragePooling1D(2)(optlearner_architecture)
    print('optlearner_architecture shape: ' +
          str(optlearner_architecture.shape))
    print('optlearner_architecture shape: ' +
          str(optlearner_architecture.shape))
    optlearner_architecture = Flatten()(optlearner_architecture)
    print('optlearner_architecture shape: ' +
          str(optlearner_architecture.shape))
    optlearner_architecture = Dropout(0.5)(optlearner_architecture)

    # Learn the parameters of a univariate Gausssian distribution for each SMPL parameter
    mu_architecture = Dense(2**7, activation="relu")(optlearner_architecture)
    delta_d_hat_mu = Dense(85, activation="linear",
                           name="delta_d_hat_mu")(mu_architecture)
    print('delta_d_hat_mu shape: ' + str(delta_d_hat_mu.shape))
    sigma_architecture = Dense(2**7,
                               activation="relu")(optlearner_architecture)
    delta_d_hat_sigma = Dense(
        85, activation="softplus",
        name="delta_d_hat_sigma_uncorrected")(sigma_architecture)
    delta_d_hat_sigma = Lambda(lambda x: x + 1e-5,
                               name="delta_d_hat_sigma")(delta_d_hat_sigma)
    print('delta_d_hat_sigma shape: ' + str(delta_d_hat_sigma.shape))
    #exit(1)

    # Draw a sample from the distribution
    delta_d_hat = normal_sample(delta_d_hat_mu, delta_d_hat_sigma)
    delta_d_hat = Lambda(lambda x: x, name="delta_d_hat")(delta_d_hat)

    # Calculate the log probability of the sample and the entropy of the distribution
    #log_prob = normal_log_prob(delta_d_hat, delta_d_hat_mu, delta_d_hat_sigma)
    log_prob = normal_log_prob_pos(delta_d_hat, delta_d_hat_mu,
                                   delta_d_hat_sigma)
    #entropy = normal_entropy(delta_d_hat_sigma)
    #entropy = normal_entropy_pos(delta_d_hat_sigma)
    false_entropy = Lambda(lambda x: 1. / (1. + x),
                           name="false_entropy")(delta_d_hat_sigma)

    # Calculate the (batched) MSE between the learned and ground truth offset in the parameters
    delta_d_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(delta_d)
    delta_d_hat_square_diff = Lambda(lambda x: K.square(x[0] - x[1]),
                                     name="delta_d_hat_square_diff")(
                                         [delta_d_NOGRAD, delta_d_hat])
    delta_d_hat_mse = Lambda(lambda x: K.mean(K.square(x[0] - x[1]), axis=1))(
        [delta_d_NOGRAD, delta_d_hat])
    delta_d_hat_mse = Reshape(target_shape=(1, ),
                              name="delta_d_hat_mse")(delta_d_hat_mse)
    print("delta_d_hat loss shape: " + str(delta_d_hat_mse.shape))

    # Construct the loss function
    #advantage = Lambda(lambda x: K.abs(x[0]) - K.abs(x[0] - x[1]), name="adv_1_norm")([delta_d_NOGRAD, delta_d_hat])   # 1-norm advantage
    advantage = Lambda(lambda x: K.square(x[0]) - K.square(x[0] - x[1]),
                       name="adv_2_norm")([delta_d_NOGRAD,
                                           delta_d_hat])  # 2-norm advantage
    #advantage = Lambda(lambda x: K.mean(K.square(x[0]), axis=-1) - K.mean(K.square(x[0] - x[1]), axis=-1), name="adv_2_norm")([delta_d_NOGRAD, delta_d_hat])   # 2-norm advantage
    advantage_NOGRAD = Lambda(lambda x: K.stop_gradient(x),
                              name="adv_NOGRAD")(advantage)
    #eligibility = Multiply(name="eligibility")([log_prob, advantage_NOGRAD])
    eligibility = Multiply(name="eligibility")([log_prob, advantage])
    #false_loss_delta_d_hat = distributional_model_loss(eligibility, entropy, weighting=1.)
    weighting = 1.
    false_loss_delta_d_hat = Lambda(lambda x: -x[0] + weighting * x[1],
                                    name="dist_perf")(
                                        [eligibility, false_entropy])

    #penalty_func_mu = Lambda(lambda x: 0.1*K.exp(K.square(x) - np.pi**2))(delta_d_hat_mu)
    penalty_func_sigma = Lambda(lambda x: K.square(x))(delta_d_hat_sigma)
    #false_loss_delta_d_hat = Lambda(lambda x: K.mean(x[0] + x[1] + x[2], axis=1), name="reg_dist_goodness")([false_loss_delta_d_hat, penalty_func_mu, penalty_func_sigma])
    false_loss_delta_d_hat = Lambda(
        lambda x: K.mean(x[0] + x[1], axis=1),
        name="reg_dist_goodness")([false_loss_delta_d_hat, penalty_func_sigma])
    #false_loss_delta_d_hat = Lambda(lambda x: K.mean(x[0] + x[1], axis=1), name="reg_dist_goodness")([false_loss_delta_d_hat, penalty_func_mu])
    #false_loss_delta_d_hat = Lambda(lambda x: K.mean(x[0] + x[1], axis=1), name="reg_dist_goodness")([false_loss_delta_d_hat, delta_d_hat_square_diff])
    #false_loss_delta_d_hat = Lambda(lambda x: K.mean(x, axis=1), name="mean_dist_perf_unshaped")(false_loss_delta_d_hat)
    false_loss_delta_d_hat = Reshape(
        target_shape=(1, ), name="mean_dist_perf")(false_loss_delta_d_hat)

    # Get metrics
    #false_sin_loss_delta_d_hat = get_sin_metric(delta_d_NOGRAD, delta_d_hat)
    false_sin_loss_delta_d_hat = get_sin_metric(delta_d_NOGRAD,
                                                delta_d_hat,
                                                average=False)
    false_sin_loss_delta_d_hat = Lambda(
        lambda x: x, name="delta_d_hat_sin_output")(false_sin_loss_delta_d_hat)
    print("delta_d_hat sin loss shape: " +
          str(false_sin_loss_delta_d_hat.shape))

    # Prevent model from using the delta_d_hat gradient in final loss
    delta_d_hat_NOGRAD = Lambda(lambda x: K.stop_gradient(x),
                                name='optlearner_output_NOGRAD')(delta_d_hat)

    # False loss designed to pass the learned offset as a gradient to the embedding layer
    false_loss_smpl = Multiply(name="smpl_diff")(
        [optlearner_params, delta_d_hat_NOGRAD])
    print("smpl loss shape: " + str(false_loss_smpl.shape))

    return [optlearner_input, gt_params, gt_pc], [
        optlearner_params, false_loss_delta_d, optlearner_pc, false_loss_pc,
        false_loss_delta_d_hat, false_sin_loss_delta_d_hat, false_loss_smpl,
        delta_d, delta_d_hat, dist_angles, delta_d_hat_mu, delta_d_hat_sigma,
        delta_d_hat_mse
    ]
示例#11
0
def OptLearnerStaticArchitecture(param_trainable, init_wrapper, emb_size=1000):
    """ Optimised learner network architecture """
    # An embedding layer is required to optimise the parameters
    optlearner_input = Input(shape=(1, ), name="embedding_index")

    # Initialise the embedding layers
    emb_layers = init_emb_layers(optlearner_input, emb_size, param_trainable,
                                 init_wrapper)
    optlearner_params = Concatenate(name="parameter_embedding")(emb_layers)
    optlearner_params = Reshape(target_shape=(85, ),
                                name="learned_params")(optlearner_params)
    print("optlearner parameters shape: " + str(optlearner_params.shape))
    #exit(1)

    # Ground truth parameters and point cloud are inputs to the model as well
    gt_params = Input(shape=(85, ), name="gt_params")
    gt_pc = Input(shape=(6890, 3), name="gt_pc")
    print("gt parameters shape: " + str(gt_params.shape))
    print("gt point cloud shape: " + str(gt_pc.shape))

    # Compute the true offset (i.e. difference) between the ground truth and learned parameters
    delta_d = Lambda(lambda x: x[0] - x[1],
                     name="delta_d")([gt_params, optlearner_params])
    print("delta_d shape: " + str(delta_d.shape))
    #exit(1)

    # Calculate the (batched) MSE between the learned parameters and the ground truth parameters
    false_loss_delta_d = Lambda(lambda x: K.mean(K.square(x), axis=1))(delta_d)
    print("delta_d loss shape: " + str(false_loss_delta_d.shape))
    #exit(1)
    false_loss_delta_d = Reshape(target_shape=(1, ),
                                 name="delta_d_mse")(false_loss_delta_d)
    print("delta_d loss shape: " + str(false_loss_delta_d.shape))

    # Load SMPL model and get necessary parameters
    smpl_params = load_params(
        './keras_rotationnet_v2_demo_for_hidde/basicModel_f_lbs_10_207_0_v1.0.0.pkl'
    )
    _, _, input_info = get_parameters()
    input_betas = Lambda(lambda x: x[:, 72:82])(optlearner_params)
    input_pose_rodrigues = Lambda(lambda x: x[:, 0:72])(optlearner_params)
    input_trans = Lambda(lambda x: x[:, 82:85])(optlearner_params)

    # Get the point cloud corresponding to these parameters
    optlearner_pc = Points3DFromSMPLParams(input_betas, input_pose_rodrigues,
                                           input_trans, smpl_params,
                                           input_info)
    print("optlearner point cloud shape: " + str(optlearner_pc.shape))
    #exit(1)

    # Get the (batched) Euclidean loss between the learned and ground truth point clouds
    pc_euclidean_diff = Lambda(lambda x: x[0] - x[1])([gt_pc, optlearner_pc])
    pc_euclidean_dist = Lambda(lambda x: K.sum(K.square(x), axis=-1))(
        pc_euclidean_diff)
    print('pc euclidean dist ' + str(pc_euclidean_dist.shape))
    #exit(1)
    false_loss_pc = Lambda(lambda x: K.mean(x, axis=1))(pc_euclidean_dist)
    false_loss_pc = Reshape(target_shape=(1, ),
                            name="pc_mean_euc_dist")(false_loss_pc)
    print("point cloud loss shape: " + str(false_loss_pc.shape))

    # Learn the offset in parameters from the difference between the ground truth and learned point clouds
    pc_euclidean_diff_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(
        pc_euclidean_diff
    )  # This is added to avoid influencing embedding layer parameters by a "bad" gradient network
    print("shape of output: " + str(pc_euclidean_diff_NOGRAD.shape))
    vertices = [1850, 1600, 2050, 1300, 5350, 5050, 5500]
    pc_euclidean_diff_NOGRAD = Lambda(lambda x: K.tf.gather(
        x, np.array(vertices).astype(np.int32), axis=-2))(
            pc_euclidean_diff_NOGRAD)
    print("shape of output: " + str(pc_euclidean_diff_NOGRAD.shape))
    #exit(1)
    pc_euclidean_diff_NOGRAD = Flatten()(pc_euclidean_diff_NOGRAD)

    optlearner_architecture = Dense(
        2**11, activation="relu")(pc_euclidean_diff_NOGRAD)
    optlearner_architecture = BatchNormalization()(optlearner_architecture)
    optlearner_architecture = Dropout(0.5)(optlearner_architecture)
    print('optlearner_architecture ' + str(optlearner_architecture.shape))
    #exit(1)
    delta_d_hat = Dense(85, activation="linear",
                        name="delta_d_hat")(optlearner_architecture)
    print('delta_d_hat shape ' + str(delta_d_hat.shape))
    #exit(1)

    # Calculate the (batched) MSE between the learned and ground truth offset in the parameters
    delta_d_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(delta_d)
    false_loss_delta_d_hat = Lambda(
        lambda x: K.sum(K.square(x[0] - x[1]), axis=1))(
            [delta_d_NOGRAD, delta_d_hat])
    false_loss_delta_d_hat = Reshape(
        target_shape=(1, ), name="delta_d_hat_mse")(false_loss_delta_d_hat)
    print("delta_d_hat loss shape: " + str(false_loss_delta_d_hat.shape))
    false_sin_loss_delta_d_hat = get_sin_metric(delta_d_NOGRAD, delta_d_hat)
    print("delta_d_hat sin loss shape: " +
          str(false_sin_loss_delta_d_hat.shape))

    # Prevent model from using the delta_d_hat gradient in final loss
    delta_d_hat_NOGRAD = Lambda(lambda x: K.stop_gradient(x),
                                name='optlearner_output_NOGRAD')(delta_d_hat)

    # False loss designed to pass the learned offset as a gradient to the embedding layer
    false_loss_smpl = Lambda(lambda x: x[1] * x[0], name="smpl_diff")(
        [optlearner_params, delta_d_hat_NOGRAD])
    print("smpl loss shape: " + str(false_loss_smpl.shape))

    return [optlearner_input, gt_params, gt_pc], [
        optlearner_params, false_loss_delta_d, optlearner_pc, false_loss_pc,
        false_loss_delta_d_hat, false_sin_loss_delta_d_hat, false_loss_smpl,
        delta_d, delta_d_hat, delta_d_hat_NOGRAD
    ]