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, nr_labels in zip(neighbourhoods, nr_of_labels): 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_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, nr_of_labels): 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, nr_of_labels) 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)] ]) colors = np.ones((verts.shape[0],3))*np.array([255,200,200])[None,:] 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, lighting={ 'ambient':0.9, 'diffuse':0.5, 'roughness':0.5 } ), 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, # lighting={ 'ambient':0.8, 'diffuse':0.5, 'roughness':0.5 } ), 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(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],),) # ) 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=[25, 25, 25], nr_of_labels=[3,6,12])