lot of things

custom_utils/custom_utils.py

@@ -66,6 +66,33 @@ class UV2Mesh():
         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
@@ -79,6 +106,21 @@ def _maskFaces(V, F, 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 = []
@@ -148,3 +190,10 @@ 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)