From 3e4a54d629091202ad67bbc68364467e15f16174 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Mon, 3 Jul 2023 22:56:50 +0200 Subject: [PATCH] update preprocessing --- data/preprocessing.py | 260 ++++++++++++++++++++++++------------------ 1 file changed, 147 insertions(+), 113 deletions(-) diff --git a/data/preprocessing.py b/data/preprocessing.py index dff4038..977c1f2 100644 --- a/data/preprocessing.py +++ b/data/preprocessing.py @@ -36,94 +36,6 @@ from smpl.smpl_torch import SMPL from custom_utils.custom_utils import _maskFaces - -def clean_dataset(folder): - ids = os.listdir(folder) - - for id in tqdm(ids): - - if not os.path.isdir(os.path.join(folder, id)): - continue - - files = os.listdir(os.path.join(folder, id)) - - for file in files: - if file.startswith('glmarks') or file.find('nricp') != -1: - os.remove(os.path.join(folder, id, file)) - - -def generate_frames_file(folder, file_name): - ids = os.listdir(folder) - frames = [] - - for id in tqdm(ids): - - if not os.path.isdir(os.path.join(folder, id)): - continue - - garments = [f.replace('.obj','') for f in os.listdir(os.path.join(folder, id)) if f.endswith('.obj')] - info = loadInfo(os.path.join(folder, id, 'info.mat')) - - for garment in garments: - for i in range(info['poses'].shape[1]): - frames.append(Frame(id, i, garment)) - - with open(os.path.join(folder, file_name), 'w') as f: - for i, frame in enumerate(frames): - f.write(str(frame.id) + ' ' + str(frame.frame_nr) + ' ' + str(frame.garment) + '\n') - - -def split_data(folder, ratio): - with open(os.path.join(folder, 'frames_all.txt'), "r") as f: - frames = f.readlines() - - ids = sorted(list(set([frame.split()[0] for frame in frames]))) - train_last_id = round(len(ids) * ratio) - - train_frames = [] - valid_frames = [] - - for frame in tqdm(frames): - id = frame.split()[0] - if id in ids[:train_last_id]: - train_frames.append(frame) - else: - valid_frames.append(frame) - - with open(os.path.join(folder, f'frames_train_{ratio}.txt'), 'w') as f: - f.writelines(train_frames) - - with open(os.path.join(folder, f'frames_valid_{round(1-ratio,1)}.txt'), 'w') as f: - f.writelines(valid_frames) - - -def genereated_registered_and_mask(folder, is_skirt = False): - ids = os.listdir(folder) - - for id in tqdm(ids): - garment_dir = os.path.join(folder, id) - if not os.path.isdir(garment_dir): - continue - garments = [f.replace('.pc16','') for f in os.listdir(garment_dir) if f.find('_') == -1 and not f.endswith('.mat') and not f.endswith('.obj') and not f.startswith('Combined')] - - assert len(garments) == 1 - - if os.path.exists(os.path.join(garment_dir, garments[0] + '_registered.pc16')): - continue - - garment = garments[0] - - verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V'] - map = sio.loadmat(os.path.join(garment_dir, garment + '_map.mat'))['map'].reshape((-1)) - verts = verts[:, map, :] # ---------------> This is the garment after registration - raw_mask = sio.loadmat(os.path.join(garment_dir, garment + '_mask.mat'))['mask'].reshape((-1)) - mask = raw_mask > 0 - # boundaries = raw_mask == 2 - - writePC2(os.path.join(garment_dir, garment + '_registered.pc16'), verts, True) - np.save(os.path.join(garment_dir, garment + '_mask.npy'), mask) - - def create_extras_file(folder): # create a new pytables file root = '/'.join(folder.split('/')[:-1]) @@ -261,14 +173,15 @@ def create_mesh_body_posed_no_orient(folder, out_folder): continue info = sio.loadmat(os.path.join(folder, id, 'info.mat')) - - pose, shape, gender = info['poses'], info['shape'], info['gender'] + # if pose.shape[0] == 72: pose = np.transpose(pose, [1, 0]).astype(np.float32) pose[:, :3] = 0 beta = np.squeeze(np.concatenate((shape, gender), axis=1)).astype(np.float32) + smpl = SMPL(path='smpl', device='cuda', batch_size=pose.shape[0]) + T = smpl(torch.tensor(pose, device=smpl.device), torch.tensor(np.tile(beta[np.newaxis, :], [pose.shape[0], 1]), device=smpl.device))[0] T = T.cpu().numpy() @@ -456,6 +369,9 @@ def create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=0, end=1) pose = torch.tensor(pose, dtype=torch.float32, device='cuda') shape = torch.tensor(shape, dtype=torch.float32, device='cuda') + + smpl = SMPL(path='smpl', device='cuda', batch_size=pose.shape[0]) + for garment in garments: if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')): @@ -627,6 +543,49 @@ 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_test(folder, out_folder): + 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) + + ids_skirt = list(ids_skirt) + ids_non_skirt = list(ids_non_skirt) + + with open(os.path.join(out_folder, 'frames_test.txt'), 'w') as f: + for frame in frames_non_skirt: + if frame.id in ids_non_skirt: + f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n') + + with open(os.path.join(out_folder, 'frames_test_skirt.txt'), 'w') as f: + for frame in frames_skirt: + if frame.id in ids_skirt: + 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) @@ -674,25 +633,26 @@ def create_frames_file_partial(folder, out_folder, split=0.9, part=0.1): 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') + for i in range(100, 200, 5): + frame: Frame = Frame('01512', i, '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 = readPC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{frame.id}.pc16'), True)['V'] - mesh_body_posed = mesh_body_posed[int(mesh_body_posed.shape[0]/2)+frame.frame_nr] + mesh_mask = np.load(f'{out_folder}/mesh_mask/{frame.id}_{frame.garment}.npy') + mesh_body_posed = readPC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{frame.id}.pc16'), True)['V'] + mesh_body_posed = mesh_body_posed[::2][frame.frame_nr] - F = np.load('smpl/smpl_faces_skirt.npy') - _, faces = _maskFaces(np.zeros((14475, 3)), F, mesh_mask) + 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() + 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 @@ -734,17 +694,88 @@ def calculate_average(folder, out_folder, frames_file='frames_train_skirt.txt'): np.save('custom_data/avg/train_skirt_no_shape_avg_uv.npy', avg_garment) +def remove_irregular_poses(folder, folder_reg, out_folder): + ids = sorted(os.listdir(folder)) + + for id in tqdm(ids): + + garment_dir = os.path.join(folder, id) + if not os.path.isdir(garment_dir): + continue + + garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')] + info = loadInfo(os.path.join(folder, id, 'info.mat')) + + for garment in garments: + + verts_count = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V'].shape[0] + if info['poses'].shape[1] != verts_count: + shutil.rmtree(os.path.join(folder, id)) + if os.path.exists(os.path.join(out_folder, 'mesh_static', f'{id}_{garment}.pc16')): + os.remove(os.path.join(out_folder, 'mesh_static', f'{id}_{garment}.pc16')) + if os.path.exists(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16')): + os.remove(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16')) + if os.path.exists(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16')): + os.remove(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16')) + if os.path.exists(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy')): + os.remove(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy')) + # os.remove(os.path.join(out_folder, 'mesh_unposed_no_shape', f'{id}_{garment}.npy')) + + if len(garments) == 2: + if os.path.exists(os.path.join(out_folder, 'mesh_static', f'{id}_{garments[1]}.pc16')): + os.remove(os.path.join(out_folder, 'mesh_static', f'{id}_{garments[1]}.pc16')) + if os.path.exists(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garments[1]}.pc16')): + os.remove(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garments[1]}.pc16')) + if os.path.exists(os.path.join(out_folder, 'mesh_mask', f'{id}_{garments[1]}.pc16')): + os.remove(os.path.join(out_folder, 'mesh_mask', f'{id}_{garments[1]}.npy')) + break + # os.remove(os.path.join(out_folder, 'mesh_unposed_no_shape', f'{id}_{garments[1]}.npy')) + + +def double_check_irregular_poses(folder, folder_reg, out_folder): + # ids = sorted(os.listdir(folder)) + + # for id in tqdm(ids): + + # garment_dir = os.path.join(folder, id) + # if not os.path.isdir(garment_dir): + # continue + + # garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')] + # info = loadInfo(os.path.join(folder, id, 'info.mat')) + + # for garment in garments: + # body_count = readPC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16'), True)['V'].shape[0] + # verts_count = readPC2(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), True)['V'].shape[0] + + # assert body_count == verts_count + + frames = sorted(os.listdir(out_folder + '/mesh_posed_no_orient')) + frames = [(id.split('_')[0], id.split('_')[1].split('.')[0]) for id in frames] + + for frame in tqdm(frames): + id = frame[0] + garment_dir = os.path.join(folder, id) + if not os.path.isdir(garment_dir): + continue + + garment = frame[1] + body_count = readPC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16'), True)['V'].shape[0] + verts_count = readPC2(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), True)['V'].shape[0] + + assert body_count == verts_count + 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/test' + folder_reg = '/data-local/127/cloth3d_registration/test' + out_folder = '/data-local/127_2/cloth3d_processed_test' - folder = '/data-net/dataset/cloth3d/train' - folder_reg = '/data-local/127/cloth3d_registration/train' - out_folder = '/data-local/127/cloth3d_processed' + # folder = '/data-net/dataset/cloth3d/train' + # folder_reg = '/data-local/127/cloth3d_registration/train' + # out_folder = '/data-local/127/cloth3d_processed' smpl = SMPL(path='smpl', device='cuda') @@ -754,6 +785,8 @@ if __name__ == '__main__': # create_extras_file(folder) # check_irregular_poses(folder, folder_reg, out_folder) + # remove_irregular_poses(folder, folder_reg, out_folder) + # double_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) @@ -762,9 +795,10 @@ if __name__ == '__main__': # create_mesh_body_posed_no_orient(folder, out_folder) # create_mesh_normals(folder, out_folder) # create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=args.start, end=args.end) - calculate_average(folder, out_folder) + # calculate_average(folder, out_folder) # create_frames_file(folder, out_folder) + create_frames_file_test(folder, out_folder) # create_frames_file_partial(folder, out_folder, part=0.5) # folder = '/data-local/127_2/cape_processed/test/female'