plot prediction curvatures

This commit is contained in:
2023-01-28 21:27:00 +00:00
parent 19cd69295d
commit 408b984b4e
+36 -26
View File
@@ -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,8 +91,6 @@ 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)
@@ -106,23 +105,24 @@ def generate_curvatures(folder, frame, neighbourhoods=[25], nr_of_labels=6):
frame_labels[(curvature >= bins[j]) * (curvature < bins[j+1])] = labels[j]
frame_labels[curvature >= bins[-2]] = labels[-1]
curvatures.append(frame_labels)
return curvatures
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,:]
pred = np.load(os.path.join(results_folder, frame + '.npy'))
curvature_labels_all = generate_curvatures(folder, frame, neighbourhoods)
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,7 +145,6 @@ 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],
@@ -156,17 +154,29 @@ def visualize_curvature_labels(folder, frame, neighbourhoods):
i = faces[:,0],
j = faces[:,1],
k = faces[:,2],
vertexcolor=colors_binned
vertexcolor=colors_binned_all[0]
),
row=1, col=i+2)
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])
visualize_curvature_labels(folder, results_folder, '00100_9')