updated pipeline, added uv_mask_loss

custom_utils/custom_utils.py

@@ -21,7 +21,7 @@ class UV2Mesh():
         super(UV2Mesh, self).__init__()
         self.res = res
         uv_map = np.load('custom_utils/stretched_uv_map_3.npy')
-        self.uv_map_discreet = np.round(uv_map * self.res).astype(np.int32)
+        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]
@@ -35,6 +35,7 @@ class UV2Mesh():
         # 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]