add collision plotting

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():
+class UV2Mesh(tf.keras.Model):
-    def __init__(self, res=256, is_skirt=False, is_smpl_orig=False):
+    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,8 +70,13 @@ 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():
@@ -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()):
@@ -235,3 +282,72 @@ def get_summary_html(model):
     for line in lines:
         html += "<pre>" + line + "</pre>"
     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