diff --git a/curvature_analysis.py b/curvature_analysis.py index c3aea26..ba861cb 100644 --- a/curvature_analysis.py +++ b/curvature_analysis.py @@ -19,6 +19,7 @@ from custom_utils.custom_utils import _maskFaces 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 export_to_ply(mesh, file): ply = tm.exchange.ply.export_ply(mesh, encoding='binary', vertex_normal=True, include_attributes=True) @@ -67,13 +68,13 @@ def get_bins(values, nr_of_bins): return bins -def generate_curvatures(folder, frame, neighbourhoods=[25], nr_of_labels=6): - id = frame.split('_')[0] - frame_nr = int(frame.split('_')[1].split('.')[0]) +def generate_curvatures(v, faces, mask, neighbourhood=25, nr_of_labels=6): + # id = frame.split('_')[0] + # frame_nr = int(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) - v = np.load(os.path.join(folder, 'mesh_posed', f'{id}_{frame_nr}.npz'))['mesh_posed'][mask,:] + # mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy')).astype(bool) + # _, faces = _maskFaces(np.zeros((14475,3)), F, mask) + # v = np.load(os.path.join(folder, 'mesh_posed', f'{id}_{frame_nr}.npz'))['mesh_posed'][mask,:] m = tm.Trimesh(v, faces, process=False, validate=False) @@ -90,39 +91,38 @@ def generate_curvatures(folder, frame, neighbourhoods=[25], nr_of_labels=6): cloud = PyntCloud.from_file(temp_file) - curvatures = [] - for neighbourhood in neighbourhoods: - k_neighbors = np.argsort(dist[mask, :][:, mask], axis=-1)[:, :neighbourhood] #cloud.get_neighbors(k=neighbourhood) - #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 = get_bins(curvature, nr_of_labels) - labels = np.linspace(0, bins.shape[0] - 2, bins.shape[0] - 1) + k_neighbors = np.argsort(dist[mask, :][:, mask], axis=-1)[:, :neighbourhood] #cloud.get_neighbors(k=neighbourhood) + #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 = 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) - for j in range(bins.shape[0] - 2): - frame_labels[(curvature >= bins[j]) * (curvature < bins[j+1])] = labels[j] - frame_labels[curvature >= bins[-2]] = labels[-1] - - curvatures.append(frame_labels) - - return curvatures + frame_labels = np.zeros((curvature.shape[0]), dtype=np.float32) + for j in range(bins.shape[0] - 2): + frame_labels[(curvature >= bins[j]) * (curvature < bins[j+1])] = labels[j] + frame_labels[curvature >= bins[-2]] = labels[-1] + + return frame_labels -def visualize_curvature_labels(folder, frame, neighbourhoods): +def visualize_curvature_labels(folder, results_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_all = generate_curvatures(folder, frame, neighbourhoods) + pred = np.load(os.path.join(results_folder, frame + '.npy')) + + 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,:]) colors_binned_all = [] - for curvature_labels in curvature_labels_all: + for curvature_labels in [curvature_labels_gt, curvature_labels_pred]: 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): @@ -130,8 +130,7 @@ def visualize_curvature_labels(folder, frame, neighbourhoods): colors_binned[curvature_labels == i, :] = colors[i] colors_binned_all.append(colors_binned) - - fig = make_subplots(rows=1, cols=len(colors_binned_all)+1, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(len(colors_binned_all)+1)] ]) + 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 @@ -146,27 +145,38 @@ def visualize_curvature_labels(folder, frame, neighbourhoods): ), row=1, col=1) - for i, colors_binned in enumerate(colors_binned_all): - 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 - ), - row=1, col=i+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[0] + ), + row=1, col=2) - + fig.add_trace(go.Mesh3d( + # 8 vertices of a cube + 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=colors_binned_all[1], + text=f'dice: {dice}' + ), + 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), + # 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],), @@ -174,8 +184,8 @@ def visualize_curvature_labels(folder, frame, neighbourhoods): fig.update_layout(scene = helper, scene2 = helper, scene3 = helper, - scene4 = helper, - scene5 = helper, + # scene4 = helper, + # scene5 = helper, ) @@ -211,7 +221,7 @@ if __name__ == '__main__': # k_neighbours(k=25) folder = '/home/workspace/local-data/train' - results_folder = '' + results_folder = '/home/workspace/data/results/post/abl_overfit_jumpsuit_2-' frames_file = 'frames_jumpsuit_30k.txt' frames = [] @@ -226,4 +236,4 @@ if __name__ == '__main__': else: dist = np.load('custom_data/dijkstra_skirt.npy') - visualize_curvature_labels(folder, '00005_1', neighbourhoods=[10, 25, 50, 100]) \ No newline at end of file + visualize_curvature_labels(folder, results_folder, '00100_9') \ No newline at end of file