plot

curvature_analysis.py

@@ -133,6 +133,7 @@ def generate_curvatures(v, faces, mask, neighbourhood=25, nr_of_labels=6):
        
     return frame_labels
 
+import tensorflow as tf
 
 def visualize_curvature_labels(folder, results_folder, frame):
     id = frame.split('_')[0]
@@ -146,7 +147,8 @@ def visualize_curvature_labels(folder, results_folder, frame):
 
     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,:])
+    # dice = dice_dense(curvature_labels_pred[None,:], curvature_labels_gt[None,:])
+    dice = dice_dense(tf.one_hot(curvature_labels_pred, 7)[None,:], tf.one_hot(curvature_labels_gt, 7)[None,:]).numpy()
     colors_binned_all = []
 
     for curvature_labels in [curvature_labels_gt, curvature_labels_pred]:
@@ -161,16 +163,16 @@ def visualize_curvature_labels(folder, results_folder, frame):
 
     fig.add_trace(go.Mesh3d(
                 # 8 vertices of a cube
-                x=verts[:,0],
-                y=verts[:,1],
-                z=verts[:,2],
+                x=-pred[:,0],
+                y=pred[:,2],
+                z=pred[:,1],
                 # 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
+                vertexcolor=np.ones((verts.shape[0],3))*np.array([255,200,200])[None,:]
             ),
-            row=1, col=1)
+            row=1, col=2)
 
     fig.add_trace(go.Mesh3d(
                 # 8 vertices of a cube
@@ -183,7 +185,7 @@ 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
@@ -247,7 +249,7 @@ def init_wandb():
     run.log_artifact(artifact)
 
 if __name__ == '__main__':
-    init_wandb()
+    # init_wandb()
 
     is_skirt = False
     skirt_modifier = '_skirt' if is_skirt else ''
@@ -273,6 +275,8 @@ if __name__ == '__main__':
     i = 0
     while True:
         frame = np.random.choice(frames)
+        # frame = '03265_193'
+        frame = '03102_49'
         fig, dice = visualize_curvature_labels(folder, results_folder, frame)
         wandb.log({f'{frame}-dice:{dice}': fig}, step=i)
         i += 1