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]
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 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 GroupedConv1DOptLearnerArchitecture(param_trainable, init_wrapper, smpl_params, input_info, faces, emb_size=1000, input_type="3D_POINTS", groups=[], DROPOUT=0.0): """ Optimised learner network architecture """ # Get network inputs optlearner_params, gt_params, gt_pc = standard_input(emb_size, param_trainable, init_wrapper) # Assert all pose parameters are specified in group hierarchy assert np.all([joint in np.flatten(groups) for joint in range(0, 24)]) # Compute the true offset (i.e. difference) between the ground truth and learned parameters delta_d = get_delta_d(gt_params, optlearner_params) # Calculate the (batched) MSE between the learned parameters and the ground truth parameters false_loss_delta_d = get_delta_d_loss(delta_d) # Load SMPL model and get necessary parameters optlearner_pc = get_pc(optlearner_params, smpl_params, input_info, faces) print("optlearner_pc shape: " + str(optlearner_pc.shape)) # Get the (batched) Euclidean loss between the learned and ground truth point clouds false_loss_pc = get_pc_loss(gt_pc, optlearner_pc) # Gather sets of points and compute their cross product to get mesh normals vertex_list = get_vertices() vertex_diff_NOGRAD = get_vertex_diff(gt_pc, optlearner_pc, vertex_list) gt_normals, opt_normals = get_normals_from_vertices(gt_pc, optlearner_pc, vertex_list) diff_normals = get_normals_diff(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, dist_angles = get_normals_angle_diff(gt_normals, opt_normals) diff_angles_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(diff_angles) dist_angles_NOGRAD = Lambda(lambda x: K.stop_gradient(x))(dist_angles) # Learn the offset in parameters from the difference between the ground truth and learned mesh normals 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")(diff_normals_NOGRAD) if input_type == "COMBINED": deep_opt_input = Dense(2**9, activation="relu")(mesh_diff_NOGRAD) print('deep_opt_input shape: '+str(deep_opt_input.shape)) # Predict pose in the specified groups if output_type == "6D": indices_ordering, group_outputs = grouped_6D_network(deep_opt_input, groups, DROPOUT) # Get delta_d pose parameters in orthogonal 6D representation delta_d_pose = delta_d_to_ortho6d(delta_d_NOGRAD, indices_ordering) # Get delta_d_hat pose parameters delta_d_hat_pose = Concatenate(axis=-1)(group_outputs) # Get loss for pose parameters false_loss_delta_d_hat = delta_d_hat_mse(delta_d_pose, delta_d_hat_pose) # Convert delta_d_hat to Rodrigues form delta_d_hat_rod_pose = ortho6d_pose_to_rodrigues(delta_d_hat_pose) # Predict shape and translation parameters shape_params = shape_network(deep_opt_input, DROPOUT) indices_ordering += [i for i in range(72, 85)] print("indices_ordering: " + strindices_ordering)) # Get loss on shape and translation parameters # Process indices ordering to re-order array reordered_indices = reorder_indices(indices_ordering) # Re-order predictions delta_d_hat = collect_and_order_outputs(group_outputs, reordered_indices) elif output_type == "Rodrigues": indices_ordering, group_outputs = grouped_network(deep_opt_input, groups, DROPOUT) # Predict shape and translation parameters shape_params = shape_network(deep_opt_input, DROPOUT) group_outputs.append(shape_params) indices_ordering += [i for i in range(72, 85)] print("indices_ordering: " + strindices_ordering)) # Process indices ordering to re-order array reordered_indices = reorder_indices(indices_ordering) # Re-order predictions delta_d_hat = collect_and_order_outputs(group_outputs, reordered_indices) # Calculate loss on predictions false_loss_delta_d_hat = delta_d_hat_mse(delta_d_NOGRAD, delta_d_hat) else: print("Output type '{}' not recognised. Exiting.".format(output_type)) exit(1) # Metrics #false_sin_loss_delta_d_hat = paramwise_sin_metric(delta_d_NOGRAD, delta_d_hat) false_sin_loss_delta_d_hat = get_angular_distance_metric(delta_d_NOGRAD, delta_d_hat) per_param_mse = param_metric(delta_d_NOGRAD, delta_d_hat, sine_metric=True) # 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]
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 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 ]
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 ]
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 ]