large refactor

custom_utils/custom_utils.py

@@ -1,3 +1,4 @@
+from faulthandler import disable
 import numpy as np
 import tensorflow as tf
 import plotly.graph_objects as go
@@ -135,14 +136,14 @@ def maskFacesBatch(V, F, mask):
 
 
 def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
-    mesh_pred_unposed[:,1] += 1
-    collision[:,1] += 2
+    mesh_pred_unposed[:,2] += 1
+    collision[:,2] += 2
     fig = go.Figure(data=[
         go.Mesh3d(
             # 8 vertices of a cube
             x=mesh_gt[:,0],
-            y=mesh_gt[:,1],
-            z=mesh_gt[:,2],
+            y=mesh_gt[:,2],
+            z=mesh_gt[:,1],
             # i, j and k give the vertices of triangles
             i = faces[:,0],
             j = faces[:,1],
@@ -151,8 +152,8 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, colli
         go.Mesh3d(
             # 8 vertices of a cube
             x=mesh_pred_unposed[:,0],
-            y=mesh_pred_unposed[:,1],
-            z=mesh_pred_unposed[:,2],
+            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],
@@ -161,8 +162,8 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, colli
         go.Mesh3d(
             # 8 vertices of a cube
             x=collision[:,0],
-            y=collision[:,1],
-            z=collision[:,2],
+            y=collision[:,2],
+            z=collision[:,1],
             # i, j and k give the vertices of triangles
             i = collision_faces[:,0],
             j = collision_faces[:,1],
@@ -202,7 +203,10 @@ def sparse_to_tensor(L):
     return tf.SparseTensor(idx, data, shape)
 
 
-def load_model(model, checkpoint):
+def load_model(model, checkpoint, disable_assert=False):
     ckpt = tf.train.Checkpoint(model=model)
     manager = tf.train.CheckpointManager(ckpt, checkpoint, max_to_keep=1)
-    ckpt.restore(manager.latest_checkpoint)
+    if not disable_assert:
+        assert manager.latest_checkpoint is not None
+        assert manager.latest_checkpoint.split('/')[-1] != 'ckpt-0'
+    ckpt.restore(manager.latest_checkpoint).assert_existing_objects_matched()#.assert_consumed()

custom_utils/fill_uv.py

@@ -5,12 +5,13 @@ import numpy as np
 import plotly.graph_objects as go
 import plotly.express as ex
 
-
+os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+import tensorflow as tf
+import tensorflow_addons as tfa
 def fill(img, mask):
     coords = np.transpose(np.stack(np.nonzero((1 - mask))), (1,0))
     img_idx = np.transpose(np.stack(np.nonzero((mask))), (1,0))
     img_data = img[img_idx[:,0], img_idx[:,1], :]
-
     interpolator =  spi.NearestNDInterpolator(x=img_idx, y=img_data)
     img_2 = interpolator(coords)
     res = np.zeros((512,256,3), dtype=np.float32)
@@ -19,6 +20,19 @@ def fill(img, mask):
 
     return res
 
+# @tf.function
+def fill_tf(img, mask):
+    img = img[tf.newaxis,:,:,:]
+    coords_int = tf.transpose(tf.stack(tf.experimental.numpy.nonzero(1-mask)), (1,0))
+    coords = tf.cast(coords_int, dtype=tf.float32)
+    img_2 = tfa.image.interpolate_bilinear(img, coords[tf.newaxis,:,:])
+    # coords_2 = list(zip(tf.ones(coords_int.shape[0]), coords_int[:,0], coords_int[:,1]))
+    # coords = tf.transpose(tf.stack([coords_int[:,0], coords_int[:,1]]))
+    out = tf.tensor_scatter_nd_update(img[0], coords_int, img_2[0])
+
+    return out
+
+
 
 def fill_label(img, mask):
     coords = np.transpose(np.stack(np.nonzero((1 - mask))), (1,0))
@@ -44,7 +58,13 @@ def fill_uv_maps_posed_body(folder, type):
         id, _ = frame.split('_')
         # mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); mask = mask['uv_mask']
         img = np.load(os.path.join(folder, type, frame)); img = img['uv_posed_body_no_orient']
-        filled = fill(img, mask)
+        # filled = fill(img, mask)
+        img = np.zeros((100, 100, 3), dtype=np.float32)
+        mask = np.zeros((100, 100), dtype=np.int32)
+        img[25:75, 25:75, :] = 1
+        mask[25:75, 25:75] = 1
+
+        filled = fill_tf(img, mask)
         np.savez_compressed(os.path.join(folder, type + '_filled', frame), uv_posed_body_no_orient_filled = filled)