from glob import glob import os import tqdm import shutil import numpy as np import plotly.express as ex import plotly.graph_objects as go from custom_utils.custom_utils import _maskFaces os.environ["CUDA_VISIBLE_DEVICES"] = '1' import tensorflow as tf gpus = tf.config.experimental.list_physical_devices('GPU') if gpus: try: # Currently, memory growth needs to be the same across GPUs for gpu in gpus: tf.config.experimental.set_memory_growth(gpu, True) logical_gpus = tf.config.experimental.list_logical_devices('GPU') print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs") except RuntimeError as e: # Memory growth must be set before GPUs have been initialized print(e) from custom_utils.smpl import SMPL # from custom_utils.smpl import SMPL folder = '/home/workspace/data/dataset/cape_release' # smpl = SMPL(orig_smpl=True) def get_frames(): smpl = SMPL(orig_smpl=True) # sequences = np.load() gender_table = np.load('/home/workspace/data/dataset/cape_release/misc/subj_genders.pkl', allow_pickle=True) test = np.load(os.path.join(folder, 'train/male/02474/longshort_ATUsquat/longshort_ATUsquat.000100.npz')) faces = np.load(os.path.join(folder, 'misc/smpl_tris.npy')) min_body = np.load(os.path.join(folder, 'minimal_body_shape/02474/02474_minimal.npy')) shape = np.load(os.path.join(folder, 'minimal_body_shape/02474/betas_2.npy')) offsets_cano = test['v_cano'] - min_body pose = test['pose'].astype(np.float32) shape = np.concatenate((shape, [0] if gender_table['02474'] == 'female' else [1])).astype(np.float32) posed_body = smpl.run(smpl.run_rodrigues(pose[np.newaxis,:]), shape[np.newaxis,:], remove=False)[0].numpy()[0] transl = test['transl'] posed_body = posed_body + transl[np.newaxis, :] offset_posed = test['v_posed'] - posed_body mask = np.abs(offsets_cano.sum(axis=-1)) > 1e-2 zero_body = smpl.run(smpl.run_rodrigues(np.zeros((1, 72), np.float32)), shape[np.newaxis,:], remove=True)[0] zero_body_root, J, _ = smpl.run(smpl.run_rodrigues(np.zeros((1, 72), np.float32)), shape[np.newaxis,:], remove=True) # _, J0, _ = smpl.run(smpl.run_rodrigues(pose[np.newaxis,:]), shape[np.newaxis,:], remove=True, subtract_root=True) posed_body = smpl.run_garment(pose[np.newaxis,:], shape[np.newaxis,:], zero_body)[0].numpy() posed_body_root = smpl.run_garment(pose[np.newaxis,:], shape[np.newaxis,:], zero_body_root) posed_body_root = posed_body_root[0].numpy() # verts, faces = _maskFaces(test['v_cano'], faces, mask) # x = test['v_cano'] # y = min_body x = min_body[smpl.to_keep, :] - J[0,0,:][np.newaxis, :] y = zero_body_root[0].numpy() faces = np.load('custom_utils/smpl_faces_orig.npy') fig = go.Figure(data=[ go.Mesh3d( # 8 vertices of a cube x=x[:,0], y=x[:,1], z=x[:,2], # i, j and k give the vertices of triangles i = faces[:,0], j = faces[:,1], k = faces[:,2] ), go.Mesh3d( # 8 vertices of a cube x=y[:,0], y=y[:,1], z=y[:,2], # i, j and k give the vertices of triangles i = faces[:,0], j = faces[:,1], k = faces[:,2] ), ]) fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1)) fig.update_layout( scene = dict( xaxis = dict(nticks=4, range=[-1.5,1.5],), yaxis = dict(nticks=4, range=[-1.5,1.5],), zaxis = dict(nticks=4, range=[-1.5,1.5],),) ) fig.show() print() folder_processed = '/home/workspace/data/dataset/cape_processed/datasets/dataset_male_4clotypes' def get_frames_processed(): # sequences = np.load()dataset/cape_processed/datasets/dataset_male_4clotypes disp = np.load(os.path.join(folder_processed, 'train/train_disp.npy')) faces = np.load(os.path.join(folder, 'misc/smpl_tris.npy')) min_body = np.load(os.path.join(folder, 'minimal_body_shape/03383/03383_minimal.npy')) x = verts fig = go.Figure(data=[ go.Mesh3d( # 8 vertices of a cube x=x[:,0], y=x[:,1], z=x[:,2], # i, j and k give the vertices of triangles i = faces[:,0], j = faces[:,1], k = faces[:,2] ), ]) fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1)) fig.update_layout( scene = dict( xaxis = dict(nticks=4, range=[-1,1],), yaxis = dict(nticks=4, range=[-1,1],), zaxis = dict(nticks=4, range=[-1,1],),) ) fig.show() print() folder_pop = '/home/workspace/data/dataset/cape_2/ps/scratch/ps_shared/qma/abps_data/CAPE_packed_for_POP_release/00096_shortlong/train' def get_frames_pop(): # sequences = np.load()dataset/cape_2/ps/scratch/ps_shared/qma/abps_data/CAPE_packed_for_POP_release/00096_shortlong/train/shortlong_ATUsquat.000009.npz test = np.load(os.path.join(folder_pop, 'shortlong_ATUsquat.000009.npz')) faces = np.load(os.path.join(folder, 'misc/smpl_tris.npy')) min_body = np.load(os.path.join(folder, 'minimal_body_shape/03383/03383_minimal.npy')) offsets_cano = test['v_cano'] - min_body mask = np.abs(offsets_cano.sum(axis=-1)) < 1e-2 # verts, faces = _maskFaces(test['v_posed'], faces, mask) x = test['v_cano'] y = min_body fig = go.Figure(data=[ go.Mesh3d( # 8 vertices of a cube x=x[:,0], y=x[:,1], z=x[:,2], # i, j and k give the vertices of triangles i = faces[:,0], j = faces[:,1], k = faces[:,2] ), go.Mesh3d( # 8 vertices of a cube x=y[:,0], y=y[:,1], z=y[:,2], # i, j and k give the vertices of triangles i = faces[:,0], j = faces[:,1], k = faces[:,2] ), ]) fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1)) fig.update_layout( scene = dict( xaxis = dict(nticks=4, range=[-1,1],), yaxis = dict(nticks=4, range=[-1,1],), zaxis = dict(nticks=4, range=[-1,1],),) ) fig.show() print() def split_data(): from CAPE.data.dataset_configs import dataset_config_dicts count = 0 cut_first = 2 sample_rate = 1 src = '/home/workspace/local-data/temp/cape_release/sequences' dst = '/home/workspace/local-data/cape/' + phase + '/' + gender os.makedirs(dst, exist_ok=True) config = dataset_config_dicts[gender] ids = os.listdir(src) for id in ids: if id not in config[phase + '_' + 'subjs']: continue garment_dir = os.path.join(src, id) sequences = os.listdir(garment_dir) for sequence in sequences: sequence_name = '_'.join(sequence.split('_')[1:]) if sequence_name in config['exclude_seqs']: continue if sequence_name not in config[phase + '_' + 'seqs']: continue cloth = sequence.split('_')[0] if cloth not in config[phase + '_' + 'cloth']: continue if [id, cloth] in config['exclude_cases']: continue frames = sorted(glob(os.path.join(garment_dir, sequence, '*.npz'))) if len(frames) == 0: print('{} no files here, skipping ..\n'.format(id+'_'+sequence)) return frames = frames[cut_first: -cut_first : sample_rate] # take every k frame if len(frames) == 0: continue os.makedirs(os.path.join(dst, id, sequence), exist_ok=True) for frame in frames: try: np.load(frame) except: # skip and report the broken file if any print('{} broken, skipping..'.format(frame)) continue frame_name = frame.split('/')[-1] count += 1 print(id + '_' + frame_name) test_name = frame.split('/') test_name.index(id) test_name.index(sequence) test_name.index(frame_name) shutil.move(frame, os.path.join(dst, id, sequence, frame_name)) print(count) def verify_split_data(): from CAPE.data.dataset_configs import dataset_config_dicts count = 0 cut_first = 2 sample_rate = 1 phase = 'test' gender = 'male' src = '/home/workspace/data/dataset/cape_release/sequences' dst = '/home/workspace/data/dataset/cape_release/' + phase + '/' + gender config = dataset_config_dicts[gender] ids = os.listdir(src) for id in ids: if id not in config[phase + '_' + 'subjs']: continue garment_dir = os.path.join(src, id) sequences = os.listdir(garment_dir) for sequence in sequences: sequence_name = '_'.join(sequence.split('_')[1:]) if sequence_name in config['exclude_seqs']: continue if sequence_name not in config[phase + '_' + 'seqs']: continue cloth = sequence.split('_')[0] if cloth not in config[phase + '_' + 'cloth']: continue if [id, cloth] in config['exclude_cases']: continue frames = sorted(glob(os.path.join(garment_dir, sequence, '*.npz'))) if len(frames) == 0: print('{} no files here, skipping ..\n'.format(id+'_'+sequence)) return frames = frames[cut_first: -cut_first : sample_rate] # take every k frame if len(frames) == 0: continue # os.makedirs(os.path.join(dst, id, sequence), exist_ok=True) for frame in frames: try: np.load(frame) except: # skip and report the broken file if any print('{} broken, skipping..'.format(frame)) continue frame_name = frame.split('/')[-1] count += 1 print(id + '_' + frame_name) # shutil.move(frame, os.path.join(dst, id, sequence, frame_name)) print(count) def merge_pose_shape(): src = os.path.join('/home/workspace/local-data/cape', phase, gender) dst = os.path.join('/home/workspace/local-data/cape_processed', phase, gender, 'pose_shape') gender_table = np.load('/home/workspace/local-data/temp/cape_release/misc/subj_genders.pkl', allow_pickle=True) os.makedirs(dst, exist_ok=True) ids = os.listdir(src) for id in ids: garment_dir = os.path.join(src, id) sequences = os.listdir(garment_dir) for sequence in sequences: # sequence_name = '_'.join(sequence.split('_')[1:]) # cloth = sequence.split('_')[0] frames = sorted(os.listdir(os.path.join(garment_dir, sequence))) for frame in frames: frame_name = frame.split('.') assert frame_name[-1] == 'npz', 'It\'s not an npz file!' frame_name = '.'.join(frame_name[:-1]) output_id = id + '_' + frame_name # continue data = np.load(os.path.join(garment_dir, sequence, frame)) pose = data['pose'] shape = np.load(os.path.join('/home/workspace/data/dataset/cape_release_old/minimal_body_shape', id, 'betas_2.npy')) seq_gender = [0] if gender_table[id] == 'female' else [1] shape = np.concatenate((shape, seq_gender)).astype(np.float32) # body = smpl.run(smpl.run_rodrigues(pose[np.newaxis, :]), shape[np.newaxis,:], remove=False)[0][0] pose_shape = np.concatenate((pose, shape)) np.save(os.path.join(dst, id + '_' + frame_name + '.npy'), pose_shape) print(id + '_' + frame) def create_posed_body(): smpl = SMPL(orig_smpl=True) src = os.path.join('/home/workspace/local-data/cape', phase, gender) dst = os.path.join('/home/workspace/local-data/cape_processed', phase, gender, 'mesh_posed_body_no_orient') gender_table = np.load('/home/workspace/data/dataset/cape_release_old/misc/subj_genders.pkl', allow_pickle=True) os.makedirs(dst, exist_ok=True) ids = os.listdir(src) for id in ids: garment_dir = os.path.join(src, id) sequences = os.listdir(garment_dir) for sequence in sequences: # sequence_name = '_'.join(sequence.split('_')[1:]) # cloth = sequence.split('_')[0] frames = sorted(os.listdir(os.path.join(garment_dir, sequence))) for frame in frames: data = np.load(os.path.join(garment_dir, sequence, frame)) pose = data['pose'].astype(np.float32) shape = np.load(os.path.join('/home/workspace/data/dataset/cape_release_old/minimal_body_shape', id, 'betas_2.npy')) seq_gender = [0] if gender_table[id] == 'female' else [1] shape = np.concatenate((shape, seq_gender)).astype(np.float32) # body = smpl.run(smpl.run_rodrigues(pose[np.newaxis, :]), shape[np.newaxis,:], remove=True)[0][0] frame_name = frame.split('.') assert frame_name[-1] == 'npz', 'It\'s not an npz file!' frame_name = '.'.join(frame_name[:-1]) # np.savez_compressed(os.path.join(dst, id + '_' + frame_name + '.npz'), mesh_body_posed=body) pose_2 = pose.copy() pose_2[:3] = 0 body_2 = smpl.run(smpl.run_rodrigues(pose_2[np.newaxis, :]), shape[np.newaxis,:], remove=True)[0][0] np.savez_compressed(os.path.join(dst, id + '_' + frame_name + '.npz'), mesh_posed_body_no_orient=body_2) print(id + '_' + frame) def create_static(): smpl = SMPL(orig_smpl=True) src = os.path.join('/home/workspace/local-data/cape', phase, gender) dst = os.path.join('/home/workspace/local-data/cape_processed', phase, gender, 'mesh_static') src_orig = os.path.join('/home/workspace/local-data/temp/cape_release/sequences') J0 = np.load('custom_data/smpl_root_joint.npy') os.makedirs(dst, exist_ok=True) ids = os.listdir(src) for id in ids: garment_dir = os.path.join(src, id) sequences = os.listdir(garment_dir) for sequence in sequences: # sequence_name = '_'.join(sequence.split('_')[1:]) # cloth = sequence.split('_')[0] avaialble_frames = sorted(os.listdir(os.path.join(src_orig, id, sequence))) first_frame = avaialble_frames[0] print(first_frame) data = np.load(os.path.join(src_orig, id, sequence, first_frame)) frame = data['v_cano'][smpl.to_keep, :] - J0[np.newaxis, :] np.savez_compressed(os.path.join(dst, id + '_' + sequence + '.npz'), mesh_static = frame) # x = data['v_posed'][smpl.to_keep, :] - J0[np.newaxis, :] - data['transl'][np.newaxis, :] # faces = np.load('custom_utils/smpl_faces_orig.npy') # fig = go.Figure(data=[ # go.Mesh3d( # # 8 vertices of a cube # x=x[:,0], # y=x[:,1], # z=x[:,2], # # i, j and k give the vertices of triangles # i = faces[:,0], # j = faces[:,1], # k = faces[:,2] # ), # ]) # fig.update_layout(scene_aspectmode='manual', # scene_aspectratio=dict(x=1, y=1, z=1)) # fig.update_layout( # scene = dict( # xaxis = dict(nticks=4, range=[-1.5,1.5],), # yaxis = dict(nticks=4, range=[-1.5,1.5],), # zaxis = dict(nticks=4, range=[-1.5,1.5],),) # ) # fig.show() # print() def create_unposed(): smpl = SMPL(orig_smpl=True) src = os.path.join('/home/workspace/local-data/cape', phase, gender) dst = os.path.join('/home/workspace/local-data/cape_processed', phase, gender, 'mesh_unposed_no_shape') gender_table = np.load('/home/workspace/data/dataset/cape_release_old/misc/subj_genders.pkl', allow_pickle=True) os.makedirs(dst, exist_ok=True) # faces = np.load('custom_utils/smpl_faces_orig.npy') J0 = np.load('custom_data/smpl_root_joint.npy') ids = os.listdir(src) for id in ids: garment_dir = os.path.join(src, id) sequences = os.listdir(garment_dir) for sequence in sequences: # sequence_name = '_'.join(sequence.split('_')[1:]) # cloth = sequence.split('_')[0] frames = sorted(os.listdir(os.path.join(garment_dir, sequence))) for frame in frames: data = np.load(os.path.join(garment_dir, sequence, frame)) frame_name = frame.split('.') assert frame_name[-1] == 'npz', 'It\'s not an npz file!' frame_name = '.'.join(frame_name[:-1]) shape = np.load(os.path.join('/home/workspace/data/dataset/cape_release_old', 'minimal_body_shape/' + id + '/betas_2.npy')) seq_gender = 0 if gender_table[id] == 'female' else 1 shape_offsets = tf.tensordot(shape[tf.newaxis, :], smpl.shapedirs[seq_gender], axes=[[1], [2]]).numpy() unposed = data['v_cano'][smpl.to_keep, :] - J0[np.newaxis, :] - np.squeeze(shape_offsets)[smpl.to_keep, :] np.savez_compressed(os.path.join(dst, id + '_' + frame_name + '.npz'), mesh_unposed_no_shape=unposed) # x = unposed # faces = np.load('custom_utils/smpl_faces_orig.npy') # fig = go.Figure(data=[ # go.Mesh3d( # # 8 vertices of a cube # x=x[:,0], # y=x[:,1], # z=x[:,2], # # i, j and k give the vertices of triangles # i = faces[:,0], # j = faces[:,1], # k = faces[:,2] # ), # ]) # fig.update_layout(scene_aspectmode='manual', # scene_aspectratio=dict(x=1, y=1, z=1)) # fig.update_layout( # scene = dict( # xaxis = dict(nticks=4, range=[-1.5,1.5],), # yaxis = dict(nticks=4, range=[-1.5,1.5],), # zaxis = dict(nticks=4, range=[-1.5,1.5],),) # ) # fig.show() print(id + '_' + frame) def create_posed(): smpl = SMPL(orig_smpl=True) src = os.path.join('/home/workspace/local-data/cape', phase, gender) dst = os.path.join('/home/workspace/local-data/cape_processed', phase, gender, 'mesh_posed') os.makedirs(dst, exist_ok=True) J0 = np.load('custom_data/smpl_root_joint.npy') ids = os.listdir(src) for id in ids: garment_dir = os.path.join(src, id) sequences = os.listdir(garment_dir) for sequence in sequences: # sequence_name = '_'.join(sequence.split('_')[1:]) # cloth = sequence.split('_')[0] frames = sorted(os.listdir(os.path.join(garment_dir, sequence))) for frame in frames: data = np.load(os.path.join(garment_dir, sequence, frame)) frame_name = frame.split('.') assert frame_name[-1] == 'npz', 'It\'s not an npz file!' frame_name = '.'.join(frame_name[:-1]) posed = data['v_posed'][smpl.to_keep, :] - J0[np.newaxis, :] - data['transl'][np.newaxis, :] np.savez_compressed(os.path.join(dst, id + '_' + frame_name + '.npz'), mesh_posed=posed) # x = posed # faces = np.load('custom_utils/smpl_faces_orig.npy') # fig = go.Figure(data=[ # go.Mesh3d( # # 8 vertices of a cube # x=x[:,0], # y=x[:,1], # z=x[:,2], # # i, j and k give the vertices of triangles # i = faces[:,0], # j = faces[:,1], # k = faces[:,2] # ), # ]) # fig.update_layout(scene_aspectmode='manual', # scene_aspectratio=dict(x=1, y=1, z=1)) # fig.update_layout( # scene = dict( # xaxis = dict(nticks=4, range=[-1.5,1.5],), # yaxis = dict(nticks=4, range=[-1.5,1.5],), # zaxis = dict(nticks=4, range=[-1.5,1.5],),) # ) # fig.show() print(id + '_' + frame) def create_mask(): pass def verify(): root = os.path.join('/home/workspace/local-data/train') # root = os.path.join('/home/workspace/local-data/cape_processed', phase, gender) mesh_static = os.listdir(os.path.join(root, 'mesh_static')) mesh_posed = sorted(os.listdir(os.path.join(root, 'mesh_posed'))) mesh_posed_body_no_orient = sorted(os.listdir(os.path.join(root, 'mesh_posed_body_no_orient'))) mesh_unposed_no_shape = sorted(os.listdir(os.path.join(root, 'mesh_unposed_no_shape'))) pose_shape = sorted(os.listdir(os.path.join(root, 'pose_shape'))) mesh_curvature = sorted(os.listdir(os.path.join(root, 'mesh_curvature'))) mesh_posed = ['.'.join(x.split('.')[:-1]) for x in mesh_posed] mesh_posed_body_no_orient = ['.'.join(x.split('.')[:-1]) for x in mesh_posed_body_no_orient] mesh_unposed_no_shape = ['.'.join(x.split('.')[:-1]) for x in mesh_unposed_no_shape] pose_shape = ['.'.join(x.split('.')[:-1]) for x in pose_shape] mesh_curvature = ['.'.join(x.split('.')[:-1]) for x in mesh_curvature] mesh_static = ['.'.join(x.split('.')[:-1]) for x in mesh_static] mesh_posed = ['.'.join(x.split('.')[:-1]) for x in mesh_posed] mesh_posed_body_no_orient = ['.'.join(x.split('.')[:-1]) for x in mesh_posed_body_no_orient] mesh_unposed_no_shape = ['.'.join(x.split('.')[:-1]) for x in mesh_unposed_no_shape] pose_shape = ['.'.join(x.split('.')[:-1]) for x in pose_shape] mesh_curvature = ['.'.join(x.split('.')[:-1]) for x in mesh_curvature] assert mesh_posed == mesh_posed_body_no_orient == mesh_unposed_no_shape == pose_shape == mesh_curvature assert set(mesh_posed) == set(mesh_posed_body_no_orient) == set(mesh_unposed_no_shape) == set(pose_shape) == set(mesh_static) == set(mesh_curvature) pass def pack_data(): in_folder = '/home/workspace/local-data/cape_processed/test/female' out_folder = '/home/workspace/local-data/cape_processed/test/female_packed' excluded = set(['mesh_static', 'uv_static_32', 'uv_static_64', 'uv_static_128', 'uv_static_256', 'uv_static', 'mesh_mask', 'uv_mask', 'curvatures_raw']) os.makedirs(out_folder, exist_ok=True) sub_folders = sorted(os.listdir(in_folder)) frames = sorted(os.listdir(os.path.join(in_folder, 'pose_shape'))) for frame in tqdm.tqdm(frames): out = dict() id, frame_nr = frame.split('.')[0], frame.split('.')[1] out['id'] = id for sub_f in sub_folders: if sub_f in excluded: continue if os.path.exists(os.path.join(in_folder, sub_f, id + '.' + frame_nr + '.npz')): data = np.load(os.path.join(in_folder, sub_f, id + '.' + frame_nr + '.npz')) out[sub_f] = data[data.files[0]] else: data = np.load(os.path.join(in_folder, sub_f, id + '.' + frame_nr + '.npy')) out[sub_f] = data np.savez_compressed(os.path.join(out_folder, id + '.' + frame_nr + '.npz'), out) if __name__ == '__main__': phases = ['test', 'test', 'train', 'train'] genders = ['male', 'female', 'male', 'female'] for phase, gender in zip(phases, genders): # get_frames() # split_data() # verify_split_data() # merge_pose_shape() # create_posed_body() # create_static() # create_unposed() # create_posed() # verify() pass pack_data() print()