From 36a118a22ea6e479e561b3e8cf539e92692390a3 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Sat, 29 Oct 2022 20:15:02 +0000 Subject: [PATCH] add collision plotting --- custom_utils/custom_utils.py | 140 ++++++++++++++++++++++++++++++++--- 1 file changed, 128 insertions(+), 12 deletions(-) diff --git a/custom_utils/custom_utils.py b/custom_utils/custom_utils.py index 3364994..00e76b2 100644 --- a/custom_utils/custom_utils.py +++ b/custom_utils/custom_utils.py @@ -1,8 +1,13 @@ +import json import numpy as np import tensorflow as tf +import scipy as sp import plotly.graph_objects as go from plotly.subplots import make_subplots +from tensorflow_graphics.geometry.convolution.graph_pooling import pool, unpool +from tensorflow_graphics.geometry.convolution.graph_convolution import edge_convolution_template + def m2uv_array(m2uv, is_smpl_orig=False): result = np.zeros((14475 if not is_smpl_orig else 3638, 2), dtype=np.int32) @@ -18,8 +23,8 @@ def m2uv_array(m2uv, is_smpl_orig=False): return result -class UV2Mesh(): - def __init__(self, res=256, is_skirt=False, is_smpl_orig=False): +class UV2Mesh(tf.keras.Model): + def __init__(self, res=256, is_skirt=False, is_smpl_orig=False, is_adaptive=False): super(UV2Mesh, self).__init__() self.res = res uv_map = np.load('custom_utils/stretched_uv_map_' + ('skirt' if is_skirt else '3') + ('_orig' if is_smpl_orig else '') + '.npy') @@ -28,14 +33,12 @@ class UV2Mesh(): m2uv = m2uv_array(m2uv, is_smpl_orig=is_smpl_orig) self.uv_pixels = self.uv_map_discreet[m2uv] self.smpl_res = 14475 if not is_smpl_orig else 3638 - # self.batch_size = batch_size + self.is_adaptive = is_adaptive + + if self.is_adaptive: + laplacian = sp.sparse.load_npz("custom_utils/laplacian/laplacian0" + ("_skirt" if is_skirt else "") + ".npz") + self.gc = GraphConvolution((3, 3), L=laplacian, batchSize=4) - # self.x1 = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.zeros((self.batch_size, 14475, 1), dtype=tf.int32)), axis=-1) - # self.y1 = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.ones((self.batch_size, 14475, 1), dtype=tf.int32)), axis=-1) - # self.z1 = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.ones((self.batch_size, 14475, 1), dtype=tf.int32) * 2), axis=-1) - # self.x2 = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.zeros((self.batch_size, 14475, 1), dtype=tf.int32)), axis=-1) - # self.y2 = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.ones((self.batch_size, 14475, 1), dtype=tf.int32)), axis=-1) - # self.z2 = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.ones((self.batch_size, 14475, 1), dtype=tf.int32) * 2), axis=-1) @tf.function def uv_to_mesh(self, uv): @@ -67,9 +70,14 @@ class UV2Mesh(): mesh = (p1 + p2) / 2 return mesh + def __call__(self, uv): - return self.uv_to_mesh(uv) - + x = self.uv_to_mesh(uv) + if self.is_adaptive: + x = self.gc(x) + return x + + class UV2MeshLabel(): def __init__(self, res=256, is_skirt=False, is_smpl_orig=False): @@ -192,6 +200,45 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, colli return fig +def plot_3d_collision(collision_1, collision_2, collision_faces, collision_color): + fig = make_subplots(rows=1, cols=2, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(2)] ]) + + fig.add_trace(go.Mesh3d( + x=collision_1[:,0], + y=collision_1[:,2], + z=collision_1[:,1], + i = collision_faces[:,0], + j = collision_faces[:,1], + k = collision_faces[:,2], + vertexcolor=collision_color + ), + row=1, col=1) + + fig.add_trace(go.Mesh3d( + x=collision_2[:,0], + y=collision_2[:,2], + z=collision_2[:,1], + i = collision_faces[:,0], + j = collision_faces[:,1], + k = collision_faces[:,2], + vertexcolor=collision_color + ), + row=1, col=2) + + + fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1), + scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1), + ) + helper = dict(xaxis = dict(nticks=4, range=[-1,1],), + yaxis = dict(nticks=4, range=[-1,1],), + zaxis = dict(nticks=4, range=[-1,1],),) + fig.update_layout(scene = helper, + scene2 = helper, + ) + + return fig + + def batch_to_dict(batch, signature): result = dict() for i, key in enumerate(signature.keys()): @@ -234,4 +281,73 @@ def get_summary_html(model): html = '' for line in lines: html += "
" + line + "
" - return html \ No newline at end of file + return html + + +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 merge_verts(Vs, batchDim): + V = tf.concat(tf.unstack(Vs, num=batchDim), axis=0) + return V + + +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 + + +def replace_arrays(string): + out = string[:1] + i = 0 + while i < len(string) and string[i+1:].find('[') != -1: + start = string[i+1:].find('[') + end = string[i+1+start:].find(']') + out += string[i+1:i+1+start] + array = string[i+1+start:i+1+start+end+1].replace('\n', '') + out += array + i = i + 1 + start + end + + out += string[i+1:] + + return out \ No newline at end of file