From 22f77aba6a2786887446770440f49696f81b1ece Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Fri, 27 Jan 2023 19:20:16 +0000 Subject: [PATCH] visualize curv with different radius --- curvature_analysis.py | 223 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 223 insertions(+) create mode 100644 curvature_analysis.py diff --git a/curvature_analysis.py b/curvature_analysis.py new file mode 100644 index 0000000..ee59ae3 --- /dev/null +++ b/curvature_analysis.py @@ -0,0 +1,223 @@ +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, 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 = 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] + + curvatures.append(frame_labels) + + return curvatures + + +def visualize_curvature_labels(folder, frame, 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, 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' + + 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', neighbourhoods=[10, 25, 50, 100]) \ No newline at end of file