luni/cvpr
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add new options to smpl code

Browse Source
This commit is contained in:
luni
2022-10-29 20:18:09 +00:00
parent 3544027d2e
commit c14849f88a
1 changed files with 112 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
custom_utils/smpl.py
+108 -10
View File
@@ -34,7 +34,7 @@ eval_joints = {
class SMPL: class SMPL:
def __init__(self): def __init__(self, orig_smpl=False, use_one_hot=False):
self.J_regressor = [] self.J_regressor = []
self.v_template = [] self.v_template = []
self.shapedirs = [] self.shapedirs = []
@@ -82,8 +82,24 @@ class SMPL:
self.to_keep = sio.loadmat(os.path.join('custom_utils/keep_remove.mat'))['keep'].reshape(-1) self.to_keep = sio.loadmat(os.path.join('custom_utils/keep_remove.mat'))['keep'].reshape(-1)
self.weights_all = self.weights self.weights_all = self.weights
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.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 Q2R(self, w, x, y, z):
def diag(a, b): # computes the diagonal entries, 1 - 2*a**2 - 2*b**2 def diag(a, b): # computes the diagonal entries, 1 - 2*a**2 - 2*b**2
return 1 - 2 * tf.pow(a, 2) - 2 * tf.pow(b, 2) return 1 - 2 * tf.pow(a, 2) - 2 * tf.pow(b, 2)
@@ -141,8 +157,8 @@ class SMPL:
#v_template = tf.gather(self.v_template, gender, axis=0) #v_template = tf.gather(self.v_template, gender, axis=0)
v_shaped = [] 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[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=(1,27578,3))) # male 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) 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): 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 return T_out, J_out, v_posed
def run_with_weights(self, R, betas, v_cloth, weights=None, filter=None): def run_with_weights(self, R, betas, v_cloth, weights=None, filter=None):
v_shaped = [] 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[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=(1,27578,3))) # male 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) v_shaped = tf.stack(v_shaped, axis=-1)
if weights is None: if weights is None:
@@ -236,10 +253,13 @@ class SMPL:
return T_out return T_out
def get_joints_template(self, betas):
def get_joints_template(self, betas, v_shaped=None):
if v_shaped is None:
v_shaped = [] 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[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=(1,27578,3))) # male 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) v_shaped = tf.stack(v_shaped, axis=-1)
def single_run(x, J_regressor=self.J_regressor): def single_run(x, J_regressor=self.J_regressor):
@@ -254,9 +274,31 @@ class SMPL:
return tf.subtract(J, J0) return tf.subtract(J, J0)
def get_root_joint(self, betas, v_shaped=None):
def run_garment(self, pose, betas, v_cloth, filter=None): if v_shaped is None:
J = self.get_joints_template(betas) 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]
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) R = self.run_rodrigues(pose)
# if filter is not None: # if filter is not None:
@@ -294,8 +336,59 @@ class SMPL:
J_out = tf.subtract(J_out, J0) J_out = tf.subtract(J_out, J0)
T_out = tf.subtract(T_out, J0) T_out = tf.subtract(T_out, J0)
return T_out 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): def run_garment_with_weights(self, R, betas, v_cloth, weights, J=None):
if J is None: if J is None:
@@ -335,9 +428,14 @@ class SMPL:
return T_out return T_out
def run_backward(self, R, J, v_posed, weights): def run_backward(self, R, J, v_posed, weights):
batch_size = R.get_shape()[0] 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_relative = J[:, 1:, :] - tf.gather(J, tf.constant(list(self.parent.values()), tf.int32), axis=1) # batch_size x 23 x 3
J_out = [] J_out = []