added skinning and gnn discriminator
This commit is contained in:
@@ -0,0 +1,207 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user