try

auto_encoder_curv.py

@@ -95,7 +95,7 @@ def dice_dense(prediction, ground_truth, weight_map=None):
     return tf.reduce_mean(dice_score)
 
 
-@tf.function
+# @tf.function
 def train_step(batch):
     with tf.GradientTape() as tape:
         pred = model([batch['uv_curvature'], batch['uv_mask']], training=True)

curvature_analysis.py

@@ -117,13 +117,17 @@ def generate_curvatures(v, faces, mask, neighbourhood=25, nr_of_labels=6):
 
     cloud = PyntCloud.from_file(temp_file)
 
-    k_neighbors = np.argsort(dist[mask, :][:, mask], axis=-1)[:, :neighbourhood] #cloud.get_neighbors(k=neighbourhood)
+    k_neighbors = np.argsort(dist[mask, :][:, mask], axis=-1)[:, :neighbourhood]
+    # k_neighbors = cloud.get_neighbors(k=neighbourhood)
+    # tree = sp.spatial.cKDTree(verts)
+    # k_neighbors = tree.query(verts, 25)[1]
     #k_neighbors = k_neighbors_cache[id] # np.argsort(dist[mask, :][:, mask], axis=-1)[:, :neighbourhood] #cloud.get_neighbors(k=neighbourhood)
     ev = cloud.add_scalar_field("eigen_values", k_neighbors=k_neighbors)
     cloud.add_scalar_field("curvature", ev=ev)
     curvature = np.array(cloud.points["curvature(" + str(neighbourhood+1) + ")"])
     
-    bins = np.load(f'custom_data/train_bins.npy') # get_bins(curvature, nr_of_labels)
+    bins = np.load(f'custom_data/train_bins.npy')
+    # bins = get_bins(curvature, nr_of_labels)
     labels = np.linspace(0, bins.shape[0] - 2, bins.shape[0] - 1)
 
     frame_labels = np.zeros((curvature.shape[0]), dtype=np.float32)
@@ -131,7 +135,45 @@ def generate_curvatures(v, faces, mask, neighbourhood=25, nr_of_labels=6):
         frame_labels[(curvature >= bins[j]) * (curvature < bins[j+1])] = labels[j]
     frame_labels[curvature >= bins[-2]] = labels[-1]
        
-    return frame_labels
+    return frame_labels, curvature
+
+import tensorflow as tf
+from configs.config_abl_overfit_jumpsuit_2 import config
+from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision, UV2MeshLabel
+
+from custom_utils.smpl import SMPL
+from models.cvae_model import CVAE
+
+uv2mesh = UV2Mesh(res=256, is_skirt=False, is_adaptive=False, is_smpl_orig=False)
+uv2meshlabel = UV2MeshLabel(res=256, is_skirt=False)
+
+def get_prediction(folder, frame, curv):
+    smpl = SMPL(use_one_hot=False, orig_smpl=False)
+    model = CVAE(config)
+    load_model(model, 'checkpoints/' + config['name'])
+
+    id = frame.split('_')[0]
+    frame_nr = frame.split('_')[1].split('.')[0]
+
+    mesh_mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy')).astype(bool)
+    uv_mask = np.load(os.path.join(folder, 'uv_mask', f'{id}.npz'))['uv_mask'].astype(np.float32)  
+    pose_shape = np.load(os.path.join(folder, 'pose_shape', f'{id}_{frame_nr}.npy')).astype(np.float32)  
+    pose_shape[:3] = 0
+    uv_unposed = np.load(os.path.join(folder, 'uv_unposed_no_shape', f'{id}_{frame_nr}.npz'))['uv_unposed_no_shape']
+    uv_static = np.load(os.path.join(folder, 'uv_static', f'{id}.npz'))['uv_static']
+    uv_body_posed = np.load(os.path.join(folder, 'uv_posed_body_no_orient_2', f'{id}_{frame_nr}.npz'))['uv_posed_body_no_orient']
+
+    pred, _, _ = model([uv_unposed[None,:,:,:], uv_static[None,:,:,:], uv_body_posed[None,:,:,:], curv[None,:,:,:], uv_mask[None,:,:]], training=False)
+    mesh_pred = uv2mesh(pred)
+
+    shape = pose_shape[None, 72:]
+    gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
+    shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
+
+    mesh_pred_shaped = mesh_pred + shape_offsets * mesh_mask[tf.newaxis,:,tf.newaxis]
+    mesh_pred_posed = smpl.run_garment(pose_shape[None, :72], pose_shape[None, 72:], mesh_pred_shaped)
+    
+    return mesh_pred_posed.numpy()[0,mesh_mask,:]
 
 
 def visualize_curvature_labels(folder, results_folder, frame):
@@ -144,12 +186,24 @@ def visualize_curvature_labels(folder, results_folder, frame):
     pred = np.load(os.path.join(results_folder, frame + '.npy'))
     pred = postprocess(pred, faces)
 
