207 lines
7.8 KiB
Python
207 lines
7.8 KiB
Python
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() |