diff --git a/curvature_label_parallel.py b/curvature_label_parallel.py index c0b7be5..9a01f4f 100644 --- a/curvature_label_parallel.py +++ b/curvature_label_parallel.py @@ -4,15 +4,15 @@ import os if __name__ == '__main__': - folder = '/home/workspace/local-data/cape/train/male' - out_folder = '/home/workspace/local-data/cape_processed/train/male' - ids = os.listdir(folder) + folder = '/home/workspace/local-data/train' + out_folder = '/home/workspace/local-data/train' + ids = os.listdir(folder + '/mesh_posed') count = 0 - step = 1 + step = 15000 while count < len(ids): # os.system('nohup python MeshUtils/curvature_labels.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'curv_train_' + str(count) + '_' + str(count+step) + ' > curv_train_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &') - os.system('nohup python curvature_labels_cape.py -if ' + folder + ' -of ' + out_folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'curv_gen_2_cape_train_male_' + str(count) + '_' + str(count+step) + ' > nohup/curv_gen_2_cape_train_male_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &') + os.system('nohup python curvature_labels.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'curv_gen_50_train_nons_' + str(count) + '_' + str(count+step) + ' > nohup/curv_gen_50_train_nons_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &') # os.system('python Data/convert_to_uv.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + '') count += step diff --git a/curvature_labels.py b/curvature_labels.py index f99b262..c6b72ec 100644 --- a/curvature_labels.py +++ b/curvature_labels.py @@ -4,6 +4,7 @@ import tqdm import argparse import random import string +import time import numpy as np import trimesh as tm import cmapy as cm @@ -40,6 +41,8 @@ def export_to_ply(mesh, file): output_file.close() + + def get_closest_n_points(verts, idx, n): dist = sp.spatial.distance.cdist(verts[idx, :], verts) dist[dist == 0] = 1e9 @@ -80,35 +83,40 @@ def k_neighbours(k): pass +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 generate_curvatures(folder, neighbourhood): - ids = os.listdir(folder)[int(args.start):int(args.end)] + # ids = os.listdir(folder)[int(args.start):int(args.end)] + frames = sorted(os.listdir(os.path.join(folder, 'mesh_posed')))[int(args.start):int(args.end)] + os.makedirs(os.path.join(folder, f'mesh_curvature_raw_{neighbourhood}'), exist_ok=True) - curvatures_all = [] + k_neighbors_cache = calculate_dist(folder, neighbourhood) - for id in tqdm.tqdm(ids): - garment_dir = os.path.join(folder, id) + for frame in tqdm.tqdm(frames): - if not os.path.isdir(garment_dir): - continue - - # if os.path.exists(garment_dir + '/curvature_25_raw.npy'): - # continue - - curvatures_garment = [] - - garments = [f.replace('.pc16','') for f in os.listdir(garment_dir) if f.find('_') == -1 and f != 'info.mat' and not f.endswith('.obj')] - # assert len (garments) == 1 - if len (garments) != 1: - continue - garment = garments[0] + id = frame.split('_')[0] + frame_nr = frame.split('_')[1].split('.')[0] - bodies = readPC2(os.path.join(folder, id, 'Body_posed.pc16'), True)['V'] - offsetss = readPC2(os.path.join(folder, id, garment + '_registered.pc16'), True)['V'] - mask = np.squeeze(sio.loadmat(os.path.join(garment_dir, garment + '_mask.mat'))['mask'].astype(bool)) + # bodies = readPC2(os.path.join(folder, id, 'Body_posed.pc16'), True)['V'] + # offsetss = readPC2(os.path.join(folder, id, garment + '_registered.pc16'), True)['V'] + mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy')).astype(bool) + faces = np.load(os.path.join(folder, 'faces', f'{id}.npy')) + v = np.load(os.path.join(folder, 'mesh_posed', f'{id}_{frame_nr}.npz'))['mesh_posed'][mask,:] - _, faces = _maskFaces(bodies[0, :, :], F, mask) - v = bodies[0, mask, :] + offsetss[0, mask, :] m = tm.Trimesh(v, faces, process=False, validate=False) outliers = np.array([len(v) for v in m.vertex_neighbors]) == 0 @@ -117,30 +125,23 @@ def generate_curvatures(folder, neighbourhood): outliers_faces = np.concatenate((outliers[:, np.newaxis], outliers_n_neighbours), axis=-1) faces = np.concatenate((faces, outliers_faces), axis=0) - step = 0 - for body, offsets in zip(bodies, offsetss): - verts = body + offsets - verts = verts[mask, :] + verts = v - m = tm.Trimesh(verts, faces, process=False, validate=False) - export_to_ply(m, temp_file) + 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) + cloud = PyntCloud.from_file(temp_file) - k_neighbors = 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) + ")"]) - curvatures_garment.append(curvature) + 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) + ")"]) + out = np.zeros((14475), dtype=np.float32) + out[mask] = curvature + + np.savez_compressed(os.path.join(folder, f'mesh_curvature_raw_{neighbourhood}', f'{id}_{frame_nr}.npz'), mesh_curvature_raw=out) - print(step) - step += 1 - - # curvatures_all.append(np.array(curvatures_garment)) - np.save(garment_dir + '/curvature_25_raw.npy', np.array(curvatures_garment)) - - return curvatures_all + return None def get_bins(values, nr_of_bins): @@ -282,21 +283,27 @@ def read_raw_curvatures(folder): if __name__ == '__main__': parser = argparse.ArgumentParser() - parser.add_argument("-f", "--folder", help = "process the given folder", default=r'/home/workspace/data/dataset/test') + parser.add_argument("-f", "--folder", help = "process the given folder", default=r'/home/workspace/local-data/train') 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("-e", "--end", help = "end of range", default=10) parser.add_argument("-n", "--name", help = "end of range", default="lhunor") args = parser.parse_args() - F = np.load("FeatureLines/superSMPL/smpl_faces.npy") + is_skirt = False + 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 = args.folder nr_of_labels = 6 - neighbourhood = 25 - temp_file = 'temp/' + ''.join(random.choice(string.ascii_lowercase) for i in range(10)) + '.ply' - dist = np.load('MeshUtils/dijsktra.npy') - # dist = np.load('MeshUtils/dijkstra_skirt.npy') + 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') + + print('SKIRT' if is_skirt else 'NOT SKIRT') generate_curvatures(folder, neighbourhood) # curvatures = read_raw_curvatures(folder)