-    curvature_labels_gt = generate_curvatures(verts, faces, mask, 25)
-    curvature_labels_pred = generate_curvatures(pred, faces, mask, 25)
-    dice = dice_dense(curvature_labels_pred[None,:], curvature_labels_gt[None,:])
+    uv_curvature = np.load(os.path.join(folder, 'uv_curvature', f'{id}_{frame_nr}.npz'))['curvature_uv']
+    mesh_curvature = np.load(os.path.join(folder, 'mesh_curvature', f'{id}_{frame_nr}.npz'))['mesh_curvature'][mask]
+    mesh_curvature_raw = np.load(os.path.join(folder, 'mesh_curvature_raw', f'{id}_{frame_nr}.npz'))['mesh_curvature_raw'][mask]
+    uv_curvature = tf.one_hot(uv_curvature, 7)
+    pred_new = get_prediction(folder, frame, uv_curvature)
+    pred_new = postprocess(pred_new, faces)
+
+    curvature_labels_gt, curv_gt = generate_curvatures(verts, faces, mask, 25)
+    curvature_labels_pred, _ = generate_curvatures(pred, faces, mask, 25)
+    curvatures_new = (uv2meshlabel(uv_curvature[None,:,:,:])[0].numpy() * mask[:,None])[mask, :]
+    curvatures_pred_new, curv_gt_new = generate_curvatures(pred_new, faces, mask, 25)
+    dice = dice_dense(tf.one_hot(curvature_labels_pred, 7)[None,:], tf.one_hot(curvature_labels_gt, 7)[None,:])
+    dice_new = dice_dense(tf.one_hot(curvatures_pred_new, 7)[None,:], curvatures_new[None,:])
+    curvatures_new = curvatures_new.argmax(axis=-1) - 1
+
     colors_binned_all = []
 
-    for curvature_labels in [curvature_labels_gt, curvature_labels_pred]:
+    for curvature_labels in [curvature_labels_gt, curvature_labels_pred, curvatures_new, curvatures_pred_new]:
         colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, int(curvature_labels.max()+1))])
         colors_binned = np.zeros((curvature_labels.shape[0], 3))
         for i in range(int(curvature_labels.max())+1):
@@ -157,20 +211,7 @@ def visualize_curvature_labels(folder, results_folder, frame):
             colors_binned[curvature_labels == i, :] = colors[i]
         colors_binned_all.append(colors_binned)
 
-    fig = make_subplots(rows=1, cols=3, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(3)] ])
-
-    fig.add_trace(go.Mesh3d(
-                # 8 vertices of a cube
-                x=verts[:,0],
-                y=verts[:,1],
-                z=verts[:,2],
-                # i, j and k give the vertices of triangles
-                i = faces[:,0],
-                j = faces[:,1],
-                k = faces[:,2],
-                vertexcolor=np.ones((verts.shape[0],3))*200
-            ),
-            row=1, col=1)
+    fig = make_subplots(rows=1, cols=4, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(4)] ])
 
     fig.add_trace(go.Mesh3d(
                 # 8 vertices of a cube
@@ -183,8 +224,8 @@ def visualize_curvature_labels(folder, results_folder, frame):
                 k = faces[:,2],
                 vertexcolor=colors_binned_all[0]
             ),
-            row=1, col=2)
-    
+            row=1, col=1)
+
     fig.add_trace(go.Mesh3d(
                 # 8 vertices of a cube
                 x=-pred[:,0],
@@ -197,12 +238,39 @@ def visualize_curvature_labels(folder, results_folder, frame):
                 vertexcolor=colors_binned_all[1],
                 text=f'dice: {dice}'
             ),
+            row=1, col=2)
+    
+    fig.add_trace(go.Mesh3d(
+                # 8 vertices of a cube
+                x=verts[:,0],
+                y=verts[:,1],
+                z=verts[:,2],
+                # i, j and k give the vertices of triangles
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2],
+                vertexcolor=colors_binned_all[2]
+            ),
             row=1, col=3)
 
+    fig.add_trace(go.Mesh3d(
+                # 8 vertices of a cube
+                x=-pred_new[:,0],
+                y=pred_new[:,2],
+                z=pred_new[:,1],
+                # i, j and k give the vertices of triangles
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2],
+                vertexcolor=colors_binned_all[3],
+                text=f'dice: {dice_new}'
+            ),
+            row=1, col=4)
+
     fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
                     scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1),
                     scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1),
-                    # scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1),
+                    scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1),
                     # scene5_aspectmode='manual', scene5_aspectratio=dict(x=1, y=1, z=1),
                     )
     helper = dict(xaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),
@@ -212,7 +280,7 @@ def visualize_curvature_labels(folder, results_folder, frame):
     fig.update_layout(scene = helper,
                         scene2 = helper,
                         scene3 = helper,
-                        # scene4 = helper,
+                        scene4 = helper,
                         # scene5 = helper,
                         )
 
@@ -241,13 +309,13 @@ def visualize_curvature_labels(folder, results_folder, frame):
 import wandb
 
 def init_wandb():
-    run = wandb.init(project='plots', name='review_4', dir='/home/workspace/local-data/temp/')
+    run = wandb.init(project='plots', name='review_5', dir='/home/workspace/local-data/temp/')
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('curvature_analysis.py')
     run.log_artifact(artifact)
 
 if __name__ == '__main__':
-    init_wandb()
+    # init_wandb()
 
     is_skirt = False
     skirt_modifier = '_skirt' if is_skirt else ''
@@ -272,7 +340,8 @@ if __name__ == '__main__':
 
     i = 0
     while True:
-        frame = np.random.choice(frames)
+        # frame = np.random.choice(frames)
+        frame = '02806_129'
         fig, dice = visualize_curvature_labels(folder, results_folder, frame)
         wandb.log({f'{frame}-dice:{dice}': fig}, step=i)
         i += 1