230 lines
8.3 KiB
Python
230 lines
8.3 KiB
Python
import os
|
|
import sys
|
|
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 custom_utils.supersmpl_io import readPC2, readPC2Frame, writePC2, readOBJ
|
|
from custom_utils.custom_utils import _maskFaces
|
|
|
|
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(folder, frame, bins, neighbourhoods=[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)
|
|
|
|
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_local = bins[neighbourhood] if bins is not None else get_bins(curvature, nr_of_labels)
|
|
labels = np.linspace(0, bins_local.shape[0] - 2, bins_local.shape[0] - 1)
|
|
|
|
frame_labels = np.zeros((curvature.shape[0]), dtype=np.float32)
|
|
for j in range(bins_local.shape[0] - 2):
|
|
frame_labels[(curvature >= bins_local[j]) * (curvature < bins_local[j+1])] = labels[j]
|
|
frame_labels[curvature >= bins_local[-2]] = labels[-1]
|
|
|
|
curvatures.append(frame_labels)
|
|
|
|
return curvatures
|
|
|
|
|
|
def visualize_curvature_labels(folder, frame, bins, neighbourhoods):
|
|
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, bins, neighbourhoods)
|
|
colors_binned_all = []
|
|
|
|
for curvature_labels in curvature_labels_all:
|
|
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=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(
|
|
# 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=np.ones((verts.shape[0],3))*200
|
|
),
|
|
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.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(nticks=4, range=[-1,1],),
|
|
yaxis = dict(nticks=4, range=[-1,1],),
|
|
zaxis = dict(nticks=4, range=[-1,1],),)
|
|
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],),)
|
|
# )
|
|
|
|
fig.show()
|
|
|
|
print()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
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/valid_skirt'
|
|
bins = dict()
|
|
bins[10] = np.load(f'custom_data/train_bins_skirt_10.npy')
|
|
bins[25] = np.load(f'custom_data/train_bins_skirt.npy')
|
|
bins[100] = np.load(f'custom_data/train_bins_skirt_100.npy')
|
|
|
|
bins = None
|
|
|
|
|
|
nr_of_labels = 6
|
|
neighbourhood = 50
|
|
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')
|
|
|
|
visualize_curvature_labels(folder, '02572_1', bins, neighbourhoods=[10, 25, 100]) |