luni/cvpr_2023
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

update preprocessing

Browse Source
This commit is contained in:
luni
2023-07-03 22:56:50 +02:00
parent 8fb4ceb5db
commit 3e4a54d629
1 changed files with 147 additions and 113 deletions
Download Patch File Download Diff File Expand all files Collapse all files
data/preprocessing.py
+147 -113
View File
@@ -36,94 +36,6 @@ from smpl.smpl_torch import SMPL
from custom_utils.custom_utils import _maskFaces 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): def create_extras_file(folder):
# create a new pytables file # create a new pytables file
root = '/'.join(folder.split('/')[:-1]) root = '/'.join(folder.split('/')[:-1])
@@ -261,14 +173,15 @@ def create_mesh_body_posed_no_orient(folder, out_folder):
continue continue
info = sio.loadmat(os.path.join(folder, id, 'info.mat')) info = sio.loadmat(os.path.join(folder, id, 'info.mat'))
pose, shape, gender = info['poses'], info['shape'], info['gender'] pose, shape, gender = info['poses'], info['shape'], info['gender']
# if pose.shape[0] == 72:
pose = np.transpose(pose, [1, 0]).astype(np.float32) pose = np.transpose(pose, [1, 0]).astype(np.float32)
pose[:, :3] = 0 pose[:, :3] = 0
beta = np.squeeze(np.concatenate((shape, gender), axis=1)).astype(np.float32) 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 = 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() 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') pose = torch.tensor(pose, dtype=torch.float32, device='cuda')
shape = torch.tensor(shape, 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: for garment in garments:
if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')): if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')):
continue 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') 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): def create_frames_file_partial(folder, out_folder, split=0.9, part=0.1):
ids = os.listdir(folder) 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): def check_irregular_poses(folder, folder_reg, out_folder):
frame: Frame = Frame('01512', 50, 'Dress') for i in range(100, 200, 5):
mesh_posed = readPC2Frame(os.path.join(out_folder, 'mesh_posed_no_orient', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True) frame: Frame = Frame('01512', i, 'Dress')
# mesh_static = np.load(f'{out_folder}/mesh_static/{frame.id}_{frame.garment}.npy') 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_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 = 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') F = np.load('smpl/smpl_faces_skirt.npy')
_, faces = _maskFaces(np.zeros((14475, 3)), F, mesh_mask) _, 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 = 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.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( fig.update_layout(scene = dict(
xaxis = dict(range=[-1,1],), xaxis = dict(range=[-1,1],),
yaxis = dict(range=[-1,1],), yaxis = dict(range=[-1,1],),
zaxis = dict(range=[-1,1],),),) zaxis = dict(range=[-1,1],),),)
fig.update_layout(scene_aspectmode='cube') fig.update_layout(scene_aspectmode='cube')
fig.show() fig.show()
pass 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) 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__': if __name__ == '__main__':
# folder = '/data-net/dataset/cloth3d/test' folder = '/data-net/dataset/cloth3d/test'
# folder_reg = '/data-local/127/cloth3d_registration/test' folder_reg = '/data-local/127/cloth3d_registration/test'
# out_folder = '/data-local/127/cloth3d_processed_test' out_folder = '/data-local/127_2/cloth3d_processed_test'
folder = '/data-net/dataset/cloth3d/train' # folder = '/data-net/dataset/cloth3d/train'
folder_reg = '/data-local/127/cloth3d_registration/train' # folder_reg = '/data-local/127/cloth3d_registration/train'
out_folder = '/data-local/127/cloth3d_processed' # out_folder = '/data-local/127/cloth3d_processed'
smpl = SMPL(path='smpl', device='cuda') smpl = SMPL(path='smpl', device='cuda')
@@ -754,6 +785,8 @@ if __name__ == '__main__':
# create_extras_file(folder) # create_extras_file(folder)
# check_irregular_poses(folder, folder_reg, out_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) # create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder)
# verify_registration_ans_mask(folder, 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_body_posed_no_orient(folder, out_folder)
# create_mesh_normals(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) # 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(folder, out_folder)
create_frames_file_test(folder, out_folder)
# create_frames_file_partial(folder, out_folder, part=0.5) # create_frames_file_partial(folder, out_folder, part=0.5)
# folder = '/data-local/127_2/cape_processed/test/female' # folder = '/data-local/127_2/cape_processed/test/female'