import numpy as np import tensorflow as tf import plotly.graph_objects as go def m2uv_array(m2uv): result = np.zeros((14475, 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(): def __init__(self, res=256): super(UV2Mesh, self).__init__() self.res = res uv_map = np.load('custom_utils/stretched_uv_map_3.npy') self.uv_map_discreet = np.floor(uv_map * self.res).astype(np.int32) m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item() m2uv = m2uv_array(m2uv) self.uv_pixels = self.uv_map_discreet[m2uv] # self.batch_size = batch_size # 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) # input = channels_first @tf.function def uv_to_mesh(self, uv): uv = tf.transpose(uv, [0,2,3,1]) 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, 14475, 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, 14475, 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, 14475, 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, 14475, 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, 14475, 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, 14475, 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): return self.uv_to_mesh(uv) class UV2MeshLabel(): def __init__(self, res=256): super(UV2MeshLabel, self).__init__() self.res = res uv_map = np.load('custom_utils/stretched_uv_map_3.npy') self.uv_map_discreet = np.floor(uv_map * self.res).astype(np.int32) m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item() m2uv = m2uv_array(m2uv) 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 plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color): mesh_pred_unposed[:,1] += 1 collision[:,1] += 2 fig = go.Figure(data=[ go.Mesh3d( # 8 vertices of a cube x=mesh_gt[:,0], y=mesh_gt[:,1], z=mesh_gt[:,2], # i, j and k give the vertices of triangles i = faces[:,0], j = faces[:,1], k = faces[:,2] ), go.Mesh3d( # 8 vertices of a cube x=mesh_pred_unposed[:,0], y=mesh_pred_unposed[:,1], z=mesh_pred_unposed[:,2], # i, j and k give the vertices of triangles i = faces[:,0], j = faces[:,1], k = faces[:,2], ), go.Mesh3d( # 8 vertices of a cube x=collision[:,0], y=collision[:,1], z=collision[:,2], # i, j and k give the vertices of triangles i = collision_faces[:,0], j = collision_faces[:,1], k = collision_faces[:,2], vertexcolor=collision_color ) ]) fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=2, z=1)) fig.update_layout( scene = dict( xaxis = dict(nticks=4, range=[-1,1],), yaxis = dict(nticks=4, range=[-1,3],), zaxis = dict(nticks=4, range=[-1,1],),) ) 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): ckpt = tf.train.Checkpoint(model=model) manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1) ckpt.restore(manager.latest_checkpoint)