lot of things

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