added skinning

.gitignore

@@ -141,3 +141,5 @@ output/*
 models/
 models_archive/
 temp.jpg
+corrupt/
+wandb copy/

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

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

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

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, )
         )
 

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
 
 

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:

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