generate multiple plots on wandb

curvature_analysis.py

@@ -1,5 +1,8 @@
 import os
 import sys
+os.environ["CUDA_VISIBLE_DEVICES"] = '4'
+
+
 import tqdm
 import argparse
 import random
@@ -16,11 +19,34 @@ from scipy.sparse import csr_matrix, coo_matrix
 from scipy.sparse.csgraph import floyd_warshall, dijkstra
 from plotly.subplots import make_subplots
 from custom_utils.custom_utils import _maskFaces
+from postprocessing import fix_collision
 
 from custom_utils.supersmpl_io import readPC2, readPC2Frame, writePC2, readOBJ
 from custom_utils.custom_utils import _maskFaces
 from auto_encoder_curv import dice_dense
 
+def postprocess(mesh, faces, mask=None, body=None, do_smoothing=True):
+    if body is not None:
+        mesh_fixed = fix_collision(mesh, body[np.newaxis,:,:], mask[np.newaxis,:])[0].numpy()[mask, :]
+    else:
+        mesh_fixed = mesh
+
+    if do_smoothing:
+        m = tm.Trimesh(mesh_fixed, faces, process=False, validate=False)
+        outliers = np.array([len(v) for v in m.vertex_neighbors]) == 0
+        outliers = np.nonzero(outliers)[0]
+        outliers_n_neighbours = get_closest_n_points(mesh_fixed, outliers, 2)
+        outliers_faces = np.concatenate((outliers[:, np.newaxis], outliers_n_neighbours), axis=-1)
+        faces_fixed = np.concatenate((faces, outliers_faces), axis=0)
+        m = tm.Trimesh(mesh_fixed, faces_fixed, process=False, validate=False)
+        tm.smoothing.filter_humphrey(m, iterations=1, alpha=0.6, beta=0.4)
+
+        mesh = m.vertices
+    else:
+        mesh = mesh_fixed
+
+    return mesh
+
 def export_to_ply(mesh, file):
     ply = tm.exchange.ply.export_ply(mesh, encoding='binary', vertex_normal=True, include_attributes=True)
     output_file = open(file, "wb+")
@@ -97,7 +123,7 @@ def generate_curvatures(v, faces, mask, neighbourhood=25, nr_of_labels=6):
     cloud.add_scalar_field("curvature", ev=ev)
     curvature = np.array(cloud.points["curvature(" + str(neighbourhood+1) + ")"])
     
-    bins = get_bins(curvature, nr_of_labels)
+    bins = np.load(f'custom_data/train_bins.npy') # 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)
@@ -116,6 +142,7 @@ def visualize_curvature_labels(folder, results_folder, frame):
     _, faces = _maskFaces(np.zeros((14475,3)), F, mask)
     verts = np.load(os.path.join(folder, 'mesh_posed', f'{id}_{frame_nr}.npz'))['mesh_posed'][mask,:]
     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)
@@ -178,9 +205,10 @@ def visualize_curvature_labels(folder, results_folder, frame):
                     # 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(nticks=4, range=[-1,1],),
-                yaxis = dict(nticks=4, range=[-1,1],),
-                zaxis = dict(nticks=4, range=[-1,1],),)
+    helper = dict(xaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),
+                yaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),
+                zaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),)
+        
     fig.update_layout(scene = helper,
                         scene2 = helper,
                         scene3 = helper,
@@ -208,12 +236,18 @@ def visualize_curvature_labels(folder, results_folder, frame):
     #                     zaxis = dict(nticks=4, range=[-1,1],),)
     #                     )
 
-    fig.show()
+    return fig, dice
 
-    print()
+import wandb
 
+def init_wandb():
+    run = wandb.init(project='plots', name='review_4', 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()
 
     is_skirt = False
     skirt_modifier = '_skirt' if is_skirt else ''
@@ -236,4 +270,10 @@ if __name__ == '__main__':
     else:
         dist = np.load('custom_data/dijkstra_skirt.npy')
 
-    visualize_curvature_labels(folder, results_folder, '00100_9')
+    i = 0
+    while True:
+        frame = np.random.choice(frames)
+        fig, dice = visualize_curvature_labels(folder, results_folder, frame)
+        wandb.log({f'{frame}-dice:{dice}': fig}, step=i)
+        i += 1
+        print(frame)