diff --git a/data/preprocessing.py b/data/preprocessing.py index aac9549..9a4a606 100644 --- a/data/preprocessing.py +++ b/data/preprocessing.py @@ -30,7 +30,7 @@ import torch from scipy.spatial.transform import Rotation as R sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') -from data.cloth3d_io import readPC2, loadInfo, writePC2, readOBJ, quads2tris +from data.cloth3d_io import readPC2, loadInfo, writePC2, readOBJ, quads2tris, readPC2Frame from data.utils import Frame from smpl.smpl_torch import SMPL from custom_utils.custom_utils import _maskFaces @@ -160,7 +160,7 @@ def deorient(verts, pose): def create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder): - ids = os.listdir(folder) + ids = sorted(os.listdir(folder)) os.makedirs(os.path.join(out_folder, 'mesh_posed_no_orient'), exist_ok=True) os.makedirs(os.path.join(out_folder, 'mesh_mask'), exist_ok=True) @@ -179,12 +179,18 @@ def create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder): continue verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V'] + # if info['poses'].shape[1] == verts.shape[0]: + # continue reg_data = sio.loadmat(os.path.join(folder_reg, id, garment + '_weights_fixed.mat')) verts = (verts[:, reg_data['F'], :] * reg_data['W'][None, :, :, None]).sum(axis=2) mask_idx = sio.loadmat(os.path.join(folder_reg, id, garment + '_mask.mat'))['mask_idx'].reshape((-1)) - 1 mask = np.zeros((14475), dtype=bool) mask[mask_idx] = True + if info['poses'].shape[1] != verts.shape[0]: + assert False, 'Somethign wrong here' + info['poses'] = info['poses'][:, :verts.shape[0]] + verts = deorient(verts, info['poses']) # # plot as scatter plot with plotly @@ -193,13 +199,12 @@ def create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder): # fig = go.Figure(data=[go.Scatter3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], mode='markers', marker=dict(size=1))]) # fig.show() - writePC2(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), verts, True) np.save(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy'), mask) def create_pose_shape(folder, out_folder): - ids = os.listdir(folder) + ids = sorted(os.listdir(folder)) os.makedirs(os.path.join(out_folder, 'pose'), exist_ok=True) os.makedirs(os.path.join(out_folder, 'shape'), exist_ok=True) @@ -216,7 +221,7 @@ def create_pose_shape(folder, out_folder): def create_mesh_static(folder, out_folder): - ids = os.listdir(folder) + ids = sorted(os.listdir(folder)) os.makedirs(os.path.join(out_folder, 'mesh_static'), exist_ok=True) for id in tqdm(ids): @@ -247,7 +252,7 @@ def create_mesh_static(folder, out_folder): def create_mesh_body_posed_no_orient(folder, out_folder): - ids = os.listdir(folder) + ids = sorted(os.listdir(folder)) os.makedirs(os.path.join(out_folder, 'mesh_body_posed_no_orient'), exist_ok=True) for id in tqdm(ids): @@ -288,10 +293,11 @@ def calculate_normals(vertices, faces): def create_mesh_normals(folder, out_folder): - ids = os.listdir(folder) + ids = sorted(os.listdir(folder)) os.makedirs(os.path.join(out_folder, 'mesh_normals'), exist_ok=True) # F = np.load('smpl/smpl_faces.npy') # F_skirt = np.load('smpl/smpl_faces_skirt.npy') + # ids = ['00438'] for id in tqdm(ids): garment_dir = os.path.join(folder, id) @@ -306,7 +312,7 @@ def create_mesh_normals(folder, out_folder): continue verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V'] - verts = deorient(verts, info['poses']) + verts = deorient(verts, info['poses'][:, :verts.shape[0]]) _, faces, _, _ = readOBJ(os.path.join(garment_dir, garment + '.obj')) faces = quads2tris(faces) normals = calculate_normals(verts, faces.astype(np.int32)) @@ -413,7 +419,7 @@ def create_mesh_normals_cape(folder, out_folder): def verify_registration_ans_mask(folder, out_folder): - ids = os.listdir(folder) + ids = sorted(os.listdir(folder)) for id in ids: garment_dir = os.path.join(folder, id) @@ -621,9 +627,80 @@ def create_frames_file(folder, out_folder, split=0.9): f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n') +def create_frames_file_partial(folder, out_folder, split=0.9, part=0.1): + ids = os.listdir(folder) + + frames_non_skirt = [] + frames_skirt = [] + ids_non_skirt = set() + ids_skirt = set() + + for id in tqdm(ids): + garment_dir = os.path.join(folder, id) + if not os.path.isdir(garment_dir): + continue + + info = sio.loadmat(os.path.join(folder, id, 'info.mat')) + garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')] + frames = info['poses'].shape[1] + + assert info['poses'].shape[0] == 72 + + for garment in garments: + for i in range(frames): + if garment == 'Skirt' or garment == 'Dress': + frames_skirt.append(Frame(id, i, garment)) + ids_skirt.add(id) + else: + frames_non_skirt.append(Frame(id, i, garment)) + ids_non_skirt.add(id) + + nr_train_non_skirt = int(len(ids_non_skirt) * split * part) + nr_train_skirt = int(len(ids_skirt) * split * part) + + ids_skirt = list(ids_skirt) + ids_non_skirt = list(ids_non_skirt) + + with open(os.path.join(out_folder, f'frames_train_{int(part*100)}p.txt'), 'w') as f: + for frame in frames_non_skirt: + if frame.id in ids_non_skirt[:nr_train_non_skirt]: + f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n') + + with open(os.path.join(out_folder, f'frames_train_skirt_{int(part*100)}p.txt'), 'w') as f: + for frame in frames_skirt: + if frame.id in ids_skirt[:nr_train_skirt]: + f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n') + + + +def check_irregular_poses(folder, folder_reg, out_folder): + frame: Frame = Frame('01512', 50, 'Dress') + mesh_posed = readPC2Frame(os.path.join(out_folder, 'mesh_posed_no_orient', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True) + # mesh_static = np.load(f'{out_folder}/mesh_static/{frame.id}_{frame.garment}.npy') + + mesh_mask = np.load(f'{out_folder}/mesh_mask/{frame.id}_{frame.garment}.npy') + mesh_body_posed = readPC2Frame(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{frame.id}.pc16'), frame.frame_nr, True) + + F = np.load('smpl/smpl_faces_skirt.npy') + _, faces = _maskFaces(np.zeros((14475, 3)), F, mesh_mask) + + fig = go.Figure(data=[go.Mesh3d(x=mesh_posed[:,0], y=mesh_posed[:,2], z=mesh_posed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1)]) + fig.add_trace(go.Mesh3d(x=mesh_body_posed[:,0], y=mesh_body_posed[:,2], z=mesh_body_posed[:,1], i=F[:,0], j=F[:,1], k=F[:,2], color='lightblue', opacity=1)) + fig.update_layout(scene = dict( + xaxis = dict(range=[-1,1],), + yaxis = dict(range=[-1,1],), + zaxis = dict(range=[-1,1],),),) + fig.update_layout(scene_aspectmode='cube') + fig.show() + pass + if __name__ == '__main__': + # folder = '/data-net/dataset/cloth3d/test' + # folder_reg = '/data-local/127/cloth3d_registration/test' + # out_folder = '/data-local/127/cloth3d_processed_test' + folder = '/data-net/dataset/cloth3d/train' folder_reg = '/data-local/127/cloth3d_registration/train' out_folder = '/data-local/127/cloth3d_processed' @@ -635,6 +712,8 @@ if __name__ == '__main__': # split_data(folder, 0.9) # create_extras_file(folder) + # check_irregular_poses(folder, folder_reg, out_folder) + # create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder) # verify_registration_ans_mask(folder, out_folder) # create_pose_shape(folder, out_folder) @@ -644,6 +723,7 @@ if __name__ == '__main__': # create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=args.start, end=args.end) # create_frames_file(folder, out_folder) + create_frames_file_partial(folder, out_folder, part=0.5) - folder = '/data-local/127_2/cape_processed/test/female' - create_mesh_normals_cape(folder, None) \ No newline at end of file + # folder = '/data-local/127_2/cape_processed/test/female' + # create_mesh_normals_cape(folder, None) \ No newline at end of file