diff --git a/.gitignore b/.gitignore index 469e186..efc27db 100644 --- a/.gitignore +++ b/.gitignore @@ -141,3 +141,5 @@ output/* models/ models_archive/ temp.jpg +corrupt/ +wandb copy/ \ No newline at end of file diff --git a/custom_utils/NN.mat b/custom_utils/NN.mat new file mode 100644 index 0000000..f2e3ac2 Binary files /dev/null and b/custom_utils/NN.mat differ diff --git a/custom_utils/custom_utils.py b/custom_utils/custom_utils.py index 71a1e98..1a59bcb 100644 --- a/custom_utils/custom_utils.py +++ b/custom_utils/custom_utils.py @@ -89,8 +89,8 @@ def maskFacesBatch(V, F, mask): return Vs, Fs -def plot_3d(mesh_gt, mesh_pred, faces, collision, collision_faces, collision_color): - mesh_pred[:,1] += 1 +def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color): + mesh_pred_unposed[:,1] += 1 collision[:,1] += 2 fig = go.Figure(data=[ go.Mesh3d( @@ -105,9 +105,9 @@ def plot_3d(mesh_gt, mesh_pred, faces, collision, collision_faces, collision_col ), go.Mesh3d( # 8 vertices of a cube - x=mesh_pred[:,0], - y=mesh_pred[:,1], - z=mesh_pred[:,2], + x=mesh_pred_unposed[:,0], + y=mesh_pred_unposed[:,1], + z=mesh_pred_unposed[:,2], # i, j and k give the vertices of triangles i = faces[:,0], j = faces[:,1], diff --git a/custom_utils/keep_remove.mat b/custom_utils/keep_remove.mat new file mode 100644 index 0000000..213b672 Binary files /dev/null and b/custom_utils/keep_remove.mat differ diff --git a/custom_utils/smpl.py b/custom_utils/smpl.py new file mode 100644 index 0000000..a6f38eb --- /dev/null +++ b/custom_utils/smpl.py @@ -0,0 +1,347 @@ +import os +import sys + +import pickle +import scipy.io as sio +from math import * + +import numpy as np +import tensorflow as tf + +# Left-Right +eval_joints = { + 'All': [1, 2, 3, 4, 5, 7, 8, 9, 12, 15, 16, 17, 18, 19, 20, 21], + 'Head': [12, 15], # Neck, Head + 'Torso': [3, 6, 9], # Back joints + 'Shoulder': [13, 14, 16, 17], # Inner, Outer + 'Elbow': [18, 19], + 'Wrist': [20, 21], + 'Hip': [1, 2], + 'Knee': [4, 5], + 'Foot': [7, 8, 10, 11]} # Ankle, Foot + +class SMPL: + + def __init__(self): + self.J_regressor = [] + self.v_template = [] + self.shapedirs = [] + self.posedirs = [] + self.weights = [] + # Load SMPL models + # path = os.path.abspath(os.path.dirname(__file__)) + with open(os.path.join('custom_utils/super_model_f.pkl'), 'rb') as f: + if sys.version_info[0]<3: + female = pickle.load(f) + else: + female = pickle.load(f, encoding="latin1") + self.J_regressor.append(tf.constant(np.array(female['J_regressor'].todense(), dtype=np.float32))) + self.v_template.append(tf.constant(female['v_template'], dtype=np.float32)) + self.posedirs.append(tf.constant(female['posedirs'], dtype=np.float32)) + self.shapedirs.append(tf.constant(female['shapedirs'], dtype=np.float32)) + self.weights.append(tf.constant(female['weights'], dtype=np.float32)) + with open(os.path.join('custom_utils/super_model_m.pkl'), 'rb') as f: + if sys.version_info[0]<3: + male = pickle.load(f) + else: + male = pickle.load(f, encoding="latin1") + self.J_regressor.append(tf.constant(np.array(male['J_regressor'].todense(), dtype=np.float32))) + self.v_template.append(tf.constant(male['v_template'], dtype=np.float32)) + self.posedirs.append(tf.constant(male['posedirs'], dtype=np.float32)) + self.shapedirs.append(tf.constant(male['shapedirs'], dtype=np.float32)) + self.weights.append(tf.constant(male['weights'], dtype=np.float32)) + kintree_table = male['kintree_table'] + id_to_col = {kintree_table[1, i]: i for i in range(kintree_table.shape[1])} + self.parent = {i: id_to_col[kintree_table[0, i]] for i in range(1, kintree_table.shape[1])} + self.J_regressor = tf.stack(self.J_regressor, axis=0) + self.v_template = tf.stack(self.v_template, axis=0) + self.shapedirs = tf.stack(self.shapedirs, axis=0) + self.posedirs = tf.stack(self.posedirs, axis=0) + self.weights = tf.stack(self.weights, axis=0) + self.weights_one_hot = tf.reduce_mean(self.weights, axis=0) + self.weights_one_hot = tf.math.argmax(self.weights_one_hot, axis=-1) + self.weights_one_hot = tf.one_hot(self.weights_one_hot, 24) + + # Topology no head/hands/feet + F = sio.loadmat('custom_utils/ssmpl_faces.mat') + self.f = np.array(F['faces_default']) + self.f_skirt = np.array(F['faces_skirt']) + self.nn = np.array(sio.loadmat('custom_utils/NN.mat')) + self.to_keep = sio.loadmat(os.path.join('custom_utils/keep_remove.mat'))['keep'].reshape(-1) + + self.weights = tf.gather(self.weights, self.to_keep, axis=1) + + def Q2R(self, w, x, y, z): + 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) + + def tr_add(a, b, c, d): # computes triangle entries with addition + return 2 * a * b + 2 * c * d + + def tr_sub(a, b, c, d): # computes triangle entries with subtraction + return 2 * a * b - 2 * c * d + + m = [[diag(y, z), tr_sub(x, y, z, w), tr_add(x, z, y, w)], + [tr_add(x, y, z, w), diag(x, z), tr_sub(y, z, x, w)], + [tr_sub(x, z, y, w), tr_add(y, z, x, w), diag(x, y)]] + + R = tf.stack([tf.stack(m[i], axis=-1) for i in range(3)], axis=-2) + return R + + def rodrigues(self, r): + theta = tf.norm(r, axis=(1, 2), keepdims=True) + theta = tf.maximum(theta, np.finfo(np.float32).eps) + r_hat = r / theta + z_stick = tf.zeros(theta.get_shape().as_list()[0], dtype=tf.float32) + m = tf.stack( + (z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick, + -r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick), axis=1) + m = tf.reshape(m, (-1, 3, 3)) + i_cube = tf.expand_dims(tf.eye(3, dtype=tf.float32), axis=0) + tf.zeros( + (theta.get_shape().as_list()[0], 3, 3), dtype=tf.float32) + # R = cos * i_cube + (1 - cos) * dot + tf.sin(theta) * m + R = i_cube + m * tf.sin(theta) + tf.matmul(m, m) * (1 - tf.cos(theta)) + return R + + def with_zeros(self, x): + return tf.concat((x, tf.constant([[0.0, 0.0, 0.0, 1.0]], dtype=tf.float32)), axis=0) + + def pack(self, x): + return tf.concat((tf.zeros((x.get_shape().as_list()[0], 4, 3), dtype=tf.float32), x), axis=2) + + def run_rodrigues(self, pose): + batch_size = pose.shape[0] + R_out = self.rodrigues(tf.reshape(pose, (-1, 1, 3))) + R_out = tf.reshape(R_out, (batch_size, -1, 3, 3)) + #R_out = [] + #for i in range(batch_size): + # R_out.append(self.rodrigues(tf.reshape(pose[i,:], (-1, 1, 3)))) + #R_out = tf.stack(R_out, axis=0) + return R_out + + def run(self, R, betas, remove=True): + batch_size = betas.shape[0] + + #J_regressor = tf.gather(self.J_regressor, gender, axis=0) + #shapedirs = tf.gather(self.shapedirs, gender, axis=0) + #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 = tf.stack(v_shaped, axis=-1) + + def single_run(x, J_regressor=self.J_regressor, posedirs=self.posedirs, weights=self.weights, parent=self.parent): + R = x[0] + v_shaped = x[1] + g = x[2] + + J = tf.matmul(J_regressor[g], v_shaped[:, :, g]) + 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))) + v_posed = v_shaped[:, :, g] + tf.tensordot(posedirs[g], lrotmin, axes=[[2], [0]]) + 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[g], results, axes=((1), (0))) + rest_shape_h = tf.concat((v_posed, tf.ones((v_posed.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), v_posed + + gender = tf.cast(tf.greater_equal(betas[:, 10] , 0), tf.int32) + T_out, J_out, v_posed = tf.map_fn(single_run, (R, v_shaped, gender), (tf.float32, 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) + v_posed = tf.subtract(v_posed, J0) + + if remove: + T_out = tf.gather(T_out, self.to_keep, axis=1) + v_posed = tf.gather(v_posed, self.to_keep, axis=1) + + 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 = tf.stack(v_shaped, axis=-1) + + if weights is None: + w = self.weights_one_hot[np.newaxis, :, :] + if filter is not None: + w = tf.gather(tf.gather(w, self.to_keep, axis=1), filter, axis=1) + if v_cloth.shape[1] == 14475: + w = tf.gather(w, self.to_keep, axis=1) + else: + w = weights + + def single_run(x, J_regressor=self.J_regressor, w=w, parent=self.parent): + R = x[0] + v_shaped = x[1] + g = x[2] + v_cloth = x[3] + + J = tf.matmul(J_regressor[g], v_shaped[:, :, g]) + 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): + 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(w[0], results, axes=((1), (0))) + rest_shape_h = tf.concat((v_cloth, tf.ones((v_cloth.shape[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, v_shaped, gender, v_cloth), (tf.float32, tf.float32)) + + J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3)) + T_out = tf.subtract(T_out, J0) + + 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 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), tf.int32) + J = tf.map_fn(single_run, (v_shaped, gender), (tf.float32)) + J0 = tf.reshape(J[:,0,:], (-1, 1, 3)) + + return tf.subtract(J, J0) + + + def run_garment(self, pose, betas, v_cloth, filter=None): + J = self.get_joints_template(betas) + R = self.run_rodrigues(pose) + + # if filter is not None: + # weights = tf.boolean_mask(self.weights, filter) + + def single_run(x, weights=self.weights, parent=self.parent): + R = x[0] + J = x[1] + g = x[2] + 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): + 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[g, :, :], results, axes=((1), (0))) + rest_shape_h = tf.concat((v_cloth, tf.ones((v_cloth.shape[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(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 + + def run_garment_with_weights(self, R, betas, v_cloth, weights, J=None): + + if J is None: + J = self.get_joints_template(betas) + + def single_run(x, parent=self.parent): + R = x[0] + J = x[1] + v_cloth = x[2] + weights = 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): + 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, v_cloth, weights), (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 + + def run_backward(self, R, J, v_posed, weights): + batch_size = R.get_shape()[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 = [] + J_out.append(tf.expand_dims(J[:, 0, :], axis=-1)) # batch_size x 3 x 1 + results = [] + results.append(tf.expand_dims(tf.eye(3, dtype=tf.float32), axis=0) + tf.zeros((batch_size, 3, 3), dtype=tf.float32)) + for j in range(1, 24): + results.append(tf.matmul(R[:, self.parent[j], :, :], results[self.parent[j]])) + J_out.append(tf.matmul(results[j], tf.expand_dims(J_relative[:, j-1, :], axis=-1)) + J_out[self.parent[j]]) # batch_size x 3 x 1 + results = tf.transpose(tf.stack(results, axis=0), [1, 0, 2, 3]) # (batch_size x 24 x 3 x 3) + J_out = tf.transpose(tf.stack(J_out, axis=0), [1, 0, 2, 3]) # (batch_size x 24 x 3 x 1) + + results = tf.matmul(R, results) + T = tf.transpose(tf.tensordot(results, weights, axes=((1), (1))), [0, 3, 1, 2]) # (batch_size x 6890 x 3 x 3) + v_posed = tf.expand_dims(v_posed, axis=-1) # (batch_size x 6890 x 3 x 1) + J = tf.expand_dims(J, axis=-1) # (batch_size x 24 x 3 x 1) + v = tf.matmul(T, v_posed) + tf.transpose(tf.tensordot(J_out - tf.matmul(results, J), weights, axes=((1), (1))), [0, 3, 1, 2]) # (batch_size x 6890 x 3 x 1) + v = tf.squeeze(v, -1) # (batch_size x 6890 x 3) + J_out = tf.squeeze(J_out, -1) # (batch_size x 24 x 3) + + return v, J_out \ No newline at end of file diff --git a/custom_utils/smpl_2.py b/custom_utils/smpl_2.py new file mode 100644 index 0000000..af2deac --- /dev/null +++ b/custom_utils/smpl_2.py @@ -0,0 +1,345 @@ +import os +import sys + +import pickle +import scipy.io as sio +from math import * + +import numpy as np +import tensorflow as tf + +# Left-Right +eval_joints = { + 'All': [1, 2, 3, 4, 5, 7, 8, 9, 12, 15, 16, 17, 18, 19, 20, 21], + 'Head': [12, 15], # Neck, Head + 'Torso': [3, 6, 9], # Back joints + 'Shoulder': [13, 14, 16, 17], # Inner, Outer + 'Elbow': [18, 19], + 'Wrist': [20, 21], + 'Hip': [1, 2], + 'Knee': [4, 5], + 'Foot': [7, 8, 10, 11]} # Ankle, Foot + +class SMPL: + + def __init__(self): + self.J_regressor = [] + self.v_template = [] + self.shapedirs = [] + self.posedirs = [] + self.weights = [] + # Load SMPL models + # path = os.path.abspath(os.path.dirname(__file__)) + with open(os.path.join('custom_utils/super_model_f.pkl'), 'rb') as f: + if sys.version_info[0]<3: + female = pickle.load(f) + else: + female = pickle.load(f, encoding="latin1") + self.J_regressor.append(tf.constant(np.array(female['J_regressor'].todense(), dtype=np.float32))) + self.v_template.append(tf.constant(female['v_template'], dtype=np.float32)) + self.posedirs.append(tf.constant(female['posedirs'], dtype=np.float32)) + self.shapedirs.append(tf.constant(female['shapedirs'], dtype=np.float32)) + self.weights.append(tf.constant(female['weights'], dtype=np.float32)) + with open(os.path.join('custom_utils/super_model_m.pkl'), 'rb') as f: + if sys.version_info[0]<3: + male = pickle.load(f) + else: + male = pickle.load(f, encoding="latin1") + self.J_regressor.append(tf.constant(np.array(male['J_regressor'].todense(), dtype=np.float32))) + self.v_template.append(tf.constant(male['v_template'], dtype=np.float32)) + self.posedirs.append(tf.constant(male['posedirs'], dtype=np.float32)) + self.shapedirs.append(tf.constant(male['shapedirs'], dtype=np.float32)) + self.weights.append(tf.constant(male['weights'], dtype=np.float32)) + kintree_table = male['kintree_table'] + id_to_col = {kintree_table[1, i]: i for i in range(kintree_table.shape[1])} + self.parent = {i: id_to_col[kintree_table[0, i]] for i in range(1, kintree_table.shape[1])} + self.J_regressor = tf.stack(self.J_regressor, axis=0) + self.v_template = tf.stack(self.v_template, axis=0) + self.shapedirs = tf.stack(self.shapedirs, axis=0) + self.posedirs = tf.stack(self.posedirs, axis=0) + self.weights = tf.stack(self.weights, axis=0) + self.weights_one_hot = tf.reduce_mean(self.weights, axis=0) + self.weights_one_hot = tf.math.argmax(self.weights_one_hot, axis=-1) + self.weights_one_hot = tf.one_hot(self.weights_one_hot, 24) + + # Topology no head/hands/feet + F = sio.loadmat('custom_utils/ssmpl_faces.mat') + self.f = np.array(F['faces_default']) + self.f_skirt = np.array(F['faces_skirt']) + self.nn = np.array(sio.loadmat('custom_utils/NN.mat')) + self.to_keep = sio.loadmat(os.path.join('custom_utils/keep_remove.mat'))['keep'].reshape(-1) + + self.weights = tf.gather(self.weights, self.to_keep, axis=1) + # self.shapedirs = tf.gather(self.shapedirs, self.to_keep, axis=1) + # self.posedirs = tf.gather(self.posedirs, self.to_keep, axis=1) + # self.J_regressor = tf.gather(self.J_regressor, self.to_keep, axis=2) + # self.v_template = tf.gather(self.v_template, self.to_keep, axis=1) + + + def Q2R(self, w, x, y, z): + 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) + + def tr_add(a, b, c, d): # computes triangle entries with addition + return 2 * a * b + 2 * c * d + + def tr_sub(a, b, c, d): # computes triangle entries with subtraction + return 2 * a * b - 2 * c * d + + m = [[diag(y, z), tr_sub(x, y, z, w), tr_add(x, z, y, w)], + [tr_add(x, y, z, w), diag(x, z), tr_sub(y, z, x, w)], + [tr_sub(x, z, y, w), tr_add(y, z, x, w), diag(x, y)]] + + R = tf.stack([tf.stack(m[i], axis=-1) for i in range(3)], axis=-2) + return R + + def rodrigues(self, r): + theta = tf.norm(r, axis=(1, 2), keepdims=True) + theta = tf.maximum(theta, np.finfo(np.float32).eps) + r_hat = r / theta + z_stick = tf.zeros(theta.get_shape().as_list()[0], dtype=tf.float32) + m = tf.stack( + (z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick, + -r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick), axis=1) + m = tf.reshape(m, (-1, 3, 3)) + i_cube = tf.expand_dims(tf.eye(3, dtype=tf.float32), axis=0) + tf.zeros( + (theta.get_shape().as_list()[0], 3, 3), dtype=tf.float32) + # R = cos * i_cube + (1 - cos) * dot + tf.sin(theta) * m + R = i_cube + m * tf.sin(theta) + tf.matmul(m, m) * (1 - tf.cos(theta)) + return R + + def with_zeros(self, x): + return tf.concat((x, tf.ones((x.shape[0], 1, 4)) * tf.constant([[0.0, 0.0, 0.0, 1.0]], dtype=tf.float32)[tf.newaxis,:,:]), axis=1) + + def pack(self, x): + return tf.concat((tf.zeros((x.shape[0], 24, 4, 3), dtype=tf.float32), x), axis=3) + + def with_zeros_old(self, x): + return tf.concat((x, tf.constant([[0.0, 0.0, 0.0, 1.0]], dtype=tf.float32)), axis=0) + + def pack_old(self, x): + return tf.concat((tf.zeros((x.get_shape().as_list()[0], 4, 3), dtype=tf.float32), x), axis=2) + + + def run_rodrigues(self, pose): + batch_size = pose.shape[0] + R_out = self.rodrigues(tf.reshape(pose, (-1, 1, 3))) + R_out = tf.reshape(R_out, (batch_size, -1, 3, 3)) + #R_out = [] + #for i in range(batch_size): + # R_out.append(self.rodrigues(tf.reshape(pose[i,:], (-1, 1, 3)))) + #R_out = tf.stack(R_out, axis=0) + return R_out + + def run(self, R, betas, remove=True): + batch_size = betas.shape[0] + + #J_regressor = tf.gather(self.J_regressor, gender, axis=0) + #shapedirs = tf.gather(self.shapedirs, gender, axis=0) + #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 = tf.stack(v_shaped, axis=-1) + + def single_run(x, J_regressor=self.J_regressor, posedirs=self.posedirs, weights=self.weights, parent=self.parent): + R = x[0] + v_shaped = x[1] + g = x[2] + + J = tf.matmul(J_regressor[g], v_shaped[:, :, g]) + 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))) + v_posed = v_shaped[:, :, g] + tf.tensordot(posedirs[g], lrotmin, axes=[[2], [0]]) + 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[g], results, axes=((1), (0))) + rest_shape_h = tf.concat((v_posed, tf.ones((v_posed.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), v_posed + + gender = tf.cast(tf.greater_equal(betas[:, 10] , 0), tf.int32) + T_out, J_out, v_posed = tf.map_fn(single_run, (R, v_shaped, gender), (tf.float32, 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) + v_posed = tf.subtract(v_posed, J0) + + if remove: + T_out = tf.gather(T_out, self.to_keep, axis=1) + v_posed = tf.gather(v_posed, self.to_keep, axis=1) + + 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 = tf.stack(v_shaped, axis=-1) + + if weights is None: + w = self.weights_one_hot[np.newaxis, :, :] + if filter is not None: + w = tf.gather(tf.gather(w, self.to_keep, axis=1), filter, axis=1) + if v_cloth.shape[1] == 14475: + w = tf.gather(w, self.to_keep, axis=1) + else: + w = weights + + def single_run(x, J_regressor=self.J_regressor, w=w, parent=self.parent): + R = x[0] + v_shaped = x[1] + g = x[2] + v_cloth = x[3] + + J = tf.matmul(J_regressor[g], v_shaped[:, :, g]) + 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): + 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(w[0], results, axes=((1), (0))) + rest_shape_h = tf.concat((v_cloth, tf.ones((v_cloth.shape[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, v_shaped, gender, v_cloth), (tf.float32, tf.float32)) + + J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3)) + T_out = tf.subtract(T_out, J0) + + return T_out + + + + def get_joints_template(self, betas, gender, batch_size): + v_shaped = tf.tensordot(betas[tf.newaxis,:], self.shapedirs[gender], axes=[[1], [2]]) + tf.reshape(self.v_template[gender], shape=(1,27578,3)) + J = tf.matmul(self.J_regressor[gender], v_shaped) + J0 = tf.reshape(J[:,0,:], (-1, 1, 3)) + J_single = tf.subtract(J, J0)#[0] + J = tf.ones((batch_size, 24, 3)) * J_single + + return J + + + def run_garment(self, pose, betas, gender, v_cloth, filter=None): + batch_size = v_cloth.shape[0] + + J = self.get_joints_template(betas, gender, batch_size) + R = self.run_rodrigues(pose) + + if filter is not None: + weights = tf.boolean_mask(self.weights, filter, axis=1) + + R_cube_big = R + results = [] + results.append(self.with_zeros(tf.concat((R_cube_big[:,0,:,:], J[:, 0, :, tf.newaxis]), axis=2))) + for j in range(1, 24): + results.append(tf.matmul(results[self.parent[j]], self.with_zeros(tf.concat((R_cube_big[:,j,:,:], (J[:, j, :] - J[:, self.parent[j], :])[:,:,tf.newaxis]), axis=2)))) + stacked = tf.stack(results, axis=1) + results = stacked - self.pack(tf.matmul(stacked, tf.reshape(tf.concat((J, tf.zeros((batch_size, 24, 1), dtype=tf.float32)), axis=2), (batch_size, 24, 4, 1)))) + + J = tf.matmul(results, tf.expand_dims(tf.concat([J, tf.ones([batch_size, 24, 1], tf.float32)], axis=2), axis=2), transpose_b=True) + J = J[:, :, :3, :] + + T = tf.tensordot(weights[gender, :, :], results, axes=((1), (1))) + T = tf.transpose(T, [1,0,2,3]) + + rest_shape_h = tf.concat((v_cloth, tf.ones((batch_size, v_cloth.shape[1], 1), dtype=tf.float32)), axis=2) + v = tf.matmul(T, tf.reshape(rest_shape_h, (batch_size, -1, 4, 1))) + v = tf.reshape(v, (batch_size, -1, 4))[:, :, :3] + + J_out = tf.squeeze(J) + J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3)) + J_out = tf.subtract(J_out, J0) + T_out = tf.subtract(v, J0) + + return T_out + + def run_garment_with_weights(self, R, betas, v_cloth, weights, J=None): + + if J is None: + J = self.get_joints_template(betas) + + def single_run(x, parent=self.parent): + R = x[0] + J = x[1] + v_cloth = x[2] + weights = 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): + 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, v_cloth, weights), (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 + + def run_backward(self, R, J, v_posed, weights): + batch_size = R.get_shape()[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 = [] + J_out.append(tf.expand_dims(J[:, 0, :], axis=-1)) # batch_size x 3 x 1 + results = [] + results.append(tf.expand_dims(tf.eye(3, dtype=tf.float32), axis=0) + tf.zeros((batch_size, 3, 3), dtype=tf.float32)) + for j in range(1, 24): + results.append(tf.matmul(R[:, self.parent[j], :, :], results[self.parent[j]])) + J_out.append(tf.matmul(results[j], tf.expand_dims(J_relative[:, j-1, :], axis=-1)) + J_out[self.parent[j]]) # batch_size x 3 x 1 + results = tf.transpose(tf.stack(results, axis=0), [1, 0, 2, 3]) # (batch_size x 24 x 3 x 3) + J_out = tf.transpose(tf.stack(J_out, axis=0), [1, 0, 2, 3]) # (batch_size x 24 x 3 x 1) + + results = tf.matmul(R, results) + T = tf.transpose(tf.tensordot(results, weights, axes=((1), (1))), [0, 3, 1, 2]) # (batch_size x 6890 x 3 x 3) + v_posed = tf.expand_dims(v_posed, axis=-1) # (batch_size x 6890 x 3 x 1) + J = tf.expand_dims(J, axis=-1) # (batch_size x 24 x 3 x 1) + v = tf.matmul(T, v_posed) + tf.transpose(tf.tensordot(J_out - tf.matmul(results, J), weights, axes=((1), (1))), [0, 3, 1, 2]) # (batch_size x 6890 x 3 x 1) + v = tf.squeeze(v, -1) # (batch_size x 6890 x 3) + J_out = tf.squeeze(J_out, -1) # (batch_size x 24 x 3) + + return v, J_out \ No newline at end of file diff --git a/custom_utils/ssmpl_faces.mat b/custom_utils/ssmpl_faces.mat new file mode 100644 index 0000000..bbf7c78 Binary files /dev/null and b/custom_utils/ssmpl_faces.mat differ diff --git a/custom_utils/super_model_f.pkl b/custom_utils/super_model_f.pkl new file mode 100644 index 0000000..070c314 Binary files /dev/null and b/custom_utils/super_model_f.pkl differ diff --git a/custom_utils/super_model_m.pkl b/custom_utils/super_model_m.pkl new file mode 100644 index 0000000..bfb6bf0 Binary files /dev/null and b/custom_utils/super_model_m.pkl differ diff --git a/dataset.py b/dataset.py index 42bc964..55d4127 100644 --- a/dataset.py +++ b/dataset.py @@ -27,14 +27,23 @@ class DatasetSimple(tf.data.Dataset): mesh_mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32) uv_mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32) mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + '_' + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed'] + mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed'] + uv_unposed = np.load(os.path.join(folder, 'uv_unposed', id + '_' + frame_nr + '.npz')); uv_unposed = uv_unposed['uv_unposed']; uv_unposed = np.transpose(uv_unposed, [2, 0, 1]) # uv_posed, (pose, shape, uv_static), mesh_static, mesh_body_posed, mask - yield uv_posed, (pose_shape[:72], pose_shape[72:], uv_static, uv_mask, uv_body_posed), mesh_static, mesh_body_posed, mesh_mask, mesh_posed + yield uv_posed, (pose_shape[:72], pose_shape[72:], uv_static, uv_mask, uv_body_posed), mesh_static, mesh_body_posed, mesh_mask, mesh_posed, mesh_unposed, uv_unposed def __new__(cls, folder, res=256): return tf.data.Dataset.from_generator( cls._generator, - output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), (tf.TensorSpec(shape = (72), dtype = tf.float32), tf.TensorSpec(shape = (11), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32) ), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), tf.TensorSpec(shape = (14475,), dtype = tf.float32), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32) ), + output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), + (tf.TensorSpec(shape = (72), dtype = tf.float32), tf.TensorSpec(shape = (11), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32) ), + tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), + tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), + tf.TensorSpec(shape = (14475,), dtype = tf.float32), + tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), + tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), + tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), ), args=(folder, ) ) diff --git a/losses.py b/losses.py index 880e71a..7074949 100644 --- a/losses.py +++ b/losses.py @@ -124,7 +124,7 @@ def collision_loss(V, B, mask, thr=.004): D = V - tf.gather_nd(B, idx) b_vn = tf.gather_nd(B_vn, idx) dot = tf.einsum('abc,abc->ab', D, b_vn) - loss = tf.reduce_mean((tf.minimum(dot - thr, 0) * mask) ** 2) + loss = tf.reduce_sum((tf.minimum(dot - thr, 0) * mask) ** 2) / tf.reduce_sum(mask) return loss @@ -147,7 +147,7 @@ def edge_loss(x, y, mask): y_e = tf.sqrt(y_e + 1e-9) d_e = y_e - x_e d_e = d_e * e_mask - loss = tf.reduce_mean(d_e ** 2) + loss = tf.reduce_sum(d_e ** 2) / tf.reduce_sum(e_mask) return loss diff --git a/stylegan2/generator.py b/stylegan2/generator.py index 3ef11d3..b6e9515 100644 --- a/stylegan2/generator.py +++ b/stylegan2/generator.py @@ -102,7 +102,7 @@ class Generator(tf.keras.Model): w_broadcasted = self.truncation_trick(w_broadcasted, truncation_psi, truncation_cutoff) image_out = self.synthesis(w_broadcasted) - image_out = image_out + conditions[4] + # image_out = image_out + conditions[4] image_out = image_out * conditions[3][:,tf.newaxis,:,:] if ret_w_broadcasted: diff --git a/train.py b/train.py index cc535a3..1e3709a 100644 --- a/train.py +++ b/train.py @@ -15,10 +15,10 @@ from dataset import DatasetSimple from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss # from losses import collision_penalty as collision_penalty from losses import collision_loss as collision_penalty +from custom_utils.smpl import SMPL - -os.environ["CUDA_VISIBLE_DEVICES"] = "4" -run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='test_conditioned_23', notes='learning offsets, all losses from the start, smaller generator (from 128), slower learning rate for generator, add a fading out l1 loss with gt, add tanh activation for last layer, higher weights for the extra losses, same learning rate again. added extra metrics, trained from test-22-10000') +os.environ["CUDA_VISIBLE_DEVICES"] = "0" +run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='run_28', notes='learning offsets, all losses from the start, smaller generator (from 128), slower learning rate for generator, add a fading out l1 loss with gt, add tanh activation for last layer, higher weights for the extra losses, same learning rate again. added extra metrics, lower discr. features too. use reduce_sum for losses divided by sum of mask. learn unposed version then do skinning. lowered collsion coeff') artifact = wandb.Artifact('code', type='code') artifact.add_file('train.py') artifact.add_file('losses.py') @@ -36,6 +36,7 @@ def initiate_models(g_params, d_params, use_custom_cuda): uv2mesh = UV2Mesh(res=256) +smpl = SMPL() class Trainer(object): def __init__(self, t_params, name): @@ -50,9 +51,9 @@ class Trainer(object): self.max_steps = int(np.ceil(self.n_total_image / self.global_batch_size)) self.n_samples = min(t_params['batch_size'], t_params['n_samples']) self.train_res = t_params['train_res'] - self.print_step = 10 - self.save_step = 5000 - self.image_summary_step = 500 + self.print_step = 100 + self.save_step = 10000 + self.image_summary_step = 1000 self.unsupervised_coeff = 1 self.gt_loss_coeff = 1 self.reached_max_steps = False @@ -128,7 +129,8 @@ class Trainer(object): return params_copy def d_train_step(self, dist_inputs): - real_images, conditions, _, _, _, _ = dist_inputs[0] + uv_posed, conditions, _, _, _, _, _, uv_unposed = dist_inputs[0] + real_images = uv_unposed with tf.GradientTape() as d_tape: # compute losses @@ -144,7 +146,8 @@ class Trainer(object): return d_loss def d_train_step_reg(self, dist_inputs): - real_images, conditions, _, _, _, _ = dist_inputs[0] + uv_posed, conditions, _, _, _, _, _, uv_unposed = dist_inputs[0] + real_images = uv_unposed with tf.GradientTape() as d_tape: # compute losses @@ -165,7 +168,8 @@ class Trainer(object): def g_train_step(self, dist_inputs): - real_images, conditions, _, _, _, _ = dist_inputs[0] + uv_posed, conditions, _, _, _, _, _, uv_unposed = dist_inputs[0] + real_images = uv_unposed with tf.GradientTape() as g_tape: # compute losses @@ -179,7 +183,8 @@ class Trainer(object): return g_loss def g_train_step_reg(self, dist_inputs): - real_images, conditions, _, _, _, _ = dist_inputs[0] + uv_posed, conditions, _, _, _, _, _, uv_unposed= dist_inputs[0] + real_images = uv_unposed with tf.GradientTape() as g_tape: # compute losses @@ -201,25 +206,28 @@ class Trainer(object): def g_train_step_extra(self, dist_inputs): - real_images, conditions, mesh_static, mesh_body_posed, mesh_mask, mesh_posed = dist_inputs[0] + uv_posed, conditions, mesh_static, mesh_body_posed, mesh_mask, mesh_posed, mesh_unposed, uv_unposed = dist_inputs[0] + real_images = uv_unposed uv_mask = conditions[3] with tf.GradientTape() as g_tape: # compute losses fake_images, g_loss = g_logistic_non_saturating(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=conditions) - mesh_pred = uv2mesh(fake_images) - e_loss = 1e4 * edge_loss(mesh_pred, mesh_static, mesh_mask) - c_loss = 1e5 * collision_penalty(mesh_pred, mesh_body_posed, mesh_mask) + mesh_pred_unposed = uv2mesh(fake_images) + mesh_pred = smpl.run_garment(conditions[0], conditions[1], mesh_pred_unposed) + + e_loss = 1e4 * edge_loss(mesh_pred_unposed, mesh_static, mesh_mask) + c_loss = 1e2 * collision_penalty(mesh_pred, mesh_body_posed, mesh_mask) l1_loss = 1e2 * tf.reduce_mean(tf.reduce_sum((tf.math.abs(real_images - fake_images)), axis=1) * uv_mask) euclid_error = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mesh_mask[:,:,tf.newaxis] * (mesh_posed - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(mesh_mask) ms_ssim = tf.image.ssim_multiscale(tf.transpose(real_images, [0,2,3,1]) + 1, tf.transpose(fake_images, [0,2,3,1]) + 1, max_val=2) # scale loss g_loss = tf.reduce_sum(g_loss) * self.global_batch_scaler - e_loss = tf.reduce_sum(e_loss) * self.global_batch_scaler - c_loss = tf.reduce_sum(c_loss) * self.global_batch_scaler - l1_loss = tf.reduce_sum(l1_loss) * self.global_batch_scaler + # e_loss = tf.reduce_sum(e_loss) * self.global_batch_scaler + # c_loss = tf.reduce_sum(c_loss) * self.global_batch_scaler + # l1_loss = tf.reduce_sum(l1_loss) * self.global_batch_scaler ms_ssim = tf.reduce_sum(ms_ssim) * self.global_batch_scaler g_loss = g_loss + self.unsupervised_coeff * (e_loss + c_loss) + self.gt_loss_coeff * l1_loss @@ -230,7 +238,8 @@ class Trainer(object): def g_train_step_reg_extra(self, dist_inputs): # real_images, conditions, mesh_static, mesh_body_posed, mesh_mask, uv_mask = dist_inputs[0] - real_images, conditions, mesh_static, mesh_body_posed, mesh_mask, mesh_posed = dist_inputs[0] + uv_posed, conditions, mesh_static, mesh_body_posed, mesh_mask, mesh_posed, mesh_unposed, uv_unposed = dist_inputs[0] + real_images = uv_unposed uv_mask = conditions[3] with tf.GradientTape() as g_tape: @@ -240,20 +249,21 @@ class Trainer(object): self.pl_minibatch_shrink, self.pl_denorm, self.pl_decay, conditions=conditions) pl_penalty = pl_penalty * self.pl_weight * self.g_opt['reg_interval'] - mesh_pred = uv2mesh(fake_images) + mesh_pred_unposed = uv2mesh(fake_images) + mesh_pred = smpl.run_garment(conditions[0], conditions[1], mesh_pred_unposed) - e_loss = 1e4 * edge_loss(mesh_pred, mesh_static, mesh_mask) - c_loss = 1e5 * collision_penalty(mesh_pred, mesh_body_posed, mesh_mask) - l1_loss = 1e2 * tf.reduce_mean(tf.reduce_sum((tf.math.abs(real_images - fake_images)), axis=1) * uv_mask) + e_loss = 1e4 * edge_loss(mesh_pred_unposed, mesh_static, mesh_mask) + c_loss = 1e2 * collision_penalty(mesh_pred, mesh_body_posed, mesh_mask) + l1_loss = 1e2 * tf.reduce_sum(tf.reduce_sum((tf.math.abs(real_images - fake_images)), axis=1) * uv_mask) / tf.reduce_sum(uv_mask) euclid_error = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mesh_mask[:,:,tf.newaxis] * (mesh_posed - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(mesh_mask) ms_ssim = tf.image.ssim_multiscale(tf.transpose(real_images, [0,2,3,1]) + 1, tf.transpose(fake_images, [0,2,3,1]) + 1, max_val=2) # scale loss pl_penalty = tf.reduce_sum(pl_penalty) * self.global_batch_scaler g_gan_loss = tf.reduce_sum(g_gan_loss) * self.global_batch_scaler - e_loss = tf.reduce_sum(e_loss) * self.global_batch_scaler - c_loss = tf.reduce_sum(c_loss) * self.global_batch_scaler - l1_loss = tf.reduce_sum(l1_loss) * self.global_batch_scaler + # e_loss = tf.reduce_sum(e_loss) * self.global_batch_scaler + # c_loss = tf.reduce_sum(c_loss) * self.global_batch_scaler + # l1_loss = tf.reduce_sum(l1_loss) * self.global_batch_scaler ms_ssim = tf.reduce_sum(ms_ssim) * self.global_batch_scaler # combine @@ -296,9 +306,9 @@ class Trainer(object): def dist_g_train_step_extra(inputs): per_replica_losses = strategy.run(fn=self.g_train_step_extra, args=(inputs,)) mean_g_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None) - mean_e_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None) - mean_c_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[2], axis=None) - mean_l1_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[3], axis=None) + mean_e_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[1], axis=None) + mean_c_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[2], axis=None) + mean_l1_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[3], axis=None) mean_euclid_error = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[4], axis=None) mean_ms_ssim = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[5], axis=None) return mean_g_loss, mean_e_loss, mean_c_loss, mean_l1_loss, mean_euclid_error, mean_ms_ssim @@ -309,9 +319,9 @@ class Trainer(object): mean_g_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None) mean_g_gan_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None) mean_pl_penalty = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[2], axis=None) - mean_e_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[3], axis=None) - mean_c_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[4], axis=None) - mean_l1_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[5], axis=None) + mean_e_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[3], axis=None) + mean_c_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[4], axis=None) + mean_l1_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[5], axis=None) mean_euclid_error = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[6], axis=None) mean_ms_ssim = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[7], axis=None) return mean_g_loss, mean_g_gan_loss, mean_pl_penalty, mean_e_loss, mean_c_loss, mean_l1_loss, mean_euclid_error, mean_ms_ssim @@ -440,7 +450,8 @@ class Trainer(object): test_conditions = batch[1] #tf.ones((self.n_samples, self.g_params['labels_dim']), dtype=tf.dtypes.float32) summary_image = dist_gen_samples((test_z, test_conditions)) - mesh = uv2mesh(summary_image[:, :, :self.train_res, int(self.train_res/2):]) + mesh_unposed = uv2mesh(summary_image[:, :, :self.train_res, int(self.train_res/2):]) + mesh = smpl.run_garment(batch[1][0], batch[1][1], mesh_unposed) # convert to tensor image summary_image = self.convert_per_replica_image(summary_image, strategy) @@ -449,7 +460,7 @@ class Trainer(object): template = uv2mesh(batch[1][2]) # gt = uv2mesh(batch[0]) - plot_meshes(mesh_writer, garment=mesh, mesh_gt=batch[5], template=template, body=batch[3], mask=batch[4].numpy(), step=step) + plot_meshes(mesh_writer, garment_unposed=mesh_unposed, garment=mesh, mesh_gt=batch[6], template=template, body=batch[3], mask=batch[4].numpy(), step=step) # plot_meshes(mesh_writer, garment=gt, template=template, body=batch[3], mask=batch[4].numpy(), step=step) # train_summary_writer.flush() @@ -503,14 +514,16 @@ def filter_resolutions_featuremaps(resolutions, featuremaps, res): return filtered_resolutions, filtered_featuremaps -def plot_meshes(writer, garment, mesh_gt, template, body, mask, step, nr_offset=0): +def plot_meshes(writer, garment_unposed, garment, mesh_gt, template, body, mask, step, nr_offset=0): body = body.numpy() template = template.numpy() mesh_gt = mesh_gt.numpy() + garment_unposed = garment_unposed.numpy() F = np.load("custom_utils/smpl_faces.npy") garment, faces = maskFacesBatch(garment.numpy(), F, mask.astype(bool)) template = [template[i, mask.astype(bool)[i,:], :] for i in range(len(garment))] mesh_gt = [mesh_gt[i, mask.astype(bool)[i,:], :] for i in range(len(garment))] + garment_unposed = [garment_unposed[i, mask.astype(bool)[i,:], :] for i in range(len(garment))] collision = [np.concatenate((body[i], garment[i]), axis=0) for i in range(len(garment))] collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(len(garment))] @@ -524,7 +537,7 @@ def plot_meshes(writer, garment, mesh_gt, template, body, mask, step, nr_offset= writer.add_mesh('collision_' + str(i+nr_offset), vertices=collision[i][np.newaxis,:,:], faces=collision_faces[i][np.newaxis,:,:], global_step=step, colors=collision_color[i][np.newaxis,:,:]) for i in range(len(garment)): - plotly_fig = plot_3d(mesh_gt[i], garment[i], faces[i], collision[i], collision_faces[i], collision_color[i]) + plotly_fig = plot_3d(mesh_gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i]) wandb.log({"garment_" + str(i): plotly_fig}) def main(): @@ -553,7 +566,8 @@ def main(): # network params resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024] g_featuremaps = [128, 128, 128, 128, 128, 64, 32, 16, 8] - d_featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32] + d_featuremaps = [128, 128, 128, 128, 128, 64, 32, 16, 8] + # d_featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32] # featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32] # featuremaps = [64, 64, 64, 64, 64, 32, 16, 8, 4] train_resolutions, train_g_featuremaps = filter_resolutions_featuremaps(resolutions, g_featuremaps, args['train_res']) @@ -596,7 +610,7 @@ def main(): 'use_custom_cuda': args['use_custom_cuda'], 'model_base_dir': args['model_base_dir'], 'enable_unsupervised_losses': 0, - 'disable_l1_loss': 10000, + 'disable_l1_loss': 30000, # network params 'g_params': g_params, @@ -611,7 +625,7 @@ def main(): 'train_res': args['train_res'], } - trainer = Trainer(training_parameters, name=f'stylegan2-test-23-{args["train_res"]}x{args["train_res"]}') + trainer = Trainer(training_parameters, name=f'stylegan2-test-28-{args["train_res"]}x{args["train_res"]}') trainer.train(dist_dataset, strategy) return