diff --git a/data/calculate_normals.py b/data/calculate_normals.py new file mode 100644 index 0000000..a978384 --- /dev/null +++ b/data/calculate_normals.py @@ -0,0 +1,112 @@ +import os +import sys +import numpy as np +import trimesh as tm +import scipy.io as sio +import plotly.graph_objects as go +import plotly.express as px +from tqdm import tqdm +import torch +from scipy.spatial.transform import Rotation as R +import argparse + +sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') + +from data.cloth3d_io import readOBJ, readPC2 +from data.utils import Frame +from smpl.smpl_torch import SMPL + + +def calculate_normals(vertices, mesh_file, poses): + _, faces, _, _ = readOBJ(mesh_file) + rodr = smpl.rodrigues(torch.tensor(poses.T.reshape(-1, 1, 3), dtype=torch.float32)).numpy().reshape(vertices.shape[0], -1, 3, 3) + + normals = [] + for i in range(vertices.shape[0]): + r = R.from_matrix(rodr[i, 0, :, :]) + # assert r.as_matrix() == rodr[frame, 0, :, :] + out = r.apply(vertices[i], inverse=True) + + # create a plotly mesh plot of out and vertices [i] + # fig = go.Figure(data=[ + # go.Mesh3d( + # x=out[:, 0], y=out[:, 2], z=out[:, 1], + # # i, j and k give the vertices of triangles + # i=faces[:, 0], j=faces[:, 1], k=faces[:, 2], + # name='y', + # showscale=False + # ), + # go.Mesh3d( + # x=vertices[i][:, 0], y=vertices[i][:, 1], z=vertices[i][:, 2], + # # i, j and k give the vertices of triangles + # i=faces[:, 0], j=faces[:, 1], k=faces[:, 2], + # name='x', + # showscale=False + # ) + # ]) + # fig.show() + + mesh = tm.Trimesh(vertices=out, faces=faces, process=False, validate=False) + normals.append(mesh.vertex_normals) + return np.stack(normals) + + +def get_frames(folder): + uv_static_folder = os.path.join(folder, 'pose_shape') + frames = sorted(os.listdir(uv_static_folder)) + + frames = [frame.split('.')[0] for frame in frames] # .split('_')[0] + frames = [Frame(frame.split('_')[0], frame.split('_')[1]) for frame in frames] + return frames + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument("-f", "--folder", help = "end of range", default='') + parser.add_argument("-f_obj", "--folder_object", help = "end of range", default='/') + parser.add_argument("-s", "--start", help = "start of range", default=0) + 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() + + obj_folder = args.folder_object # '/data-net/dataset/cloth3d/train' + folder = args.folder #'/data-local/127/train' + + os.makedirs(os.path.join(folder, 'mesh_normals'), exist_ok=True) + frames = get_frames(folder) + smpl = SMPL('cpu', 'smpl') + + ids = {frame.id for frame in frames} + ids = sorted(list(ids))[int(args.start):int(args.end)] + + for id in tqdm(ids): + # find obj files in the folder + garment_obj_folder = os.path.join(obj_folder, id[:5]) + garments = [file.split('.')[0] for file in os.listdir(garment_obj_folder) if file.endswith('.obj')] + + if len(id) == 5: + garment = garments[0] + elif len(id) == 6: + garment = garments[1] + else: + raise ValueError('Invalid frame id') + + mesh_file = os.path.join(garment_obj_folder, f'{garment}.obj') + vertices =readPC2(os.path.join(garment_obj_folder, f'{garment}.pc16'), True)['V'] + poses = sio.loadmat(os.path.join(garment_obj_folder, f'info.mat'))['poses'] + normals = calculate_normals(vertices, mesh_file, poses) + + map = sio.loadmat(os.path.join(garment_obj_folder, f'{garment}_map.mat'))['map'][0] + mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy')) + + mesh_normals = normals[:,map,:] + mesh_normals = mesh_normals * mask[None,:, None] + + for i in range(vertices.shape[0]): + # np.savez(os.path.join(folder, 'mesh_normals', f'{id}_{i}.npz'), mesh_normals=mesh_normals) + np.savez_compressed(os.path.join(folder, 'mesh_normals', f'{id}_{i}.npz'), mesh_normals=mesh_normals[i]) + + # pass + + # print(frame.id, frame.frame_nr) + diff --git a/run_parallel.py b/run_parallel.py index 834de78..da12a64 100644 --- a/run_parallel.py +++ b/run_parallel.py @@ -4,24 +4,27 @@ import os if __name__ == '__main__': - folder = '/home/workspace/data/dataset/cloth3d/train' + # folder = '/home/workspace/data/dataset/cloth3d/train' # folder = '/home/workspace/local-data/cape_processed/train/male' # folder = '/home/workspace/local-data/valid_skirt' # folder = '/home/workspace/local-data/valid' # ids = os.listdir(folder) # ids = os.listdir(folder + '/mesh_posed_no_orient')#[280000:] skirt = 0 + obj_folder = '/data-net/dataset/cloth3d/train' + folder = '/data-local/127/train' - count = int(500 * 1e3) - step = 10000 + count = 0 + step = 100 # !!!!!!! change skirt parameter in script - while count < 510 * 1e3: + while count <1200: # os.system('nohup python convert_to_uv_cape.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_2_cape_train_male_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_2_cape_train_male_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &') # os.system('nohup python convert_to_uv.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_unposed_onehot_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_train_unposed_onehot_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &') # os.system('nohup python convert_to_uv_local.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_local_unposed_onehot_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_train_local_unposed_onehot_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &') - os.system('nohup python data/convert_to_uv.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_2_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv/train_2_' + str(count) + '_' + str(count+step) + '_again.out' + ' 2>&1 &') + # os.system('nohup python data/convert_to_uv.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_2_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv/train_2_' + str(count) + '_' + str(count+step) + '_again.out' + ' 2>&1 &') + os.system(f'nohup python data/calculate_normals.py -f {folder} -f_obj {obj_folder} -s {count} -e {(count+step)} -n normals_train_calc_{count}_{(count+step)} > nohup/normals/normals_train_calc_{count}_{(count+step)}.out' + ' 2>&1 &') # os.system('nohup python convert_to_uv_local_single.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_skirt_local_unposed_onehot_3_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_train_skirt_local_unposed_onehot_3_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &') # os.system('nohup python convert_to_uv_local_single.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_valid_skirt_local_unposed_onehot_3_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_valid_skirt_local_unposed_onehot_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &') # os.system('nohup python downscale.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'downscale_valid_posed_no_orient_' + str(count) + '_' + str(count+step) + ' > nohup/downscale_valid_posed_no_orient_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')