started writing preprocessing

.vscode/launch.json

@@ -0,0 +1,17 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Python: Current File",
+            "type": "python",
+            "request": "launch",
+            "program": "${file}",
+            "console": "integratedTerminal",
+            "justMyCode": true,
+            "cwd": "${workspaceFolder}"
+        }
+    ]
+}

custom_utils/smpl_2.py

@@ -1,346 +0,0 @@
-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, path):
-        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(path, '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(path, '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(path + '/ssmpl_faces.mat')
-        self.f = np.array(F['faces_default'])
-        self.f_skirt = np.array(F['faces_skirt'])
-        self.nn = np.array(sio.loadmat(path + '/NN.mat'))
-        self.to_keep = sio.loadmat(os.path.join(path, '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, R, betas, gender, v_cloth, filter_idx=None):
-        batch_size = v_cloth.shape[0]
-
-        J = self.get_joints_template(betas, gender, batch_size)
-        
-        if filter_idx is not None:
-            weights = tf.gather(self.weights, filter_idx, 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)
-        if J.shape[0] == 1:
-            J_out = J_out[tf.newaxis, :, :]
-        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

custom_utils/supersmpl_io.py

@@ -1,175 +0,0 @@
-import numpy as np
-from struct import pack, unpack
-
-
-"""
-Reads PC2 files, and proposed format PC16 files
-Inputs:
-- file: path to .pc2/.pc16 file
-- float16: False for PC2 files, True for PC16
-Output:
-- data: dictionary with .pc2/.pc16 file data
-NOTE: 16-bit floats lose precision with high values (positive or negative),
-	  we do not recommend using this format for data outside range [-2, 2]
-"""
-def readPC2(file, float16=False):
-	assert file.endswith('.pc2') and not float16 or file.endswith('.pc16') and float16, 'File format not consistent with specified input format'
-	data = {}
-	bytes = 2 if float16 else 4
-	dtype = np.float16 if float16 else np.float32
-	with open(file, 'rb') as f:
-		# Header
-		data['sign'] = f.read(12)
-		# data['version'] = int.from_bytes(f.read(4), 'little')
-		data['version'] = unpack('<i', f.read(4))[0]
-		# Num points
-		# data['nPoints'] = int.from_bytes(f.read(4), 'little')
-		data['nPoints'] = unpack('<i', f.read(4))[0]
-		# Start frame
-		data['startFrame'] = unpack('f', f.read(4))
-		# Sample rate
-		data['sampleRate'] = unpack('f', f.read(4))
-		# Number of samples
-		# data['nSamples'] = int.from_bytes(f.read(4), 'little')
-		data['nSamples'] = unpack('<i', f.read(4))[0]
-		# Animation data
-		size = data['nPoints']*data['nSamples']*3*bytes
-		data['V'] = np.frombuffer(f.read(size), dtype=dtype).astype(np.float32)
-		data['V'] = data['V'].reshape(data['nSamples'], data['nPoints'], 3)
-		
-	return data
-	
-"""
-Reads an specific frame of PC2/PC16 files
-Inputs:
-- file: path to .pc2/.pc16 file
-- frame: number of the frame to read
-- float16: False for PC2 files, True for PC16
-Output:
-- T: mesh vertex data at specified frame
-"""
-def readPC2Frame(file, frame, float16=False):
-	assert file.endswith('.pc2') and not float16 or file.endswith('.pc16') and float16, 'File format not consistent with specified input format'
-	assert frame >= 0 and isinstance(frame,int), 'Frame must be a positive integer'
-	bytes = 2 if float16 else 4
-	dtype = np.float16 if float16 else np.float32
-	with open(file,'rb') as f:
-		# Num points
-		f.seek(16)
-		# nPoints = int.from_bytes(f.read(4), 'little')
-		nPoints = unpack('<i', f.read(4))[0]
-		# Number of samples
-		f.seek(28)
-		# nSamples = int.from_bytes(f.read(4), 'little')
-		nSamples = unpack('<i', f.read(4))[0]
-		if frame > nSamples:
-			print("Frame index outside size")
-			print("\tN. frame: " + str(frame))
-			print("\tN. samples: " + str(nSamples))
-			return
-		# Read frame
-		size = nPoints * 3 * bytes
-		f.seek(size * frame, 1) # offset from current '1'
-		T = np.frombuffer(f.read(size), dtype=dtype).astype(np.float32)
-	return T.reshape(nPoints, 3)
-
-"""
-Writes PC2 and PC16 files
-Inputs:
-- file: path to file (overwrites if exists)
-- V: 3D animation data as a three dimensional array (N. Frames x N. Vertices x 3)
-- float16: False for writing as PC2 file, True for PC16
-This function assumes 'startFrame' to be 0 and 'sampleRate' to be 1
-NOTE: 16-bit floats lose precision with high values (positive or negative),
-	  we do not recommend using this format for data outside range [-2, 2]
-"""
-def writePC2(file, V, float16=False):
-	assert file.endswith('.pc2') and not float16 or file.endswith('.pc16') and float16, 'File format not consistent with specified input format'
-	if float16: V = V.astype(np.float16)
-	else: V = V.astype(np.float32)
-	with open(file, 'wb') as f:
-		# Create the header
-		headerFormat='<12siiffi'
-		headerStr = pack(headerFormat, b'POINTCACHE2\0',
-						1, V.shape[1], 0, 1, V.shape[0])
-		f.write(headerStr)
-		# Write vertices
-		f.write(V.tobytes())
-
-"""
-Reads OBJ files
-Only handles vertices, faces and UV maps
-Input:
-- file: path to .obj file
-Outputs:
-- V: 3D vertices
-- F: 3D faces
-- Vt: UV vertices
-- Ft: UV faces
-Correspondence between mesh and UV map is implicit in F to Ft correspondences
-If no UV map data in .obj file, it shall return Vt=None and Ft=None
-"""
-def readOBJ(file):
-	V, Vt, F, Ft = [], [], [], []
-	with open(file, 'r') as f:
-		T = f.readlines()
-	for t in T:
-		# 3D vertex
-		if t.startswith('v '):
-			v = [float(n) for n in t.replace('v ','').split(' ')]
-			V += [v]
-		# UV vertex
-		elif t.startswith('vt '):
-			v = [float(n) for n in t.replace('vt ','').split(' ')]
-			Vt += [v]
-		# Face
-		elif t.startswith('f '):
-			idx = [n.split('/') for n in t.replace('f ','').split(' ')]
-			f = [int(n[0]) - 1 for n in idx]
-			F += [f]
-			# UV face
-			if '/' in t:
-				f = [int(n[1]) - 1 for n in idx]
-				Ft += [f]
-	V = np.array(V, np.float32)
-	Vt = np.array(Vt, np.float32)
-	if Ft: assert len(F) == len(Ft), 'Inconsistent .obj file, mesh and UV map do not have the same number of faces' 
-	else: Vt, Ft = None, None
-	return V, F, Vt, Ft
-
-
-
-"""
-Writes OBJ files
-Only handles vertices, faces and UV maps
-Inputs:
-- file: path to .obj file (overwrites if exists)
-- V: 3D vertices
-- F: 3D faces
-- Vt: UV vertices
-- Ft: UV faces
-Correspondence between mesh and UV map is implicit in F to Ft correspondences
-If no UV map data as input, it will write only 3D data in .obj file
-"""
-def writeOBJ(file, V, F, Vt=None, Ft=None):
-	if not Vt is None:
-		assert len(F) == len(Ft), 'Inconsistent data, mesh and UV map do not have the same number of faces'
-		
-	with open(file, 'w') as file:
-		# Vertices
-		for v in V:
-			line = 'v ' + ' '.join([str(_) for _ in v]) + '\n'
-			file.write(line)
-		# UV verts
-		if not Vt is None:
-			for v in Vt:
-				line = 'vt ' + ' '.join([str(_) for _ in v]) + '\n'
-				file.write(line)
-		# 3D Faces / UV faces
-		if Ft:
-			F = [[str(i+1)+'/'+str(j+1) for i,j in zip(f,ft)] for f,ft in zip(F,Ft)]
-		else:
-			F = [[str(i + 1) for i in f] for f in F]		
-		for f in F:
-			line = 'f ' + ' '.join(f) + '\n'
-			file.write(line)

custom_utils/unpose.py

@@ -17,7 +17,7 @@ from tqdm import tqdm
 from datetime import datetime
 
 
-from custom_utils.supersmpl_io import readPC2, writePC2
+from data.supersmpl_io import readPC2, writePC2
 from custom_utils.smpl import SMPL
 from custom_utils.smpl_2 import SMPL as SMPL_2
 

data/cloth3d_io.py

@@ -0,0 +1,391 @@
+import numpy as np
+import scipy.io as sio
+from struct import pack, unpack
+
+
+"""
+Reads PC2 files, and proposed format PC16 files
+Inputs:
+- file: path to .pc2/.pc16 file
+- float16: False for PC2 files, True for PC16
+Output:
+- data: dictionary with .pc2/.pc16 file data
+NOTE: 16-bit floats lose precision with high values (positive or negative),
+	  we do not recommend using this format for data outside range [-2, 2]
+"""
+def readPC2(file, float16=False):
+	assert file.endswith('.pc2') and not float16 or file.endswith('.pc16') and float16, 'File format not consistent with specified input format'
+	data = {}
+	bytes = 2 if float16 else 4
+	dtype = np.float16 if float16 else np.float32
+	with open(file, 'rb') as f:
+		# Header
+		data['sign'] = f.read(12)
+		# data['version'] = int.from_bytes(f.read(4), 'little')
+		data['version'] = unpack('<i', f.read(4))[0]
+		# Num points
+		# data['nPoints'] = int.from_bytes(f.read(4), 'little')
+		data['nPoints'] = unpack('<i', f.read(4))[0]
+		# Start frame
+		data['startFrame'] = unpack('f', f.read(4))
+		# Sample rate
+		data['sampleRate'] = unpack('f', f.read(4))
+		# Number of samples
+		# data['nSamples'] = int.from_bytes(f.read(4), 'little')
+		data['nSamples'] = unpack('<i', f.read(4))[0]
+		# Animation data
+		size = data['nPoints']*data['nSamples']*3*bytes
+		data['V'] = np.frombuffer(f.read(size), dtype=dtype).astype(np.float32)
+		data['V'] = data['V'].reshape(data['nSamples'], data['nPoints'], 3)
+		
+	return data
+	
+"""
+Reads an specific frame of PC2/PC16 files
+Inputs:
+- file: path to .pc2/.pc16 file
+- frame: number of the frame to read
+- float16: False for PC2 files, True for PC16
+Output:
+- T: mesh vertex data at specified frame
+"""
+def readPC2Frame(file, frame, float16=False):
+	assert file.endswith('.pc2') and not float16 or file.endswith('.pc16') and float16, 'File format not consistent with specified input format'
+	assert frame >= 0 and isinstance(frame,int), 'Frame must be a positive integer'
+	bytes = 2 if float16 else 4
+	dtype = np.float16 if float16 else np.float32
+	with open(file,'rb') as f:
+		# Num points
+		f.seek(16)
+		# nPoints = int.from_bytes(f.read(4), 'little')
+		nPoints = unpack('<i', f.read(4))[0]
+		# Number of samples
+		f.seek(28)
+		# nSamples = int.from_bytes(f.read(4), 'little')
+		nSamples = unpack('<i', f.read(4))[0]
+		if frame > nSamples:
+			print("Frame index outside size")
+			print("\tN. frame: " + str(frame))
+			print("\tN. samples: " + str(nSamples))
+			return
+		# Read frame
+		size = nPoints * 3 * bytes
+		f.seek(size * frame, 1) # offset from current '1'
+		T = np.frombuffer(f.read(size), dtype=dtype).astype(np.float32)
+	return T.reshape(nPoints, 3)
+
+"""
+Writes PC2 and PC16 files
+Inputs:
+- file: path to file (overwrites if exists)
+- V: 3D animation data as a three dimensional array (N. Frames x N. Vertices x 3)
+- float16: False for writing as PC2 file, True for PC16
+This function assumes 'startFrame' to be 0 and 'sampleRate' to be 1
+NOTE: 16-bit floats lose precision with high values (positive or negative),
+	  we do not recommend using this format for data outside range [-2, 2]
+"""
+def writePC2(file, V, float16=False):
+	assert file.endswith('.pc2') and not float16 or file.endswith('.pc16') and float16, 'File format not consistent with specified input format'
+	if float16: V = V.astype(np.float16)
+	else: V = V.astype(np.float32)
+	with open(file, 'wb') as f:
+		# Create the header
+		headerFormat='<12siiffi'
+		headerStr = pack(headerFormat, b'POINTCACHE2\0',
+						1, V.shape[1], 0, 1, V.shape[0])
+		f.write(headerStr)
+		# Write vertices
+		f.write(V.tobytes())
+
+"""
+Reads OBJ files
+Only handles vertices, faces and UV maps
+Input:
+- file: path to .obj file
+Outputs:
+- V: 3D vertices
+- F: 3D faces
+- Vt: UV vertices
+- Ft: UV faces
+Correspondence between mesh and UV map is implicit in F to Ft correspondences
+If no UV map data in .obj file, it shall return Vt=None and Ft=None
+"""
+def readOBJ(file):
+	V, Vt, F, Ft = [], [], [], []
+	with open(file, 'r') as f:
+		T = f.readlines()
+	for t in T:
+		# 3D vertex
+		if t.startswith('v '):
+			v = [float(n) for n in t.replace('v ','').split(' ')]
+			V += [v]
+		# UV vertex
+		elif t.startswith('vt '):
+			v = [float(n) for n in t.replace('vt ','').split(' ')]
+			Vt += [v]
+		# Face
+		elif t.startswith('f '):
+			idx = [n.split('/') for n in t.replace('f ','').split(' ')]
+			f = [int(n[0]) - 1 for n in idx]
+			F += [f]
+			# UV face
+			if '/' in t:
+				f = [int(n[1]) - 1 for n in idx]
+				Ft += [f]
+	V = np.array(V, np.float32)
+	Vt = np.array(Vt, np.float32)
+	if Ft: assert len(F) == len(Ft), 'Inconsistent .obj file, mesh and UV map do not have the same number of faces' 
+	else: Vt, Ft = None, None
+	return V, F, Vt, Ft
+
+
+
+"""
+Writes OBJ files
+Only handles vertices, faces and UV maps
+Inputs:
+- file: path to .obj file (overwrites if exists)
+- V: 3D vertices
+- F: 3D faces
+- Vt: UV vertices
+- Ft: UV faces
+Correspondence between mesh and UV map is implicit in F to Ft correspondences
+If no UV map data as input, it will write only 3D data in .obj file
+"""
+def writeOBJ(file, V, F, Vt=None, Ft=None):
+	if not Vt is None:
+		assert len(F) == len(Ft), 'Inconsistent data, mesh and UV map do not have the same number of faces'
+		
+	with open(file, 'w') as file:
+		# Vertices
+		for v in V:
+			line = 'v ' + ' '.join([str(_) for _ in v]) + '\n'
+			file.write(line)
+		# UV verts
+		if not Vt is None:
+			for v in Vt:
+				line = 'vt ' + ' '.join([str(_) for _ in v]) + '\n'
+				file.write(line)
+		# 3D Faces / UV faces
+		if Ft:
+			F = [[str(i+1)+'/'+str(j+1) for i,j in zip(f,ft)] for f,ft in zip(F,Ft)]
+		else:
+			F = [[str(i + 1) for i in f] for f in F]		
+		for f in F:
+			line = 'f ' + ' '.join(f) + '\n'
+			file.write(line)
+
+
+def loadInfo(filename):
+	'''
+	this function should be called instead of direct sio.loadmat
+	as it cures the problem of not properly recovering python dictionaries
+	from mat files. It calls the function check keys to cure all entries
+	which are still mat-objects
+	'''
+	data = sio.loadmat(filename, struct_as_record=False, squeeze_me=True)
+	del data['__globals__']
+	del data['__header__']
+	del data['__version__']
+	return _check_keys(data)
+
+def _check_keys(dict):
+	'''
+	checks if entries in dictionary are mat-objects. If yes
+	todict is called to change them to nested dictionaries
+	'''
+	for key in dict:
+		if isinstance(dict[key], sio.matlab.mio5_params.mat_struct):
+			dict[key] = _todict(dict[key])
+	return dict        
+
+def _todict(matobj):
+	'''
+	A recursive function which constructs from matobjects nested dictionaries
+	'''
+	dict = {}
+	for strg in matobj._fieldnames:
+		elem = matobj.__dict__[strg]
+		if isinstance(elem, sio.matlab.mio5_params.mat_struct):
+			dict[strg] = _todict(elem)
+		elif isinstance(elem, np.ndarray) and np.any([isinstance(item, sio.matlab.mio5_params.mat_struct) for item in elem]):
+			dict[strg] = [None] * len(elem)
+			for i,item in enumerate(elem):
+				if isinstance(item, sio.matlab.mio5_params.mat_struct):
+					dict[strg][i] = _todict(item)
+				else:
+					dict[strg][i] = item
+		else:
+			dict[strg] = elem
+	return dict
+
+
+# import os
+# import sys
+# import numpy as np
+# from PIL import Image
+
+# sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+# from cloth3d.smpl_np import SMPLModel
+# from cloth3d.util import loadInfo, zRotMatrix, proj, mesh2UV, uv_to_pixel
+# from superSMPL.IO import readOBJ, readPC2Frame
+
+# class DataReader:
+
+#     def __init__(self, folder, smpl_path):
+#         # Data paths
+#         self.SRC = folder
+#         # SMPL model
+#         smpl_path = smpl_path
+#         self.smpl = {
+#             'f': SMPLModel(os.path.join(smpl_path, 'super_model_f.pkl')),
+#             'm': SMPLModel(os.path.join(smpl_path, 'super_model_m.pkl'))
+#         }
+        
+#     """ 
+#     Read sample info 
+#     Input:
+#     - sample: name of the sample e.g.:'01_01_s0'
+#     """
+#     def read_info(self, sample):
+#         info_path = os.path.join(self.SRC, sample, 'info')
+#         return loadInfo(info_path)
+        
+#     """ Human data """
+#     """
+#     Read SMPL parameters for the specified sample and frame
+#     Inputs:
+#     - sample: name of the sample
+#     - frame: frame number
+#     """
+#     def read_smpl_params(self, sample, frame):
+#         # Read sample data
+#         info = self.read_info(sample)
+#         # SMPL parameters
+#         gender = 'm' if info['gender'] else 'f'
+#         pose = info['poses'][:, frame].reshape(self.smpl[gender].pose_shape)
+#         shape = info['shape']
+#         trans = info['trans'][:, frame].reshape(self.smpl[gender].trans_shape)
+#         return gender, pose, shape, trans
+    
+#     """
+#     Computes human mesh for the specified sample and frame
+#     Inputs:
+#     - sample: name of the sample
+#     - frame: frame number
+#     - absolute: True for absolute vertex locations, False for locations relative to SMPL root joint	
+#     Outputs:
+#     - V: human mesh vertices
+#     - F: mesh faces
+#     """
+#     def read_human(self, sample, frame, zrot, use_trans=True):
+#         # SMPL parameters
+#         gender, pose, shape, trans = self.read_smpl_params(sample, frame)
+#         # Compute SMPL
+#         V, J = self.smpl[gender].set_params(pose=pose, beta=shape, trans=trans if use_trans else None)
+#         V -= J[0:1]
+#         # Apply rotation on z-axis
+#         zRot = zRotMatrix(zrot)
+#         return zRot.dot(V.T).T, self.smpl[gender].faces, zRot.dot(trans.T).T[None,:], zRot.dot((J -  J[0:1] + trans[None,:]).T).T[None,:], zRot.dot((J -  J[0:1]).T).T[None,:]
+    
+#     """ Garment data """
+#     """
+#     Reads garment vertices location for the specified sample, garment and frame
+#     Inputs:
+#     - sample: name of the sample
+#     - garment: type of garment (e.g.: 'Tshirt', 'Jumpsuit', ...)
+#     - frame: frame number
+#     - TODO: absolute: True for absolute vertex locations, False for locations relative to SMPL root joint	
+#     Outputs:
+#     - V: 3D vertex locations for the specified sample, garment and frame
+#     """
+#     def read_garment_vertices(self, sample, garment, frame, trans, zrot):
+#         # Read garment vertices (relative to root joint)
+#         pc16_path = os.path.join(self.SRC, sample, garment + '.pc16')
+#         V = readPC2Frame(pc16_path, frame, True)
+#         # Read sample data
+#         if trans is not None:
+#             # Transform to absolute
+#             V += trans[None]
+#         zRot = zRotMatrix(zrot)
+#         return zRot.dot(V.T).T
+
+#     """
+#     Reads garment faces for the specified sample and garment
+#     Inputs:
+#     - sample: name of the sample
+#     - garment: type of garment (e.g.: 'Tshirt', 'Jumpsuit', ...)
+#     Outputs:
+#     - F: mesh faces	
+#     """
+#     def read_garment_topology(self, sample, garment):
+#         # Read OBJ file
+#         obj_path = os.path.join(self.SRC, sample, garment + '.obj')
+#         return readOBJ(obj_path)[1]
+
+#     """	
+#     Reads garment UV map for the specified sample and garment
+#     Inputs:
+#     - sample: name of the sample
+#     - garment: type of garment (e.g.: 'Tshirt', 'Jumpsuit', ...)
+#     Outputs:
+#     - Vt: UV map vertices
+#     - Ft: UV map faces		
+#     """
+#     def read_garment_UVMap(self, sample, garment):
+#         # Read OBJ file
+#         obj_path = os.path.join(self.SRC, sample, garment + '.obj')
+#         return readOBJ(obj_path)[2:]		
+
+#     """
+#     Reads vertex colors of the specified sample and garment
+#     Inputs:
+#     - sample: name of the sample
+#     - garment: type of garment (e.g.: 'Tshirt', 'Jumpsuit', ...)
+#     - F: mesh faces
+#     - Vt: UV map vertices
+#     - Ft: UV map faces
+#     Output
+#     - C: RGB colors
+#     """
+#     def read_garment_vertex_colors(self, sample, garment, F, Vt, Ft):
+#         # Read sample texture info
+#         texture = self.read_info(sample)['outfit'][garment]['texture']
+#         # Plain color
+#         if texture['type'] == 'color': return (255*texture['data']).astype(np.int32)
+#         # Image texture
+#         img_path = os.path.join(self.SRC, sample, garment + '.png')
+#         img = Image.open(img_path)
+#         # Get color of each UV vertex
+#         color = np.array([img.getpixel(uv_to_pixel(vt)) for vt in Vt])
+#         # Compute correspondece between V (mesh) and Vt (UV map)
+#         m2uv = mesh2UV(F, Ft)
+#         return np.array([color[list(m2uv[idx])].mean(0) for idx in sorted(list(m2uv.keys()))], np.uint8)
+        
+#     """ Scene data """
+#     """
+#     Read camera location and compute projection matrix
+#     Input:
+#     - sample: name of the sample
+#     Output:
+#     - P: camera projection matrix (3 x 4)
+#     """
+#     def read_camera(self, sample):
+#         # Read sample data
+#         info = self.read_info(sample)
+#         # Camera location
+#         camLoc = info['camLoc']
+#         # Camera projection matrix
+#         return proj(camLoc)
+        
+# # TESTING
+# if __name__ == '__main__':
+#     sample = '135_02_s8'
+#     frame = 0
+#     garment = 'Tshirt'
+    
+#     reader = DataReader()
+#     F = reader.read_garment_topology(sample, garment)
+#     Vt, Ft = reader.read_garment_UVMap(sample, garment)
+#     C = reader.read_garment_vertex_colors(sample, garment, F, Vt, Ft)

data/preprocessing.py

@@ -0,0 +1,61 @@
+import os
+import sys
+
+import shutil
+import numpy as np
+import scipy.io as sio
+
+from tqdm import tqdm
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+from data.cloth3d_io import readPC2, loadInfo
+from data.utils import Frame
+
+
+
+
+def clean_dataset(folder):
+    ids = os.listdir(folder)
+
+    for id in tqdm(ids):
+
+        if not os.path.isdir(os.path.join(folder, id)):
+            continue
+
+        files = os.listdir(os.path.join(folder, id))
+
+        for file in files:
+            if file.startswith('glmarks') or file.find('nricp') != -1:
+                os.remove(os.path.join(folder, id, file))
+
+ 
+def generate_frames_file(folder, file_name):
+    ids = os.listdir(folder)
+    frames = []
+
+    for id in tqdm(ids):
+
+        if not os.path.isdir(os.path.join(folder, id)):
+            continue
+
+        garments = [f.replace('.obj','') for f in os.listdir(os.path.join(folder, id)) if f.endswith('.obj')]
+        info = loadInfo(os.path.join(folder, id, 'info.mat'))
+
+        for garment in garments:
+            for i in range(info['poses'].shape[1]):
+                frames.append(Frame(id, i, garment))
+
+    with open(os.path.join(folder, file_name), 'w') as f:
+        for i, frame in enumerate(frames):
+            f.write(str(frame.id) + ' ' + str(frame.frame_nr) + ' ' + str(frame.garment) + '\n')
+
+
+
+
+
+if __name__ == '__main__':
+
+    folder = '/home/workspace/data/dataset/cloth3d/train'
+
+    # clean_dataset(folder)
+    generate_frames_file(folder, 'frames_all.txt')

data/processing/make_super.py

@@ -0,0 +1,112 @@
+import sys
+
+import numpy as np
+import pickle as pkl
+from scipy import sparse
+
+from time import time
+
+from data.supersmpl_io import writeOBJ, readOBJ
+
+
+def makeGraph(F):
+	G = {}
+	for f in F:
+		N = len(f)
+		for i,v in enumerate(f):
+			if not v in G: G[v] = {}
+			G[v][f[(i+1) % N]] = 1
+	return G
+
+def extrapolateMissings(M, G):
+	for i in range(6890, M.shape[0]):
+		if np.isnan(M[i]).all():
+			M[i] = neighAvg(M, G, i)
+		elif np.isnan(M[i]).any():
+			print("HOW?")
+			print(M[i])
+	if np.isnan(M).any(): extrapolateMissings(M, G)
+
+def neighAvg(M, G, i):
+	neigh = set(G[i].keys())
+	avg = 0
+	count = 0
+	for n in neigh:
+		if not np.isnan(M[n]).any():
+			avg += M[n]
+			count += 1
+	if count: return avg / count
+	return np.nan
+
+
+if __name__ == '__main__':
+
+	g = 'f'
+	# Read Super topology
+	V, F, Vt, Ft = readOBJ('super_f_avg.obj')
+	N = V.shape[0]
+	G = makeGraph(F)
+
+	# Read PKL
+	print("Loading .pkl file")
+	with open('custom_data/model_' + g + '.pkl', 'rb') as f:
+		M = pkl.load(f, encoding='latin1')
+		
+	# Fields
+	# print("Checking fields and shapes")
+	# for k in M:
+	# 	print(k, M[k].shape)
+	# Faces
+	print("New faces...")
+	M['f'] = F
+	# J_regressor (24, 6890)
+	# Sparse matrix, just change shape to add zeros
+	print("New regressor")
+	A = np.zeros((24, N), np.float32)
+	A[:,:6890] = M['J_regressor'].todense()
+	M['J_regressor'] = sparse.csr_matrix(A)
+
+	# Weights (6890, 24)
+	print("New weights")
+	start = time()
+	A = np.empty((N, 24), np.float32)
+	A[:] = np.nan
+	A[:6890] = M['weights'][:6890]
+	extrapolateMissings(A, G)
+	M['weights'] = A
+	print(time() - start)
+
+	# Posedirs (6890, 3, 207)
+	print("New posedirs")
+	start = time()
+	A = np.empty((N, 3, 207), np.float32)
+	A[:] = np.nan
+	A[:6890] = M['posedirs'][:6890]
+	extrapolateMissings(A, G)
+	M['posedirs'] = A
+	print(time() - start)
+
+	# v_tempalte (6890, 3)
+	print("New template")
+	start = time()
+	A = np.empty((N, 3), np.float32)
+	A[:] = np.nan
+	A[:6890] = M['v_template'][:6890]
+	extrapolateMissings(A, G)
+	M['v_template'] = A
+	print(time() - start)
+
+	# Shapedirs (6890, 3, 10)
+	print("New shapedirs")
+	start = time()
+	A = np.empty((N, 3, 10), np.float32)
+	A[:] = np.nan
+	A[:6890] = M['shapedirs'][:6890]
+	extrapolateMissings(A, G)
+	M['shapedirs'] = A
+	print(time() - start)
+
+	# Dump
+	print("Saving new model")
+	with open('super_model_' + g + '.pkl', 'wb') as f:
+		pkl.dump(M, f)

data/processing/preprocessing_old.py

@@ -0,0 +1,310 @@
+import glob
+import os
+
+from data.supersmpl_io import readPC2
+os.environ["CUDA_VISIBLE_DEVICES"] = "3"
+
+import sys
+import shutil
+import tqdm
+import random
+import math
+import numpy as np
+import scipy as sp
+import scipy.io as sio
+
+from custom_utils.smpl import SMPL
+import plotly.graph_objects as go
+from custom_utils.custom_utils import _maskFaces
+
+
+def create_no_orient_body(folder):
+    ids = os.listdir(folder)
+    F = np.load("custom_utils/smpl_faces.npy")
+
+    smpl = SMPL()
+    for id in tqdm.tqdm(ids):
+        if not os.path.isdir(os.path.join(folder, id)):
+            continue
+
+        # if os.path.exists(os.path.join(folder, id, 'mesh_posed_body_no_orient')):
+        #     continue
+
+        info = sio.loadmat(os.path.join(folder, id, 'info.mat'))
+
+        pose, shape, gender = info['poses'], info['shape'], info['gender']
+
+        # if gender == 1:
+        #     continue 
+
+        pose = np.transpose(pose, [1, 0]).astype(np.float32)
+
+        if len(os.listdir(os.path.join(folder, id, 'mesh_posed_body_no_orient'))) == pose.shape[0]:
+            continue
+
+        pose[:, :3] = 0
+        beta = np.squeeze(np.concatenate((shape, gender), axis=1)).astype(np.float32)
+        R = smpl.run_rodrigues(pose)
+        T, _, _ = smpl.run(R, np.tile(beta[np.newaxis, :], [pose.shape[0], 1]), True)
+        T = T.numpy()
+
+
+        os.makedirs(os.path.join(folder, id, 'mesh_posed_body_no_orient'), exist_ok=True)
+        for i in range(pose.shape[0]):
+            np.savez_compressed(os.path.join(folder, id, 'mesh_posed_body_no_orient', id + '_' + str(i) + '.npz'), mesh_posed_body_no_orient = T[i, :, :])
+
+
+
+def create_no_shape_unposed(folder):
+    import tensorflow as tf
+    ids = os.listdir(folder)
+
+    smpl = SMPL()
+    for id in tqdm.tqdm(ids):
+        garment_dir = os.path.join(folder, id)
+        if not os.path.isdir(garment_dir):
+            continue
+        # if os.path.exists(os.path.join(folder, id, 'mesh_unposed_no_shape')):
+        #     continue
+
+
+        info = sio.loadmat(os.path.join(folder, id, 'info.mat'))
+
+        # if len(os.listdir(os.path.join(folder, id, 'mesh_unposed_no_shape_one_hot'))) == info['poses'].shape[1]:
+        #     continue
+
+        shape = info['shape']
+        gender = int(info['gender'][0,0] > 0)
+        garments = [f.replace('.pc16','') for f in os.listdir(garment_dir) if f.find('_') == -1 and f != 'info.mat' and not f.endswith('.obj')]
+        assert len (garments) == 1
+        assert gender == info['gender'][0,0]
+        # if len (garments) != 1:
+        #     continue
+        garment = garments[0]
+        mask = np.load(os.path.join(folder, id, garment + '_mask.npy'))
+        shape_offsets = tf.tensordot(shape.astype(np.float32), smpl.shapedirs[gender], axes=[[1], [2]]).numpy()
+
+        os.makedirs(os.path.join(folder, id, 'mesh_unposed_no_shape_one_hot'), exist_ok=True)
+        for i in range(info['poses'].shape[1]):
+            mesh_unposed = np.load(os.path.join(garment_dir, 'mesh_unposed_one_hot', id + '_' + str(i) + '.npz'))
+            mesh_unposed = mesh_unposed['mesh_unposed']
+            mesh_unposed[mask, :] = mesh_unposed[mask, :] - np.squeeze(shape_offsets)[smpl.to_keep, :][mask, :]
+            np.savez_compressed(os.path.join(folder, id, 'mesh_unposed_no_shape_one_hot', id + '_' + str(i) + '.npz'), mesh_unposed_no_shape = mesh_unposed)
+
+        
+
+
+def check_if_skirt(folder):
+    ids = os.listdir(folder)
+
+    for id in tqdm.tqdm(ids):
+        garment_dir = os.path.join(folder, id)
+        if not os.path.isdir(garment_dir):
+            continue
+
+        garments = [f.replace('.pc16','') for f in os.listdir(garment_dir) if f.find('_') == -1 and f != 'info.mat' and not f.endswith('.obj')]
+        # assert len (garments) == 1
+        # garment = garments[0]
+
+        # assert garment != 'Skirt'
+
+        assert "Skirt" not in garments
+
+def gather_ids(folder, file_name):
+    ids = os.listdir(folder)
+
+    to_print = []
+
+    for id in tqdm.tqdm(ids):
+        garment_dir = os.path.join(folder, id)
+        if not os.path.isdir(garment_dir):
+            continue
+
+        if len(id) == 5:
+            to_print.append(id)
+
+    to_print.sort()
+
+    with open(file_name, 'w') as f:
+        f.write('\n'.join(to_print))
+
+
+def verify_poses(folder):
+    ids = os.listdir(folder)
+
+    for id in tqdm.tqdm(ids):
+        if not os.path.isdir(os.path.join(folder, id)):
+            continue
+
+        garments = [f.replace('_registered.pc16','') for f in os.listdir(os.path.join(folder, id)) if f.endswith('_registered.pc16') and not f.startswith('Combined')]
+        assert len (garments) == 1
+
+        info = sio.loadmat(os.path.join(folder, id, 'info.mat'))
+        poses = info['poses']
+        verts = readPC2(os.path.join(folder, id, garments[0] + '_registered.pc16'), True)['V']
+        # print(id)
+        # assert type(poses) == np.ndarray and poses.shape[0] == 72 and poses.shape[1] == verts.shape[0], 'mv ' + os.path.join(folder, id) + ' /home/workspace/data/dataset/invalid/'
+        if not (type(poses) == np.ndarray and poses.shape[0] == 72 and poses.shape[1] == verts.shape[0]):
+            shutil.move(os.path.join(folder, id), '/home/workspace/data/dataset/invalid/')
+            print(id)
+
+
+def fix_split_garment_generated_data(folder):
+    ids = os.listdir(folder)
+
+    for id in tqdm.tqdm(ids):
+        if not os.path.isdir(os.path.join(folder, id)):
+            continue
+
+        garments = [f.replace('_registered.pc16','') for f in os.listdir(os.path.join(folder, id)) if f.endswith('_registered.pc16') and not f.startswith('Combined')]
+        assert len (garments) == 1
+
+        if len(id) == 6 and id.endswith('2'):
+            frames_folder = os.path.join(folder, id, 'mesh_posed_body_no_orient') # os.path.join(folder, id, 'mesh_unposed_no_shape')
+            frames = os.listdir(frames_folder)
+            frame_id = frames[0].split('_')[0]
+            if len(frame_id) == 6:
+                assert frame_id == id, 'Wrong id found!'
+            if len(frame_id) == 5 and id.startswith(frame_id):
+                for frame in frames:
+                    frame_nr = frame.split('.')[0].split('_')[1]
+                    shutil.move(os.path.join(frames_folder, frame), os.path.join(frames_folder, id + '_' + frame_nr + '.npz'))
+
+
+def create_garment_type_table():
+    folder = '/home/workspace/data/dataset'
+    # sub_folders = ['train', 'validation', 'test', 'train_skirt', 'valid_skirt', 'test_skirt']
+    # sub_folders = ['train', 'validation', 'train_skirt', 'valid_skirt']
+    sub_folders = ['test_skirt']
+    # sub_folders = ['test', 'test_skirt']
+
+    garment_type_table = dict()
+
+    for sub_f in sub_folders:
+        ids = os.listdir(os.path.join(folder, sub_f))
+
+        for id in tqdm.tqdm(ids):
+            if not os.path.isdir(os.path.join(folder, sub_f, id)):
+                continue
+
+            garments = [f.replace('_registered.pc16','') for f in os.listdir(os.path.join(folder, sub_f, id)) if f.endswith('_registered.pc16') and not f.startswith('Combined')]
+            assert len (garments) == 1
+
+            if id in garment_type_table.keys():
+                print(id + ': ' + garments[0] + ' ' + garment_type_table[id])
+
+            garment_type_table[id] = garments[0]
+
+    np.save('custom_data/garment_type_table_test.npy', garment_type_table)
+
+
+def generate_faces():
+    folder = '/home/workspace/data/dataset/test_skirt_ready'
+    output_folder = '/home/workspace/data/results/support/faces_skirt'
+    F = np.load("custom_utils/smpl_faces_skirt.npy")
+
+    os.makedirs(output_folder, exist_ok=True)
+
+    ids = os.listdir(os.path.join(folder, 'pose_shape'))
+    ids = [id.split('_')[0] for id in ids]
+    ids = set(ids)
+
+    for id in tqdm.tqdm(ids):
+        mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy')).astype(bool)
+
+        _, faces = _maskFaces(np.zeros((14475, 3), np.float32), F, mask)
+        np.save(os.path.join(output_folder, id + '.npy'), faces.astype(np.int32))
+
+
+def generate_frames_file(folder, ratio=0.1):
+    all_frames = os.listdir(os.path.join(folder, 'pose_shape'))
+    ids = [frame.split('_')[0] for frame in all_frames]
+    all_frames = [frame.split('.')[0] for frame in all_frames]
+    ids = sorted(list(set(ids)))
+
+    garment_type_table = np.load('custom_data/garment_type_table_train.npy', allow_pickle=True).item()
+
+
+    random.shuffle(ids)
+
+    sub_ids = ids[:int(len(ids) * ratio)]
+
+    selected = []
+
+    for id in tqdm.tqdm(sub_ids):
+        if garment_type_table[id] != 'Dress':
+            continue
+        if len(selected) > 30000:
+            break
+
+        # frames = glob.glob(os.path.join(folder, 'pose_shape', id + '_*'))
+        frames = [frame for frame in all_frames if frame.startswith(id + '_')]
+        frames = [frame.split('/')[-1].split('.')[0] for frame in frames]
+        selected.extend(frames)
+
+    with open(os.path.join(folder, 'frames_dress_30k.txt'), 'w') as f:
+        for i, frame in enumerate(selected):
+            f.write(frame + '\n')
+
+
+def generate_frames_file_filtered_cape():
+    folder = '/home/workspace/local-data/cape_processed/train/female'
+
+    all_frames = os.listdir(os.path.join(folder, 'pose_shape'))
+    ids = [frame.split('.')[0] for frame in all_frames]
+    all_frames = ['.'.join(frame.split('.')[:-1]) for frame in all_frames]
+    ids = sorted(list(set(ids)))
+
+    seqs_train_00032_SS = ['00032_shortshort_hips', '00032_shortshort_shoulders_mill']
+    seqs_test_00032_SS = [ '00032_shortshort_tilt_twist_left']
+    
+    seqs_train_00032_SL = ['00032_shortlong_hips', '00032_shortlong_shoulders_mill']
+    seqs_test_00032_SL = [ '00032_shortlong_tilt_twist_left']
+
+    seqs_train_00032_SL_2 = ['00032_shortlong_hips', '00032_shortlong_tilt_twist_left']
+    seqs_test_00032_SL_2 = ['00032_shortlong_shoulders_mill']
+
+    seqs_train_03223_SL_2 = ['03223_shortlong_hips', '03223_shortlong_tilt_twist_left']
+    seqs_test_03223_SL_2 = ['03223_shortlong_shoulders_mill']
+
+    seqs_train_03223_SS_2 = ['03223_shortshort_hips', '03223_shortshort_tilt_twist_left']
+    seqs_test_03223_SS_2 = ['03223_shortshort_shoulders_mill']
+
+    ids = [id for id in ids if id in seqs_train_03223_SS_2]
+
+    random.shuffle(ids)
+
+    sub_ids = ids
+
+    selected = []
+
+    for id in tqdm.tqdm(sub_ids):
+        # frames = glob.glob(os.path.join(folder, 'pose_shape', id + '_*'))
+        frames = [frame for frame in all_frames if frame.startswith(id + '.')]
+        # frames = [frame.split('/')[-1].split('.')[0] for frame in frames]
+        selected.extend(frames)
+
+    random.shuffle(selected)
+
+    with open(os.path.join(folder, 'frames_03223_SS_train_2.txt'), 'w') as f:
+        for i, frame in enumerate(selected):
+            f.write(frame + '\n')
+
+
+if __name__ == '__main__':
+
+    src_folder = '/home/workspace/local-data/train_skirt' #cloth3d_skirt
+    # src_folder = r'D:\Work\Dataset\validation_2'
+    dest_folder = '/home/workspace/data/dataset/validation' 
+
+
+    # create_no_orient_body(src_folder)
+    # fix_split_garment_generated_data(src_folder)
+    # create_no_shape_unposed(src_folder)
+    # check_if_skirt(src_folder)
+    # gather_ids(src_folder, 'ids/test_non_skirt_ids.txt')
+    # verify_poses(src_folder)
+    # create_garment_type_table()
+    # generate_faces()
+    generate_frames_file(src_folder, ratio=1)
+    # generate_frames_file_filtered_cape()

data/utils.py

@@ -0,0 +1,14 @@
+
+
+class Frame:
+    def __init__(self, id, frame_nr, garment):
+        self.id = id
+        self.frame_nr = frame_nr
+        self.garment = garment
+
+    def __hash__(self):
+        return hash((self.id, self.frame_nr))
+
+    def __eq__(self, other):
+        return (self.id, self.frame_nr) == (other.id, other.frame_nr)
+