diff --git a/custom_utils/smpl.py b/custom_utils/smpl.py index 5e7dc34..f888576 100644 --- a/custom_utils/smpl.py +++ b/custom_utils/smpl.py @@ -34,7 +34,7 @@ eval_joints = { class SMPL: - def __init__(self): + def __init__(self, orig_smpl=False, use_one_hot=False): self.J_regressor = [] self.v_template = [] self.shapedirs = [] @@ -82,7 +82,23 @@ class SMPL: self.to_keep = sio.loadmat(os.path.join('custom_utils/keep_remove.mat'))['keep'].reshape(-1) self.weights_all = self.weights - self.weights = tf.gather(self.weights, self.to_keep, axis=1) + if use_one_hot: + self.weights = tf.gather(self.weights_one_hot, self.to_keep, axis=0) + self.weights = tf.tile(self.weights[tf.newaxis,:,:], [2,1,1]) + else: + self.weights = tf.gather(self.weights, self.to_keep, axis=1) + + self.shape = (1,27578,3) + + if orig_smpl: + self.to_keep = self.to_keep[self.to_keep < 6890] + self.J_regressor = self.J_regressor[:,:,:6890] + self.v_template = self.v_template[:,:6890,:] + self.shapedirs = self.shapedirs[:,:6890,:,:] + self.posedirs = self.posedirs[:,:6890,:,:] + self.weights_all = self.weights_all[:,:6890,:] + self.weights = tf.gather(self.weights_all, self.to_keep, axis=1) + self.shape = (1, 6890, 3) def Q2R(self, w, x, y, z): def diag(a, b): # computes the diagonal entries, 1 - 2*a**2 - 2*b**2 @@ -141,8 +157,8 @@ class SMPL: #v_template = tf.gather(self.v_template, gender, axis=0) v_shaped = [] - v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[0], axes=[[1], [2]]) + tf.reshape(self.v_template[0], shape=(1,27578,3))) # female - v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=(1,27578,3))) # male + v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[0], axes=[[1], [2]]) + tf.reshape(self.v_template[0], shape=self.shape)) # female + v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=self.shape)) # male v_shaped = tf.stack(v_shaped, axis=-1) def single_run(x, J_regressor=self.J_regressor, posedirs=self.posedirs, weights=self.weights_all, parent=self.parent): @@ -187,11 +203,12 @@ class SMPL: return T_out, J_out, v_posed + def run_with_weights(self, R, betas, v_cloth, weights=None, filter=None): v_shaped = [] - v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[0], axes=[[1], [2]]) + tf.reshape(self.v_template[0], shape=(1,27578,3))) # female - v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=(1,27578,3))) # male + v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[0], axes=[[1], [2]]) + tf.reshape(self.v_template[0], shape=self.shape)) # female + v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=self.shape)) # male v_shaped = tf.stack(v_shaped, axis=-1) if weights is None: @@ -236,11 +253,14 @@ class SMPL: return T_out - def get_joints_template(self, betas): - v_shaped = [] - v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[0], axes=[[1], [2]]) + tf.reshape(self.v_template[0], shape=(1,27578,3))) # female - v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=(1,27578,3))) # male - v_shaped = tf.stack(v_shaped, axis=-1) + + def get_joints_template(self, betas, v_shaped=None): + + if v_shaped is None: + v_shaped = [] + v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[0], axes=[[1], [2]]) + tf.reshape(self.v_template[0], shape=self.shape)) # female + v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=self.shape)) # male + v_shaped = tf.stack(v_shaped, axis=-1) def single_run(x, J_regressor=self.J_regressor): v_shaped = x[0] @@ -254,9 +274,31 @@ class SMPL: return tf.subtract(J, J0) + def get_root_joint(self, betas, v_shaped=None): + + if v_shaped is None: + v_shaped = [] + v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[0], axes=[[1], [2]]) + tf.reshape(self.v_template[0], shape=self.shape)) # female + v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=self.shape)) # male + v_shaped = tf.stack(v_shaped, axis=-1) - def run_garment(self, pose, betas, v_cloth, filter=None): - J = self.get_joints_template(betas) + def single_run(x, J_regressor=self.J_regressor): + v_shaped = x[0] + g = x[1] + J = tf.matmul(J_regressor[g], v_shaped[:, :, g]) + return J + + gender = tf.cast(tf.greater(betas[:, 10] , 0.5), tf.int32) + J = tf.map_fn(single_run, (v_shaped, gender), (tf.float32)) + J0 = tf.reshape(J[:,0,:], (-1, 1, 3)) + + return J[:,0,:] + + # return J0 + + + def run_garment(self, pose, betas, v_cloth, v_shaped=None): + J = self.get_joints_template(betas, v_shaped) R = self.run_rodrigues(pose) # if filter is not None: @@ -269,7 +311,7 @@ class SMPL: v_cloth = x[3] R_cube_big = R - + results = [] results.append(self.with_zeros(tf.concat((R_cube_big[0], tf.reshape(J[0, :], (3, 1))), axis=1))) for j in range(1, 24): @@ -294,8 +336,59 @@ class SMPL: J_out = tf.subtract(J_out, J0) T_out = tf.subtract(T_out, J0) + return T_out + + def run_garment_orig(self, R, betas, v_cloth, J=None, weights=None): + batch_size = R.shape[0] + + if J is None: + J = self.get_joints_template(betas) + + if weights is None: + weights = tf.gather(self.weights_one_hot, self.to_keep, axis=0) + # else: + # weights = tf.gather(weights, self.to_keep, axis=0) + + def single_run(x, J_regressor=self.J_regressor, posedirs=self.posedirs, weights=weights, parent=self.parent): + R = x[0] + J = x[1] + g = x[2] + v_cloth = x[3] + + R_cube_big = R + # R_cube = R_cube_big[1:] + # I_cube = tf.expand_dims(tf.eye(3, dtype=tf.float32), axis=0) + tf.zeros((R_cube.get_shape()[0], 3, 3), dtype=tf.float32) + # lrotmin = tf.squeeze(tf.reshape((R_cube - I_cube), (-1, 1))) + + results = [] + results.append(self.with_zeros(tf.concat((R_cube_big[0], tf.reshape(J[0, :], (3, 1))), axis=1))) + for j in range(1, 24): + results.append(tf.matmul(results[parent[j]], self.with_zeros(tf.concat((R_cube_big[j], tf.reshape(J[j, :] - J[parent[j], :], (3, 1))), axis=1)))) + stacked = tf.stack(results, axis=0) + results = stacked - self.pack(tf.matmul(stacked, tf.reshape(tf.concat((J, tf.zeros((24, 1), dtype=tf.float32)), axis=1), (24, 4, 1)))) + + J = tf.matmul(results, tf.expand_dims(tf.concat([J, tf.ones([24, 1], tf.float32)], axis=1), axis=1), transpose_b=True) + J = tf.gather(J, tf.constant([0, 1, 2], tf.int32), axis=1) + + T = tf.tensordot(weights, results, axes=((1), (0))) + rest_shape_h = tf.concat((v_cloth, tf.ones((v_cloth.get_shape().as_list()[0], 1), dtype=tf.float32)), axis=1) + v = tf.matmul(T, tf.reshape(rest_shape_h, (-1, 4, 1))) + v = tf.reshape(v, (-1, 4))[:, :3] + + return v, tf.squeeze(J) + + gender = tf.cast(tf.greater_equal(betas[:, 10] , 0), tf.int32) + T_out, J_out = tf.map_fn(single_run, (R, J, gender, v_cloth), (tf.float32, tf.float32)) + + J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3)) + J_out = tf.subtract(J_out, J0) + T_out = tf.subtract(T_out, J0) + + return T_out, J_out + + def run_garment_with_weights(self, R, betas, v_cloth, weights, J=None): if J is None: @@ -335,9 +428,14 @@ class SMPL: return T_out + def run_backward(self, R, J, v_posed, weights): batch_size = R.get_shape()[0] + if weights.shape[0] != 14475: + weights = tf.gather(self.weights_one_hot, self.to_keep, axis=0) + + J_relative = J[:, 1:, :] - tf.gather(J, tf.constant(list(self.parent.values()), tf.int32), axis=1) # batch_size x 23 x 3 J_out = []