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