395 lines
15 KiB
Python
395 lines
15 KiB
Python
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 |