save

curvature_analysis.py

@@ -22,18 +22,11 @@ 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, save_stl
+from custom_utils.custom_utils import _maskFaces, save_stl, write_obj
 from auto_encoder_curv import dice_dense
 
 
-def write_obj(file, vertices, faces, colors):
+
-    with open(file, 'w') as file:
-        for v, c in zip(vertices, colors):
-            line = 'v ' + ' '.join([str(_) for _ in v]) + ' ' + ' '.join([str(_) for _ in c]) + '\n'
-            file.write(line)
-        for f in faces:
-            line = 'f ' + ' '.join([str(_) for _ in f]) + '\n'
-            file.write(line)
         
 
 def postprocess(mesh, faces, mask=None, body=None, do_smoothing=True):
@@ -261,15 +254,105 @@ def init_wandb():
     artifact.add_file('curvature_analysis.py')
     run.log_artifact(artifact)
 
-if __name__ == '__main__':
-    init_wandb()
 
-    is_skirt = False
+def visualize_neighbourhoods(folder, frame):
+    id = frame.split('_')[0]
+    frame_nr = frame.split('_')[1].split('.')[0]
+
+    mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy')).astype(bool)
+    _, 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,:]
+
+    curvature_labels_10 = generate_curvatures(verts, faces, mask, 10)
+    curvature_labels_25 = generate_curvatures(verts, faces, mask, 25)
+    curvature_labels_100 = generate_curvatures(verts, faces, mask, 100)
+    colors_binned_all = []
+
+    for curvature_labels in [curvature_labels_10, curvature_labels_25, curvature_labels_100]:
+        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):
+            # print(np.count_nonzero(curvature_labels == i))
+            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=colors_binned_all[0]
+            ),
+            row=1, col=1)
+
+    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[1]
+            ),
+            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.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),
+                    # 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),
+                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,
+                        # scene4 = helper,
+                        # scene5 = helper,
+                        )
+    
+    write_obj(f'examples/curvature_analysis/{id}_{frame_nr}_10.obj', verts, faces.astype(np.int32), colors_binned_all[0].astype(np.int32))
+    write_obj(f'examples/curvature_analysis/{id}_{frame_nr}_25.obj', verts, faces.astype(np.int32), colors_binned_all[1].astype(np.int32))
+    write_obj(f'examples/curvature_analysis/{id}_{frame_nr}_100.obj', verts, faces.astype(np.int32), colors_binned_all[2].astype(np.int32))
+    write_obj(f'examples/curvature_analysis/{id}_{frame_nr}_mesh.obj', verts, faces.astype(np.int32), np.ones((verts.shape[0], 3), dtype=np.int32) * 200)
+
+
+    fig.show()
+    return fig
+
+
+if __name__ == '__main__':
+    # init_wandb()
+
+    is_skirt = True
     skirt_modifier = '_skirt' if is_skirt else ''
     F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy").astype(np.int32)
     # k_neighbours(k=25)
 
-    folder = '/home/workspace/local-data/test_ready'
+    folder = '/home/workspace/local-data_128/test_skirt_ready_2'
     results_folder = '/home/workspace/data/results/iccv/iccv_pyr_32_64_128_256_512-120000'
 
     # frames_file = 'frames_temp.txt'
@@ -291,7 +374,8 @@ if __name__ == '__main__':
         frame = np.random.choice(frames)
         # frame = '03265_193'
         # frame = '03102_49'
-        fig, dice = visualize_curvature_labels(folder, results_folder, frame)
+        # fig, dice = visualize_curvature_labels(folder, results_folder, frame)
-        wandb.log({f'{frame}-dice:{dice}': fig}, step=i)
+        # wandb.log({f'{frame}-dice:{dice}': fig}, step=i)
+        visualize_neighbourhoods(folder, frame)
         i += 1
         print(frame)

custom_utils/custom_utils.py

@@ -442,4 +442,14 @@ def save_stl(file, mesh, faces):
         for j in range(3):
             obj.vectors[i][j] = mesh[f[j],:]    
 
-    obj.save(file)
+    obj.save(file)
+
+
+def write_obj(file, vertices, faces, colors):
+    with open(file, 'w') as file:
+        for v, c in zip(vertices, colors):
+            line = 'v ' + ' '.join([str(_) for _ in v]) + ' ' + ' '.join([str(_) for _ in c]) + '\n'
+            file.write(line)
+        for f in faces:
+            line = 'f ' + ' '.join([str(_) for _ in f]) + '\n'
+            file.write(line)