general update

custom_utils/custom_utils.py

@@ -4,8 +4,8 @@ import plotly.graph_objects as go
 from plotly.subplots import make_subplots
 
 
-def m2uv_array(m2uv):
-    result = np.zeros((14475, 2), dtype=np.int32)
+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)
@@ -19,14 +19,15 @@ def m2uv_array(m2uv):
 
 
 class UV2Mesh():
-    def __init__(self, res=256, is_skirt=False):
+    def __init__(self, res=256, is_skirt=False, is_smpl_orig=False):
         super(UV2Mesh, self).__init__()
         self.res = res
-        uv_map = np.load('custom_utils/stretched_uv_map_' + ('skirt' if is_skirt else '3') + '.npy')
+        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 '') + '.npy', allow_pickle=True).item()
-        m2uv = m2uv_array(m2uv)
+        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.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)
@@ -45,12 +46,12 @@ class UV2Mesh():
             
         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_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)
@@ -71,13 +72,13 @@ class UV2Mesh():
         
 
 class UV2MeshLabel():
-    def __init__(self, res=256, is_skirt=False):
+    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') + '.npy')
+        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 '') + '.npy', allow_pickle=True).item()
-        m2uv = m2uv_array(m2uv)
+        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