import os import sys import tqdm import argparse import random import string import time import networkx as nx import tables as tb 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 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') from data.cloth3d_io import readPC2, readPC2Frame, writePC2, readOBJ, read_garment_mask, read_posed_garment_frame from custom_utils.custom_utils import _maskFaces # from FeatureLines.superSMPL.smpl import SMPL def get_body(): body = readOBJ('custom_data/superSMPL.obj') return body[0] def get_skirt_body(): body = readOBJ('custom_data/superSMPL_skirt.obj') return body[0] 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 histedges_equalN(x, nbin): npt = len(x) return np.interp(np.linspace(0, npt, nbin + 1), np.arange(npt), np.sort(x)) def k_neighbours(k): if is_skirt: body = get_skirt_body() else: body = get_body() adj = np.ones((14475, 14475)) * 9999 # ??? diagonal = np.linspace(0, 14474, 14475, dtype=np.int32) m = tm.Trimesh(body, F, process=False, validate=False) edges = m.edges_unique lengths = m.edges_unique_length adj[edges[:, 0], edges[:, 1]] = lengths adj[edges[:, 1], edges[:, 0]] = lengths adj[diagonal, diagonal] = 0 # G = nx.convert_matrix.from_numpy_array(adj) adj = csr_matrix(adj) dist = dijkstra(adj, directed=False, unweighted=False, limit=0.1) skirt_modifier = '_skirt' if is_skirt else '' np.save('custom_data/dijkstra' + skirt_modifier + '.npy', dist) # dist = floydWarshall(adj, 14475) # idx_sort = np.argsort(dist, axis=-1) # closest = idx_sort[:, :k] pass class Data: def __init__(self): self.k_neighbors = None self.mask = None self.faces = None class Cache: def __init__(self, folder, neighbourhood, frames, include_neighbours=True): self.cache = dict() self.folder = folder self.neighbourhood = neighbourhood self.frames = frames self.include_neighbours = include_neighbours def read(self, main_id): # frames = [frame for frame in self.frames if (frame.split()[0] == main_id.split()[0] and frame.split()[2] == main_id.split()[1])] # for frame in frames: id = main_id.split()[0] garment = main_id.split()[1] self.cache[f'{id} {garment}'] = Data() is_skirt = (garment == 'Dress' or garment == 'Skirt') mask = read_garment_mask(self.folder, id, garment) _, faces = _maskFaces(np.zeros((14475,3)), F if not is_skirt else F_skirt, mask) v = read_posed_garment_frame(folder, id, garment, 0)[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) # sorted_faces = [sorted(f) for f in faces.astype(np.int64)] # for outlier_face in outliers_faces: # assert sorted(outlier_face) not in sorted_faces faces = np.concatenate((faces, outliers_faces), axis=0) m = tm.Trimesh(v, faces, process=False, validate=False) if self.include_neighbours: if not is_skirt: k_neighbors = np.argsort(dist[mask, :][:, mask], axis=-1)[:, :self.neighbourhood] else: k_neighbors = np.argsort(dist_skirt[mask, :][:, mask], axis=-1)[:, :self.neighbourhood] self.cache[f'{id} {garment}'].k_neighbors = k_neighbors G = tm.graph.vertex_adjacency_graph(m) for i in range(v.shape[0]): for k in range(25): assert nx.has_path(G, i, k_neighbors[i, k]), f'No path between {i} and {k_neighbors[i, k]}' self.cache[f'{id} {garment}'].mask = mask self.cache[f'{id} {garment}'].faces = faces def __getitem__(self, id): if id not in self.cache.keys(): del self.cache self.cache = dict() self.read(id) return self.cache[id] def calculate_dist(folder, neighbourhood, frames): dist = dict() # frames = frames[int(args.start):int(args.end)] for frame in tqdm.tqdm(frames): id = frame.split()[0] garment = frame.split()[2] if id in dist.keys(): continue mask = read_garment_mask(folder, id, garment) #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) dist[id] = k_neighbors return dist def generate_curvatures(folder, neighbourhood, frames): # ids = os.listdir(folder)[int(args.start):int(args.end)] # frames = frames[int(args.start):int(args.end)] os.makedirs(os.path.join(folder + '_extras', f'mesh_curvature_raw'), exist_ok=True) cache = Cache(folder, neighbourhood, frames) # filters = tb.Filters(complevel=9, complib='blosc2') for frame in tqdm.tqdm(frames): id = frame.split()[0] frame_nr = int(frame.split()[1]) garment = frame.split()[2] mask = cache[f'{id} {garment}'].mask faces = cache[f'{id} {garment}'].faces v = read_posed_garment_frame(folder, id, garment, frame_nr)[mask,:] m = tm.Trimesh(v, faces, process=False, validate=False) export_to_ply(m, temp_file) cloud = PyntCloud.from_file(temp_file) k_neighbors = cache[f'{id} {garment}'].k_neighbors 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) + ")"]) out = np.zeros((14475), dtype=np.float32) out[mask] = curvature.real if str(curvature.dtype) == 'complex64' else curvature np.savez_compressed(os.path.join(folder + '_extras', f'mesh_curvature_raw', f'{id}_{frame_nr}_{garment}.npz'), mesh_curvature_raw=out) # group = h5file.get_node(f'/{id}/{frame_nr}/{garment}') # if f'curvature_raw' in group: # array = h5file.get_node(f'/{id}/{frame_nr}/{garment}/curvature_raw') # array[:] = out # else: # carray = h5file.create_carray(group, "curvature_raw", atom=tb.Float32Atom(), shape=out.shape, filters=filters) # carray[:] = out return None 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_curvature_labels(folder, bins, curvatures, frames, neighbourhood): cache = Cache(folder, neighbourhood, frames, include_neighbours=False) os.makedirs(os.path.join(folder + '_extras', f'curvature_mesh_global'), exist_ok=True) os.makedirs(os.path.join(folder + '_extras', f'curvature_mesh_garment'), exist_ok=True) labels = np.linspace(0, bins.shape[0] - 2, bins.shape[0] - 1) # filters = tb.Filters(complevel=9, complib='blosc2') for i, frame in enumerate(tqdm.tqdm(frames)): id = frame.split()[0] frame_nr = int(frame.split()[1]) garment = frame.split()[2] curvature = curvatures[i][cache[f'{id} {garment}'].mask] frame_labels = np.zeros((curvature.shape[0]), dtype=np.uint8) 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] out = np.zeros((14475), dtype=np.uint8) out[cache[f'{id} {garment}'].mask] = frame_labels + 1 # group = h5file.get_node(f'/{id}/{frame_nr}/{garment}') # if f'curvature_mesh_global' in group: # array = h5file.get_node(f'/{id}/{frame_nr}/{garment}/curvature_mesh_global') # array[:] = out # else: # carray = h5file.create_carray(group, "curvature_mesh_global", atom=tb.UInt8Atom(), shape=out.shape, filters=filters) # carray[:] = out np.savez_compressed(os.path.join(folder + '_extras', f'curvature_mesh_global', f'{id}_{frame_nr}_{garment}.npz'), mesh_curvature_raw=out) # writePC2(os.path.join(garment_dir, "curvature_label_" + str(nr_of_labels) + "_" + str(neighbourhood) + ".pc16"), unmasked_label, True) # redo it with local bins bins = get_bins([curvature], nr_of_labels) frame_labels = np.zeros((curvature.shape[0]), dtype=np.uint8) 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] out = np.zeros((14475), dtype=np.uint8) out[cache[f'{id} {garment}'].mask] = frame_labels + 1 # group = h5file.get_node(f'/{id}/{frame_nr}/{garment}') # if f'curvature_mesh_garment' in group: # array = h5file.get_node(f'/{id}/{frame_nr}/{garment}/curvature_mesh_garment') # array[:] = out # else: # carray = h5file.create_carray(group, "curvature_mesh_garment", atom=tb.UInt8Atom(), shape=out.shape, filters=filters) # carray[:] = out np.savez_compressed(os.path.join(folder + '_extras', f'curvature_mesh_garment', f'{id}_{frame_nr}_{garment}.npz'), mesh_curvature_raw=out) def visualize_curvature_labels(folder, frames): cache = Cache(folder, 0, frames, include_neighbours=False) for frame in tqdm.tqdm(frames): id = frame.split()[0] frame_nr = int(frame.split()[1]) garment = frame.split()[2] mask = cache[id].mask faces = cache[id].faces verts = read_posed_garment_frame(folder, id, garment, frame_nr)[mask, :] curv = h5file.get_node(f'/{id}/{frame_nr}/{garment}/curvature_mesh_global').read()[mask] colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, int(curv.max()))]) colors_binned = np.zeros((curv.shape[0], 3)) for i in range(1, int(curv.max())+1): # print(np.count_nonzero(curv == i)) colors_binned[curv == i, :] = colors[i-1] x = verts meshes = [go.Mesh3d( x=x[:,0], y=x[:,2], z=x[:,1], 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() def read_raw_curvatures(folder, frames, neighbourhood): curvatures_all = [] for frame in tqdm.tqdm(frames): id = frame.split()[0] frame_nr = int(frame.split()[1]) garment = frame.split()[2] curv = h5file.get_node(f'/{id}/{frame_nr}/{garment}/curvature_raw').read() curvatures_all.append(curv) return curvatures_all def read_raw_curvatures_2(folder, frames, neighbourhood): curvatures_all = [] for frame in tqdm.tqdm(frames): id = frame.split()[0] frame_nr = int(frame.split()[1]) garment = frame.split()[2] # curv = np.load(os.path.join(folder, f'mesh_curvature_raw_{neighbourhood}', f'{frame}.npz'))['mesh_curvature_raw'] curv = np.load(os.path.join('/'.join(folder.split('/')[:-1]), 'backup', f'mesh_curvature_raw', f'{id}_{frame_nr}_{garment}.npz'))['mesh_curvature_raw'] curvatures_all.append(curv) return curvatures_all if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument("-f", "--folder", help = "process the given folder", default=r'/home/workspace/data/dataset/cloth3d/train') parser.add_argument("-s", "--start", help = "start of range", default=300) parser.add_argument("-e", "--end", help = "end of range", default=600) parser.add_argument("-n", "--name", help = "end of range", default="lhunor") args = parser.parse_args() F = np.load("custom_utils/smpl_faces.npy").astype(np.int32) F_skirt = np.load("custom_utils/smpl_faces_skirt.npy").astype(np.int32) # k_neighbours(k=25) folder = args.folder frames_file = 'frames_all.txt' nr_of_labels = 6 neighbourhood = 25 temp_file = '/home/workspace/local-data/temp/' + ''.join(random.choice(string.ascii_lowercase) for i in range(10)) + '.ply' dist = np.load('custom_data/dijkstra.npy') dist_skirt = np.load('custom_data/dijkstra_skirt.npy') print(f'params: -f {args.folder} -s {args.start} -e {args.end} -n {args.name} -neigh {neighbourhood}') # frames = sorted(os.listdir(os.path.join(folder, 'mesh_posed')))[int(args.start):int(args.end)] # frames = [frame.split('.')[0] for frame in frames] frames = [] with open(os.path.join(folder, frames_file), 'r') as f: lines = f.readlines() lines = [line.strip() for line in lines] frames = lines[int(args.start):int(args.end)] root = '/'.join(folder.split('/')[:-1]) type = folder.split('/')[-1] # h5file = tb.open_file(os.path.join(root, f'{type}_extras.h5'), mode='a') start = time.perf_counter() # frames = [f'05235 {i} Tshirt' for i in range(300)] generate_curvatures(folder, neighbourhood, frames) # curvatures = read_raw_curvatures_2(folder, frames, neighbourhood) # bins = get_bins(curvatures, nr_of_labels) # np.save(f'custom_data/{type}_bins{skirt_modifier}_{neighbourhood}.npy', bins) # print(bins) # bins = np.load(f'custom_data/{type}_bins_{neighbourhood}.npy') # generate_curvature_labels(folder, bins, curvatures, frames, neighbourhood) # visualize_curvature_labels(folder, frames) print(time.perf_counter() - start) pass