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
]
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 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
]