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) for k, v in m2uv.items(): if len(v) == 2: result[k] = list(v) elif len(v) == 1: x = v.pop() result[k] = [x, x] else: assert False return result 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') self.uv_map_discreet = np.floor(uv_map * self.res).astype(np.int32) m2uv = np.load('custom_utils/m2uv' + ('_skirt' if is_skirt else '') + ('_orig' if is_smpl_orig else '') + '.npy', allow_pickle=True).item() 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.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) @tf.function def uv_to_mesh(self, uv): if uv.shape[1] == 3: uv = tf.transpose(uv, [0,2,3,1]) if uv.shape[1] != 512: uv = tf.image.resize(uv, (512, 256), tf.image.ResizeMethod.LANCZOS5) batch_size = uv.shape[0] x1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.zeros((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1) y1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1) z1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32) * 2), axis=-1) x2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.zeros((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1) y2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1) z2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32) * 2), axis=-1) x1 = tf.gather_nd(uv, x1_ind, batch_dims=1) y1 = tf.gather_nd(uv, y1_ind, batch_dims=1) z1 = tf.gather_nd(uv, z1_ind, batch_dims=1) x2 = tf.gather_nd(uv, x2_ind, batch_dims=1) y2 = tf.gather_nd(uv, y2_ind, batch_dims=1) z2 = tf.gather_nd(uv, z2_ind, batch_dims=1) p1 = tf.stack((x1, y1, z1), axis=-1) p2 = tf.stack((x2, y2, z2), axis=-1) mesh = (p1 + p2) / 2 return mesh def __call__(self, 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): super(UV2MeshLabel, 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') self.uv_map_discreet = np.floor(uv_map * self.res).astype(np.int32) m2uv = np.load('custom_utils/m2uv' + ('_skirt' if is_skirt else '') + ('_orig' if is_smpl_orig else '') + '.npy', allow_pickle=True).item() m2uv = m2uv_array(m2uv, is_smpl_orig=is_smpl_orig) self.uv_pixels = self.uv_map_discreet[m2uv] # input = channels_first # @tf.function def uv_to_mesh(self, uv): batch_size = uv.shape[0] x1_ind = tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]) x2_ind = tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]) x1 = tf.gather_nd(uv, x1_ind, batch_dims=1) x2 = tf.gather_nd(uv, x2_ind, batch_dims=1) mesh = (x1 + x2) / 2 return mesh def __call__(self, uv): return self.uv_to_mesh(uv) def _maskFaces(V, F, mask): V_mask = mask F_mask = np.take(V_mask, F.reshape(-1), axis=0) F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3 V_ind = np.zeros(V_mask.shape) V_ind[V_mask] = np.arange(V_mask.sum()) F_new = np.take(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3)) V_new = V[V_mask, :] return V_new, F_new # def _maskFaces_tf(V, F, mask): # V_mask = mask # F_mask = tf.gather(V_mask, F.reshape(-1), axis=0) # F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3 # F_mask = tf.reduce_sum(tf.reshape(F_mask, (-1, 3)), axis=1) == 3 # V_ind = tf.zeros(V_mask.shape) # V_ind[V_mask] = tf.range(V_mask.sum()) # F_new = tf.gather(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3)) # V_new = V[V_mask, :] # return V_new, F_new def maskFacesBatch(V, F, mask): Vs = [] Fs = [] for i in range(mask.shape[0]): v, f = _maskFaces(V[i], F, mask[i]) Vs.append(v) Fs.append(f) return Vs, Fs def flip_axis(x, ind=None): y = x.copy() if ind is not None: y[ind:,1] = x[ind:,2] y[ind:,2] = x[ind:,1] else: y[:,1] = x[:,2] y[:,2] = x[:,1] return y def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color, flip_xy=False): fig = make_subplots(rows=1, cols=3, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(3)] ]) if flip_xy: mesh_gt = flip_axis(mesh_gt) mesh_pred_unposed = flip_axis(mesh_pred_unposed) collision = flip_axis(collision, ind=14475) fig.add_trace(go.Mesh3d( # 8 vertices of a cube x=mesh_gt[:,0], y=mesh_gt[:,2], z=mesh_gt[:,1], # i, j and k give the vertices of triangles i = faces[:,0], j = faces[:,1], k = faces[:,2] ), row=1, col=1) fig.add_trace(go.Mesh3d( # 8 vertices of a cube x=mesh_pred_unposed[:,0], y=mesh_pred_unposed[:,2], z=mesh_pred_unposed[:,1], # i, j and k give the vertices of triangles i = faces[:,0], j = faces[:,1], k = faces[:,2] ), row=1, col=2) fig.add_trace(go.Mesh3d( # 8 vertices of a cube x=collision[:,0], y=collision[:,2], z=collision[:,1], # i, j and k give the vertices of triangles i = collision_faces[:,0], j = collision_faces[:,1], k = collision_faces[:,2], vertexcolor=collision_color ), row=1, col=3) 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), scene3_aspectmode='manual', scene3_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, scene3 = helper, ) return fig def plot_3d_collision(gt, collision_1, collision_2, collision_faces, collision_color, flip_xy=False): fig = make_subplots(rows=1, cols=3, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(3)] ]) if flip_xy: collision_1 = flip_axis(collision_1, ind=14475) collision_2 = flip_axis(collision_2, ind=14475) gt = flip_axis(gt, ind=14475) fig.add_trace(go.Mesh3d( x=gt[:,0], y=gt[:,2], z=gt[:,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_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=2) 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=3) 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), scene3_aspectmode='manual', scene3_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, scene3 = helper, ) return fig def batch_to_dict(batch, signature): result = dict() for i, key in enumerate(signature.keys()): result[key] = batch[i] return result conditions = ['pose_shape', 'uv_static', 'uv_mask', 'uv_shaped_body'] def nest_conditions(batch): batch['conditions'] = {d: batch.pop(d) for d in conditions} return batch 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 load_model(model, checkpoint, disable_assert=False, assert_consumed_on=False, iteration=''): ckpt = tf.train.Checkpoint(model=model) manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1) if not disable_assert: assert manager.latest_checkpoint is not None, checkpoint assert manager.latest_checkpoint.split('/')[-1] != 'ckpt-0', checkpoint if assert_consumed_on: ckpt.restore(manager.latest_checkpoint).assert_consumed() else: if iteration != '': name_base = manager.latest_checkpoint.split('-')[0] ckpt_ind = [i for i in range(len(manager.checkpoints)) if manager.checkpoints[i].endswith(name_base + '-' + iteration)][0] ckpt.restore(manager.checkpoints[ckpt_ind]).assert_existing_objects_matched()#.assert_consumed() else: ckpt.restore(manager.latest_checkpoint).assert_existing_objects_matched()#.assert_consumed() def get_summary_str(model): lines = [] model.summary(print_fn=lines.append, line_length=150,expand_nested=True,show_trainable=True) # Add initial spaces to avoid markdown formatting in TensorBoard return lines def get_summary_html(model): lines = get_summary_str(model) html = '' for line in lines: html += "
" + line + "
" 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 def group_inputs_for_levels(batch, config): levels_input = [] for res in config['pyramid']['resolutions']: res_modifier = '_' + str(res) if res != 512 else '' if config['use_skinning']: uv_orig = batch['uv_unposed' + res_modifier] else: uv_orig = batch['uv_posed' + res_modifier] levels_input.append([uv_orig, batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], batch['uv_curvature' + res_modifier], batch['uv_mask' + res_modifier], batch['uv_mask']]) conds = [batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature']] return levels_input, conds class Accumulator(): def __init__(self): self.dict = dict() def update(self, dict): for key in dict.keys(): if key in self.dict.keys(): self.dict[key].append(dict[key]) else: self.dict[key] = [dict[key]] def get(self): result = dict() for key in self.dict.keys(): result[key] = np.mean(self.dict[key]) return result def reset(self): self.dict = dict()