added skinning and gnn discriminator

custom_utils/graph_utils.py

@@ -0,0 +1,207 @@
+import scipy as sp
+import numpy as np
+import scipy.io as sio
+import tensorflow as tf
+from tensorflow.keras.activations import tanh, relu
+from tensorflow.keras.layers import Dense, BatchNormalization, Add, Flatten
+from custom_utils.custom_utils import UV2Mesh
+
+from tensorflow_graphics.geometry.convolution.graph_pooling import pool, unpool
+from tensorflow_graphics.geometry.convolution.graph_convolution import edge_convolution_template
+
+
+def sparse_to_tensor(L):
+    idx = np.mat([L.row, L.col]).transpose().astype(np.float32)
+    data = L.data.astype(dtype=np.float32)
+    shape = L.shape
+    return tf.SparseTensor(idx, data, shape)
+
+def merge_verts(Vs, batchDim):
+    V = tf.concat(tf.unstack(Vs, num=batchDim), axis=0)
+    return V
+
+
+def repeat_laplacian(L, times):
+    idx = np.repeat(np.mat([L.row, L.col]), repeats=times, axis=1)
+    idx_offset = np.linspace(0, L.shape[0]  * (times-1), times).astype(np.int32)
+    idx_offset = np.repeat(idx_offset, L.row.shape[0], axis=0)
+    idx += idx_offset[np.newaxis,:]
+    data = np.repeat(L.data, repeats=times, axis=0)
+    shape = (L.shape[0] * times, L.shape[1] * times)
+    return sp.sparse.coo_matrix((data, (np.squeeze(np.asarray(idx[0,:])), np.squeeze(np.asarray(idx[1,:])))), shape)
+
+
+
+class GraphConvolution(tf.keras.layers.Layer):
+    def __init__(self, shape, activation=None, name='gcn', L=None, batchSize=None, normalization=None):
+        super(GraphConvolution, self).__init__()
+        self.shape = shape
+        self.w = self.add_weight(name=name + '/weight', shape=(2, *shape), initializer=tf.initializers.glorot_normal, trainable=True)
+        self.b = self.add_weight(name=name + '/bias', shape=shape[-1], initializer=tf.initializers.zeros, trainable=True)
+        self.activation = activation or (lambda x: x)
+        self.normalization = normalization
+        self.batchSize = batchSize
+        self.indices = tf.cast(tf.linspace(0, self.batchSize * L.shape[1], self.batchSize+1), dtype=tf.int32)
+        self.L = repeat_laplacian(L, self.batchSize)
+
+    def gather(self):
+        return [self.w, self.b]
+
+    def _single(self, X):
+        X0 = tf.einsum('ab,bc->ac', X, self.w[0]) # Node
+        X1 = tf.einsum('ab,bc->ac', X, self.w[1]) # Neigh
+        # Neigh conv
+        X1 = edge_convolution_template(
+            X1, sparse_to_tensor(self.L), sizes=None,
+            edge_function=lambda x, y: y, 
+            reduction='weighted', 
+            edge_function_kwargs={}
+        ) 
+        X = X0 + X1 + self.b
+        return X
+
+    def call(self, X, training):
+        X = merge_verts(X, self.batchSize)
+        X = self._single(X)
+        X = tf.stack([X[self.indices[i-1]:self.indices[i]] for i in range(1, len(self.indices))])
+        X = self.activation(X)
+        if self.normalization:
+            X = self.normalization(X, training=training)
+        return X
+
+
+class GraphUnPool(tf.keras.layers.Layer):
+    def __init__(self, pool_map):
+        super(GraphUnPool, self).__init__()
+        self.pool_map = pool_map
+
+    def call(self, X):
+        # print(X.dtype)
+        # print(self.pool_map.dtype)
+        return unpool(X, self.pool_map, None)
+        # return tf.gather(X, self.pool_map, axis=1)
+
+
+
+class GraphPool(tf.keras.layers.Layer):
+    def __init__(self, pool_map, type, batch_size):
+        super(GraphPool, self).__init__()
+        self.type = type
+        self.batch_size = batch_size
+        self.indices = tf.cast(tf.linspace(0, batch_size * pool_map.shape[0], batch_size+1), dtype=tf.int32)
+        # self.pool_map = repeat_laplacian(pool_map, self.batch_size)
+        self.pool_map = np.squeeze(pool_map)
+
+    def _single(self, x):
+        return pool(x, sparse_to_tensor(self.pool_map), algorithm=self.type, sizes=None)
+
+    def call(self, x, training=None):
+        # x = merge_verts(x, self.batch_size)
+        # x = self._single(x)
+        # x = tf.stack([x[self.indices[i-1]:self.indices[i]] for i in range(1, len(self.indices))])
+
+        N, M = len(self.pool_map), len(set(self.pool_map)) # N = HQ verts, M = LQ verts
+        # Assign correspondences as V0_lq = [V0_hq, V1_hq, ..., VN_hq]
+        d = [[] for _ in range(M)] # list of lists
+        for i in range(N): d[self.pool_map[i]] += [i]
+        # max sublist len
+        m = max([len(d[i]) for i in range(M)])
+        # expand all sublist to max len
+        for i in range(M):
+            while len(d[i]) < m: 
+                # repeating elements will NOT change MIN/MAX output
+                d[i] += [d[i][0]] 
+ 
+        x = tf.reduce_max(tf.gather(x, d, axis=1), axis=2)
+        return x
+
+
+
+
+
+class GNNDiscriminator(tf.keras.Model):
+    def __init__(self, batchSize, isSkirt=False, n_feature=256, pool_type="max", global_average=True):
+        super(GNNDiscriminator, self).__init__()
+
+        self.uv2mesh = UV2Mesh(res=256)        
+
+        laplacians = [sio.loadmat("custom_utils/laplacian/ssmpl" + ("_skirt" if isSkirt else "") + str(i) + ".mat")['L'].tocoo() for i in range(5)] 
+        poolings = [sio.loadmat("custom_utils/pooling/pool" + str(i) + str(i+1) + ".mat")['idx'] for i in range(5)]
+   
+        self.conv1 = GraphConvolution((3, int(n_feature/2**2)), activation=relu, L=laplacians[0], batchSize=batchSize, normalization=BatchNormalization(), name='gcn_1')
+        self.pool1 = GraphPool(pool_map=poolings[0], type=pool_type, batch_size=batchSize)
+
+        self.conv2 = GraphConvolution((int(n_feature/2**2), int(n_feature/2**2)), activation=relu, L=laplacians[1], batchSize=batchSize, normalization=BatchNormalization(), name='gcn_2')
+        self.pool2 = GraphPool(pool_map=poolings[1], type=pool_type, batch_size=batchSize)
+
+        self.custom_layers = []
+        for i in range(2,5):
+            self.custom_layers.append(GraphConvolution((int(n_feature/2**(4- (i if i == 2 else i-1) )), int(n_feature/2**(4-i))), activation=relu, L=laplacians[i], batchSize=batchSize, normalization=BatchNormalization()))
+            for j in range(2**(i-2)):
+                conv1 = GraphConvolution((int(n_feature/2**(4-i)), int(n_feature/2**(4-i))), activation=relu, L=laplacians[i], batchSize=batchSize, normalization=BatchNormalization())
+                conv2 = GraphConvolution((int(n_feature/2**(4-i)), int(n_feature/2**(4-i))), activation=relu, L=laplacians[i], batchSize=batchSize, normalization=BatchNormalization())
+                self.custom_layers.append([conv1, conv2])
+            self.custom_layers.append(GraphPool(pool_map=poolings[i], type=pool_type, batch_size=batchSize))
+            
+            self.custom_layers.append(Dense(256, relu))
+            self.custom_layers.append(BatchNormalization())
+            self.custom_layers.append(Dense(int(n_feature/2**(4-i)), relu))
+            self.custom_layers.append(BatchNormalization())
+
+        self.flatten = Flatten()
+        self.last = Dense(1)
+
+    def call(self, inputs, training=None):
+        inputs = self.uv2mesh(inputs)
+
+        x = self.conv1(inputs, training)
+        x = self.pool1(x)
+
+        x = self.conv2(x, training)
+        x = self.pool2(x)
+
+        for layer in self.custom_layers:
+            if isinstance(layer, list):
+                x1 = layer[0](x, training)
+                x2 = layer[1](x1, training)
+                x = Add()([x, x2])
+            else:
+                x = layer(x, training=training)
+
+        x = self.flatten(x)
+        x = self.last(x)
+
+        return x
+
+
+
+
+
+
+class DatasetTest(tf.data.Dataset):
+
+    def _generator():
+        for i in range(10):
+            yield np.array([i], dtype=np.int32)
+
+    def __new__(cls):
+        return tf.data.Dataset.from_generator(
+            cls._generator,
+            output_signature = tf.TensorSpec(shape = (1), dtype = tf.int32) ,
+            args=()
+        )
+
+
+
+if __name__ == '__main__':
+    import os
+    os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+
+    batch_size = 4
+    input = tf.ones((batch_size, 14475, 3))
+
+    ge = GNNDiscriminator(batchSize=batch_size, n_feature=32)
+
+    x = ge(input)
+
+    print()

custom_utils/smpl.py

@@ -121,6 +121,7 @@ class SMPL:
 
     def run(self, R, betas, remove=True):
         batch_size = betas.shape[0]
+        # R = self.run_rodrigues(pose)
 
         #J_regressor = tf.gather(self.J_regressor, gender, axis=0)
         #shapedirs = tf.gather(self.shapedirs, gender, axis=0)