update preprocessing
This commit is contained in:
+133
-99
@@ -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()
|
||||
|
||||
@@ -457,6 +370,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')):
|
||||
continue
|
||||
@@ -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,13 +633,14 @@ 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')
|
||||
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_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)
|
||||
@@ -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'
|
||||
|
||||
Reference in New Issue
Block a user