prepare for cape
This commit is contained in:
@@ -0,0 +1,271 @@
|
||||
import os
|
||||
import sys
|
||||
import tqdm
|
||||
import argparse
|
||||
import random
|
||||
import string
|
||||
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 custom_utils.supersmpl_io import readPC2, readPC2Frame, writePC2, readOBJ
|
||||
from custom_utils.custom_utils import _maskFaces
|
||||
# from FeatureLines.superSMPL.smpl import SMPL
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
|
||||
import tensorflow as tf
|
||||
gpus = tf.config.experimental.list_physical_devices('GPU')
|
||||
if gpus:
|
||||
try:
|
||||
for gpu in gpus:
|
||||
tf.config.experimental.set_memory_growth(gpu, True)
|
||||
|
||||
except RuntimeError as e:
|
||||
print(e)
|
||||
from custom_utils.smpl import SMPL
|
||||
|
||||
smpl = SMPL(orig_smpl=True)
|
||||
|
||||
|
||||
def get_skirt_body():
|
||||
body = readOBJ('FeatureLines/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):
|
||||
body = get_skirt_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
|
||||
|
||||
adj = csr_matrix(adj)
|
||||
dist = dijkstra(adj, directed=False, unweighted=False, limit=0.5)
|
||||
|
||||
np.save('MeshUtils/dijkstra_skirt.npy', dist)
|
||||
|
||||
# dist = floydWarshall(adj, 14475)
|
||||
|
||||
# idx_sort = np.argsort(dist, axis=-1)
|
||||
# closest = idx_sort[:, :k]
|
||||
|
||||
pass
|
||||
|
||||
|
||||
def generate_curvatures(folder, neighbourhood, save_folder):
|
||||
ids = os.listdir(folder)[int(args.start):int(args.end)]
|
||||
# gender_table = np.load('/home/workspace/data/dataset/cape_release/misc/subj_genders.pkl', allow_pickle=True)
|
||||
faces = np.load('custom_utils/smpl_faces_orig.npy')
|
||||
k_neighbors = np.load('custom_data/dijkstra_orig_argsort_25.npy')
|
||||
|
||||
os.makedirs(os.path.join(save_folder, 'curvatures_raw'), exist_ok=True)
|
||||
|
||||
for id in tqdm.tqdm(ids):
|
||||
garment_dir = os.path.join(folder, id)
|
||||
|
||||
sequences = os.listdir(garment_dir)
|
||||
|
||||
for sequence in sequences:
|
||||
frames = os.listdir(os.path.join(garment_dir, sequence))
|
||||
|
||||
for frame in frames:
|
||||
data = np.load(os.path.join(folder, id, sequence, frame))
|
||||
# pose = data['poses']
|
||||
# shape = np.load(os.path.join('/home/workspace/data/dataset/cape_release/minimal_body_shape', id, id + '_minimal.npy'))
|
||||
# gender = [0] if gender_table[id] == 'female' else [1]
|
||||
# shape = np.concatenate((shape, gender)).astype(np.float32)
|
||||
# body = smpl.run(smpl.run_rodrigues(pose[np.newaxis, :]), shape[np.newaxis,:], remove=False)[0][0]
|
||||
|
||||
verts = data['v_posed'][smpl.to_keep, :]
|
||||
|
||||
m = tm.Trimesh(verts, faces, process=False, validate=False)
|
||||
export_to_ply(m, temp_file)
|
||||
|
||||
cloud = PyntCloud.from_file(temp_file)
|
||||
# cloud.plot(mesh=True)
|
||||
|
||||
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) + ")"])
|
||||
|
||||
frame_name = frame.split('.')
|
||||
assert frame_name[-1] == 'npz', 'It\'s not an npz file!'
|
||||
frame_name = '.'.join(frame_name[:-1])
|
||||
|
||||
np.save(os.path.join(save_folder, 'curvatures_raw', id + '_' + frame_name + '.npy'), np.array(curvature))
|
||||
|
||||
print(id + '_' + frame)
|
||||
|
||||
|
||||
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):
|
||||
os.makedirs(os.path.join(save_folder, 'mesh_curvature'), exist_ok=True)
|
||||
|
||||
frames = os.listdir(folder)#[int(args.start):int(args.end)]
|
||||
labels = np.linspace(0, bins.shape[0] - 2, bins.shape[0] - 1)
|
||||
|
||||
|
||||
for frame in tqdm.tqdm(frames):
|
||||
curvature = np.load(os.path.join(folder, frame))
|
||||
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]
|
||||
|
||||
frame_name = frame.split('.')
|
||||
assert frame_name[-1] == 'npy', 'It\'s not an npy file!'
|
||||
frame_name = '.'.join(frame_name[:-1])
|
||||
np.savez_compressed(os.path.join(save_folder, 'mesh_curvature', frame_name + '.npz'), mesh_curvature=frame_labels)
|
||||
|
||||
# ids = os.listdir(folder)#[:10]
|
||||
# labels = np.linspace(0, bins.shape[0] - 2, bins.shape[0] - 1)
|
||||
|
||||
# i = 0
|
||||
# for id in tqdm.tqdm(ids):
|
||||
# garment_dir = os.path.join(folder, id)
|
||||
|
||||
|
||||
# curvature = curvatures[i]
|
||||
# i += 1
|
||||
|
||||
# curvature_label = np.zeros((curvature.shape[0], curvature.shape[1]))
|
||||
# frame = 0
|
||||
# for frame in range(curvature.shape[0]):
|
||||
# frame_labels = np.zeros((curvature.shape[1]), dtype=np.float32)
|
||||
# for j in range(bins.shape[0] - 2):
|
||||
# frame_labels[(curvature[frame, :] >= bins[j]) * (curvature[frame, :] < bins[j+1])] = labels[j]
|
||||
# frame_labels[curvature[frame, :] >= bins[-2]] = labels[-1]
|
||||
# curvature_label[frame, :] = frame_labels
|
||||
# frame += 1
|
||||
|
||||
# unmasked_label = np.zeros((curvature.shape[0], 14475), dtype=np.float32)
|
||||
# unmasked_label[:, mask] = curvature_label
|
||||
# np.save(os.path.join(garment_dir, "curvature_label_" + str(nr_of_labels) + "_" + str(neighbourhood) + ".npy"), unmasked_label)
|
||||
# writePC2(os.path.join(garment_dir, "curvature_label_" + str(nr_of_labels) + "_" + str(neighbourhood) + ".pc16"), unmasked_label, True)
|
||||
|
||||
|
||||
def visualize_curvature_labels(folder):
|
||||
frames = os.listdir(os.path.join(folder, 'mesh_curvature'))#[:10]
|
||||
faces = np.load('custom_utils/smpl_faces_orig.npy')
|
||||
|
||||
for _ in range(10):
|
||||
|
||||
frame = random.choice(frames)
|
||||
mesh_posed = np.load(os.path.join(folder, 'mesh_posed', frame)); mesh_posed = mesh_posed['mesh_posed']
|
||||
mesh_curvature = np.load(os.path.join(folder, 'mesh_curvature', frame)); mesh_curvature = mesh_curvature['mesh_curvature']
|
||||
|
||||
colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, int(mesh_curvature.max()+1))])
|
||||
colors_binned = np.zeros((mesh_curvature.shape[0], 3))
|
||||
for i in range(int(mesh_curvature.max())+1):
|
||||
print(np.count_nonzero(mesh_curvature == i))
|
||||
colors_binned[mesh_curvature == i, :] = colors[i]
|
||||
|
||||
x = mesh_posed
|
||||
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.5,1.5],),
|
||||
yaxis = dict(nticks=4, range=[-1.5,1.5],),
|
||||
zaxis = dict(nticks=4, range=[-1.5,1.5],),)
|
||||
)
|
||||
|
||||
fig.show()
|
||||
|
||||
print()
|
||||
|
||||
|
||||
def read_raw_curvatures(folder):
|
||||
|
||||
frames = os.listdir(folder)#[int(args.start):int(args.end)]
|
||||
|
||||
curvatures_all = []
|
||||
|
||||
for frame in tqdm.tqdm(frames):
|
||||
curv = np.load(os.path.join(folder, frame))
|
||||
curvatures_all.append(curv)
|
||||
|
||||
return curvatures_all
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("-if", "--in_folder", help = "process the given folder", default=r'/home/workspace/local-data/cape_processed/test/male/curvatures_raw')
|
||||
parser.add_argument("-of", "--out_folder", help = "process the given folder", default=r'/home/workspace/local-data/cape_processed/test/male')
|
||||
parser.add_argument("-s", "--start", help = "start of range", default=0)
|
||||
parser.add_argument("-e", "--end", help = "end of range", default=1500)
|
||||
parser.add_argument("-n", "--name", help = "end of range", default="lhunor")
|
||||
args = parser.parse_args()
|
||||
|
||||
F = np.load("custom_utils/smpl_faces_orig.npy")
|
||||
# k_neighbours(k=25)
|
||||
|
||||
folder = args.in_folder
|
||||
save_folder = args.out_folder #'/home/workspace/data/dataset/cape_processed'
|
||||
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 = np.load('MeshUtils/dijkstra_skirt.npy')
|
||||
|
||||
# generate_curvatures(folder, neighbourhood, save_folder)
|
||||
# curvatures = read_raw_curvatures(folder)
|
||||
# bins = get_bins(curvatures, nr_of_labels)
|
||||
# np.save('custom_data/cape_2_train_male_bins.npy', bins)
|
||||
bins = np.load('custom_data/cape_2_train_male_bins.npy')
|
||||
generate_curvature_labels(folder, bins)
|
||||
# visualize_curvature_labels(save_folder)
|
||||
|
||||
pass
|
||||
Reference in New Issue
Block a user