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.supersmpl_io import readPC2, readPC2Frame, writePC2, readOBJ
from custom_utils.custom_utils import _maskFaces from custom_utils.custom_utils import _maskFaces
from auto_encoder_curv import dice_dense
def export_to_ply(mesh, file): def export_to_ply(mesh, file):
ply = tm.exchange.ply.export_ply(mesh, encoding='binary', vertex_normal=True, include_attributes=True) 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 return bins
def generate_curvatures(folder, frame, neighbourhoods=[25], nr_of_labels=6): def generate_curvatures(v, faces, mask, neighbourhood=25, nr_of_labels=6):
id = frame.split('_')[0] # id = frame.split('_')[0]
frame_nr = int(frame.split('_')[1].split('.')[0]) # frame_nr = int(frame.split('_')[1].split('.')[0])
mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy')).astype(bool) # mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy')).astype(bool)
_, faces = _maskFaces(np.zeros((14475,3)), F, mask) # _, 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,:] # 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) 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) 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 = 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) #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) 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[j]) * (curvature < bins[j+1])] = labels[j]
frame_labels[curvature >= bins[-2]] = labels[-1] frame_labels[curvature >= bins[-2]] = labels[-1]
curvatures.append(frame_labels) return frame_labels
return curvatures
def visualize_curvature_labels(folder, frame, neighbourhoods): def visualize_curvature_labels(folder, results_folder, frame):
id = frame.split('_')[0] id = frame.split('_')[0]
frame_nr = frame.split('_')[1].split('.')[0] frame_nr = frame.split('_')[1].split('.')[0]
mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy')).astype(bool) mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy')).astype(bool)
_, faces = _maskFaces(np.zeros((14475,3)), F, mask) _, 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,:] 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 = [] 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 = 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)) colors_binned = np.zeros((curvature_labels.shape[0], 3))
for i in range(int(curvature_labels.max())+1): 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[curvature_labels == i, :] = colors[i]
colors_binned_all.append(colors_binned) 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 = 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.add_trace(go.Mesh3d( fig.add_trace(go.Mesh3d(
# 8 vertices of a cube # 8 vertices of a cube
@@ -146,7 +145,6 @@ def visualize_curvature_labels(folder, frame, neighbourhoods):
), ),
row=1, col=1) row=1, col=1)
for i, colors_binned in enumerate(colors_binned_all):
fig.add_trace(go.Mesh3d( fig.add_trace(go.Mesh3d(
# 8 vertices of a cube # 8 vertices of a cube
x=verts[:,0], x=verts[:,0],
@@ -156,17 +154,29 @@ def visualize_curvature_labels(folder, frame, neighbourhoods):
i = faces[:,0], i = faces[:,0],
j = faces[:,1], j = faces[:,1],
k = faces[:,2], 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), 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), scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1),
scene3_aspectmode='manual', scene3_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), # scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1),
scene5_aspectmode='manual', scene5_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],), helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
yaxis = 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, fig.update_layout(scene = helper,
scene2 = helper, scene2 = helper,
scene3 = helper, scene3 = helper,
scene4 = helper, # scene4 = helper,
scene5 = helper, # scene5 = helper,
) )
@@ -211,7 +221,7 @@ if __name__ == '__main__':
# k_neighbours(k=25) # k_neighbours(k=25)
folder = '/home/workspace/local-data/train' 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_file = 'frames_jumpsuit_30k.txt'
frames = [] frames = []
@@ -226,4 +236,4 @@ if __name__ == '__main__':
else: else:
dist = np.load('custom_data/dijkstra_skirt.npy') 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')