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 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 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) 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): 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 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): 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"] = "2" batch_size = 4 # input = tf.ones((batch_size, 14475, 3)) # ge = GNNDiscriminator(batchSize=batch_size, n_feature=32) # x = ge(input) gd = Reconstructor_V4_Decoder(batch_size) x = gd(tf.ones((batch_size, 21, 64))) print()