From cbecf668dff4693537dc8a2b1aca233c116f95dc Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Tue, 7 Mar 2023 12:54:53 +0000 Subject: [PATCH] save --- curvature_analysis.py | 114 ++++++++++++++++++++++++++++++----- custom_utils/custom_utils.py | 12 +++- 2 files changed, 110 insertions(+), 16 deletions(-) diff --git a/curvature_analysis.py b/curvature_analysis.py index 6b08b86..ec1e7d3 100644 --- a/curvature_analysis.py +++ b/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) - wandb.log({f'{frame}-dice:{dice}': fig}, step=i) + # fig, dice = visualize_curvature_labels(folder, results_folder, frame) + # wandb.log({f'{frame}-dice:{dice}': fig}, step=i) + visualize_neighbourhoods(folder, frame) i += 1 print(frame) \ No newline at end of file diff --git a/custom_utils/custom_utils.py b/custom_utils/custom_utils.py index 1f180bf..a570a11 100644 --- a/custom_utils/custom_utils.py +++ b/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) \ No newline at end of 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) \ No newline at end of file