381 lines
15 KiB
Python
381 lines
15 KiB
Python
import os
|
|
import sys
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = '4'
|
|
|
|
|
|
import tqdm
|
|
import argparse
|
|
import random
|
|
import string
|
|
import time
|
|
import numpy as np
|
|
import trimesh as tm
|
|
import cmapy as cm
|
|
import scipy as sp
|
|
import scipy.io as sio
|
|
import plotly.graph_objects as go
|
|
from pyntcloud import PyntCloud
|
|
from scipy.sparse import csr_matrix, coo_matrix
|
|
from scipy.sparse.csgraph import floyd_warshall, dijkstra
|
|
from plotly.subplots import make_subplots
|
|
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, write_obj
|
|
from auto_encoder_curv import dice_dense
|
|
|
|
|
|
|
|
|
|
|
|
def postprocess(mesh, faces, mask=None, body=None, do_smoothing=True):
|
|
if body is not None:
|
|
mesh_fixed = fix_collision(mesh, body[np.newaxis,:,:], mask[np.newaxis,:])[0].numpy()[mask, :]
|
|
else:
|
|
mesh_fixed = mesh
|
|
|
|
if do_smoothing:
|
|
m = tm.Trimesh(mesh_fixed, faces, process=False, validate=False)
|
|
outliers = np.array([len(v) for v in m.vertex_neighbors]) == 0
|
|
outliers = np.nonzero(outliers)[0]
|
|
outliers_n_neighbours = get_closest_n_points(mesh_fixed, outliers, 2)
|
|
outliers_faces = np.concatenate((outliers[:, np.newaxis], outliers_n_neighbours), axis=-1)
|
|
faces_fixed = np.concatenate((faces, outliers_faces), axis=0)
|
|
m = tm.Trimesh(mesh_fixed, faces_fixed, process=False, validate=False)
|
|
tm.smoothing.filter_humphrey(m, iterations=1, alpha=0.6, beta=0.4)
|
|
|
|
mesh = m.vertices
|
|
else:
|
|
mesh = mesh_fixed
|
|
|
|
return mesh
|
|
|
|
def export_to_ply(mesh, file):
|
|
ply = tm.exchange.ply.export_ply(mesh, encoding='binary', vertex_normal=True, include_attributes=True)
|
|
output_file = open(file, "wb+")
|
|
output_file.write(ply)
|
|
output_file.close()
|
|
|
|
|
|
def get_closest_n_points(verts, idx, n):
|
|
dist = sp.spatial.distance.cdist(verts[idx, :], verts)
|
|
dist[dist == 0] = 1e9
|
|
idx_sort = np.argsort(dist, axis=-1)
|
|
closest = idx_sort[:, :n]
|
|
return closest
|
|
|
|
|
|
def calculate_dist(folder, neighbourhood):
|
|
out = dict()
|
|
frames = sorted(os.listdir(os.path.join(folder, 'mesh_posed')))[int(args.start):int(args.end)]
|
|
|
|
for frame in tqdm.tqdm(frames):
|
|
id = frame.split('_')[0]
|
|
if id in out.keys():
|
|
continue
|
|
|
|
mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy')).astype(bool)
|
|
k_neighbors = np.argsort(dist[mask, :][:, mask], axis=-1)[:, :neighbourhood] #cloud.get_neighbors(k=neighbourhood)
|
|
out[id] = k_neighbors
|
|
|
|
return out
|
|
|
|
|
|
def histedges_equalN(x, nbin):
|
|
npt = len(x)
|
|
return np.interp(np.linspace(0, npt, nbin + 1),
|
|
np.arange(npt),
|
|
np.sort(x))
|
|
|
|
|
|
def get_bins(values, nr_of_bins):
|
|
flattened = []
|
|
for c in values:
|
|
flattened.extend(c.flatten())
|
|
flattened = np.array(flattened)
|
|
bins = histedges_equalN(flattened, nr_of_bins)
|
|
return bins
|
|
|
|
|
|
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,:]
|
|
|
|
m = tm.Trimesh(v, faces, process=False, validate=False)
|
|
|
|
outliers = np.array([len(v) for v in m.vertex_neighbors]) == 0
|
|
outliers = np.nonzero(outliers)[0]
|
|
outliers_n_neighbours = get_closest_n_points(v, outliers, 2)
|
|
outliers_faces = np.concatenate((outliers[:, np.newaxis], outliers_n_neighbours), axis=-1)
|
|
faces = np.concatenate((faces, outliers_faces), axis=0)
|
|
|
|
verts = v
|
|
|
|
m = tm.Trimesh(verts, faces, process=False, validate=False)
|
|
export_to_ply(m, temp_file)
|
|
|
|
cloud = PyntCloud.from_file(temp_file)
|
|
|
|
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 = np.load(f'custom_data/train_bins.npy') # get_bins(curvature, nr_of_labels)
|
|
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]
|
|
|
|
return frame_labels
|
|
|
|
import tensorflow as tf
|
|
|
|
|
|
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'))
|
|
pred = postprocess(pred, faces)
|
|
|
|
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,:])
|
|
dice = dice_dense(tf.one_hot(curvature_labels_pred, 7)[None,:], tf.one_hot(curvature_labels_gt, 7)[None,:]).numpy()
|
|
colors_binned_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):
|
|
# 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=-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=np.ones((verts.shape[0],3))*np.array([255,200,200])[None,:]
|
|
),
|
|
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[0]
|
|
),
|
|
row=1, col=1)
|
|
|
|
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),
|
|
)
|
|
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,
|
|
)
|
|
|
|
|
|
# meshes = [go.Mesh3d(
|
|
# x=x[:,0],
|
|
# y=x[:,1],
|
|
# z=x[:,2],
|
|
# i = faces[:,0],
|
|
# j = faces[:,1],
|
|
# k = faces[:,2],
|
|
# vertexcolor=colors_binned
|
|
# )]
|
|
|
|
# fig = go.Figure(data=meshes)
|
|
# fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1))
|
|
# fig.update_layout(
|
|
# scene = dict(
|
|
# xaxis = dict(nticks=4, range=[-1,1],),
|
|
# yaxis = dict(nticks=4, range=[-1,1],),
|
|
# zaxis = dict(nticks=4, range=[-1,1],),)
|
|
# )
|
|
|
|
return fig, dice
|
|
|
|
import wandb
|
|
|
|
def init_wandb():
|
|
run = wandb.init(project='plots', name='curvatures', dir='/home/workspace/local-data/temp/')
|
|
artifact = wandb.Artifact('code', type='code')
|
|
artifact.add_file('curvature_analysis.py')
|
|
run.log_artifact(artifact)
|
|
|
|
|
|
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_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'
|
|
# with open(os.path.join(folder, frames_file), 'r') as f:
|
|
# lines = f.readlines()
|
|
# lines = [line.strip() for line in lines]
|
|
# frames = lines
|
|
frames = os.listdir(os.path.join(folder, 'pose_shape'))
|
|
frames = [frame.split('.')[0] for frame in frames]
|
|
|
|
temp_file = '/home/workspace/local-data/temp/' + ''.join(random.choice(string.ascii_lowercase) for i in range(10)) + '.ply'
|
|
if not is_skirt:
|
|
dist = np.load('custom_data/dijkstra.npy')
|
|
else:
|
|
dist = np.load('custom_data/dijkstra_skirt.npy')
|
|
|
|
i = 0
|
|
while True:
|
|
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)
|
|
visualize_neighbourhoods(folder, frame)
|
|
i += 1
|
|
print(frame) |