lot of things

custom_utils/custom_utils.py

@@ -66,6 +66,33 @@ class UV2Mesh():
         return self.uv_to_mesh(uv)
         
 
+class UV2MeshLabel():
+    def __init__(self, res=256):
+        super(UV2MeshLabel, self).__init__()
+        self.res = res
+        uv_map = np.load('custom_utils/stretched_uv_map_3.npy')
+        self.uv_map_discreet = np.floor(uv_map * self.res).astype(np.int32)
+        m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item()
+        m2uv = m2uv_array(m2uv)
+        self.uv_pixels = self.uv_map_discreet[m2uv]
+
+    # input = channels_first
+    # @tf.function
+    def uv_to_mesh(self, uv):
+        batch_size = uv.shape[0]
+
+        x1_ind = tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1])
+        x2_ind = tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1])
+
+        x1 = tf.gather_nd(uv, x1_ind, batch_dims=1)
+        x2 = tf.gather_nd(uv, x2_ind, batch_dims=1)
+
+        mesh = (x1 + x2) / 2
+        return mesh
+
+    def __call__(self, uv):
+        return self.uv_to_mesh(uv)
+
 
 def _maskFaces(V, F, mask):
     V_mask = mask
@@ -79,6 +106,21 @@ def _maskFaces(V, F, mask):
 
     return V_new, F_new
 
+
+# def _maskFaces_tf(V, F, mask):
+#     V_mask = mask
+#     F_mask = tf.gather(V_mask, F.reshape(-1), axis=0)
+#     F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3
+#     F_mask = tf.reduce_sum(tf.reshape(F_mask, (-1, 3)), axis=1) == 3
+
+#     V_ind = tf.zeros(V_mask.shape)
+#     V_ind[V_mask] = tf.range(V_mask.sum())
+#     F_new = tf.gather(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3))
+#     V_new = V[V_mask, :]
+
+#     return V_new, F_new
+
+
 def maskFacesBatch(V, F, mask):
     Vs = []
     Fs = []
@@ -148,3 +190,10 @@ conditions = ['pose_shape', 'uv_static', 'uv_mask', 'uv_shaped_body']
 def nest_conditions(batch):
     batch['conditions'] = {d: batch.pop(d) for d in conditions}
     return batch
+
+
+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)

custom_utils/graph_utils.py

@@ -1,10 +1,14 @@
+import os
+import sys
 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
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+# 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
@@ -31,6 +35,14 @@ def repeat_laplacian(L, times):
     return sp.sparse.coo_matrix((data, (np.squeeze(np.asarray(idx[0,:])), np.squeeze(np.asarray(idx[1,:])))), shape)
 
 
+def convert_pooling(pooling, batchSize):
+    pool_map = np.zeros((np.max(pooling)+1, pooling.shape[0]), dtype=np.float32)
+    for i in range(pooling.shape[0]):
+        pool_map[pooling[i], i] = 1.0
+    pool_map = np.tile(pool_map[np.newaxis, :, :], [batchSize, 1, 1])
+    pool_map = tf.sparse.from_dense(pool_map)
+    return pool_map
+
 
 class GraphConvolution(tf.keras.layers.Layer):
     def __init__(self, shape, activation=None, name='gcn', L=None, batchSize=None, normalization=None):
@@ -174,8 +186,76 @@ class GNNDiscriminator(tf.keras.Model):
         return x
 
 
+class Reconstructor_V4_Decoder(tf.keras.Model):
+    def __init__(self, batchSize, isSkirt=False, n_feature=64):
+        super(Reconstructor_V4_Decoder, self).__init__()
 
 
+        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)]
+        poolings = [convert_pooling(poolings[i][0,:], batchSize) for i in range(len(poolings))]
+
+        self.custom_layers = []
+        for i in range(4,1,-1):
+            self.custom_layers.append(GraphUnPool(pool_map=poolings[i]))
+            self.custom_layers.append(GraphConvolution((int(n_feature/2**(4-(i if i == 4 else i+1))), int(n_feature/2**(4-i)) ), activation=relu, L=laplacians[i], batchSize=batchSize, normalization=BatchNormalization()))
+            for _ 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.pool2 = GraphUnPool(pool_map=poolings[1])
+        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.pool1 = GraphUnPool(pool_map=poolings[0])
+        self.conv1 = GraphConvolution((int(n_feature/2**2), int(n_feature/2**3)), activation=relu, L=laplacians[0], batchSize=batchSize, normalization=BatchNormalization(), name='gcn_1')
+
+        self.conv0 = GraphConvolution((int(n_feature/2**3), 3), L=laplacians[0], batchSize=batchSize, name='gcn_0')
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        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.pool2(x)
+        x = self.conv2(x, training)
+
+        x = self.pool1(x)
+        x = self.conv1(x, training)
+
+        x = self.conv0(x, training)
+
+        return x
+
+
+
+class Reconstructor_v2(tf.keras.Model):
+    def __init__(self, batchSize, isSkirt=False):
+        super(Reconstructor_v2, self).__init__()
+
+        self.flatten = tf.keras.layers.Flatten()
+
+        self.fc2 = tf.keras.layers.Dense(14475*3, kernel_regularizer=tf.keras.regularizers.l2)
+
+        self.reshape2 = tf.keras.layers.Reshape((14475, 3))
+        # self.drop = Dropout(0.2)
+        # self.fc3 = Dense(512)
+        self.fc1 = tf.keras.layers.Dense(3)
+
+    def call(self, inputs, training=None):
+        x = self.flatten(inputs)
+        x = self.fc2(x)
+        # x = self.drop(x, training=training)
+        x = self.reshape2(x)
+        # x = self.fc3(x)
+        x = self.fc1(x)
+        return x
 
 
 class DatasetTest(tf.data.Dataset):
@@ -195,13 +275,16 @@ class DatasetTest(tf.data.Dataset):
 
 if __name__ == '__main__':
     import os
-    os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+    os.environ["CUDA_VISIBLE_DEVICES"] = "2"
 
     batch_size = 4
-    input = tf.ones((batch_size, 14475, 3))
+    # input = tf.ones((batch_size, 14475, 3))
 
-    ge = GNNDiscriminator(batchSize=batch_size, n_feature=32)
+    # ge = GNNDiscriminator(batchSize=batch_size, n_feature=32)
 
-    x = ge(input)
+    # x = ge(input)
+
+    gd = Reconstructor_V4_Decoder(batch_size)
+    x = gd(tf.ones((batch_size, 21, 64)))
 
     print()