import glob import os from data.supersmpl_io import readPC2 os.environ["CUDA_VISIBLE_DEVICES"] = "3" import sys import shutil import tqdm import random import math import numpy as np import scipy as sp import scipy.io as sio from custom_utils.smpl import SMPL import plotly.graph_objects as go from custom_utils.custom_utils import _maskFaces def create_no_orient_body(folder): ids = os.listdir(folder) F = np.load("custom_utils/smpl_faces.npy") smpl = SMPL() for id in tqdm.tqdm(ids): if not os.path.isdir(os.path.join(folder, id)): continue # if os.path.exists(os.path.join(folder, id, 'mesh_posed_body_no_orient')): # continue info = sio.loadmat(os.path.join(folder, id, 'info.mat')) pose, shape, gender = info['poses'], info['shape'], info['gender'] # if gender == 1: # continue pose = np.transpose(pose, [1, 0]).astype(np.float32) if len(os.listdir(os.path.join(folder, id, 'mesh_posed_body_no_orient'))) == pose.shape[0]: continue pose[:, :3] = 0 beta = np.squeeze(np.concatenate((shape, gender), axis=1)).astype(np.float32) R = smpl.run_rodrigues(pose) T, _, _ = smpl.run(R, np.tile(beta[np.newaxis, :], [pose.shape[0], 1]), True) T = T.numpy() os.makedirs(os.path.join(folder, id, 'mesh_posed_body_no_orient'), exist_ok=True) for i in range(pose.shape[0]): np.savez_compressed(os.path.join(folder, id, 'mesh_posed_body_no_orient', id + '_' + str(i) + '.npz'), mesh_posed_body_no_orient = T[i, :, :]) def create_no_shape_unposed(folder): import tensorflow as tf ids = os.listdir(folder) smpl = SMPL() for id in tqdm.tqdm(ids): garment_dir = os.path.join(folder, id) if not os.path.isdir(garment_dir): continue # if os.path.exists(os.path.join(folder, id, 'mesh_unposed_no_shape')): # continue info = sio.loadmat(os.path.join(folder, id, 'info.mat')) # if len(os.listdir(os.path.join(folder, id, 'mesh_unposed_no_shape_one_hot'))) == info['poses'].shape[1]: # continue shape = info['shape'] gender = int(info['gender'][0,0] > 0) garments = [f.replace('.pc16','') for f in os.listdir(garment_dir) if f.find('_') == -1 and f != 'info.mat' and not f.endswith('.obj')] assert len (garments) == 1 assert gender == info['gender'][0,0] # if len (garments) != 1: # continue garment = garments[0] mask = np.load(os.path.join(folder, id, garment + '_mask.npy')) shape_offsets = tf.tensordot(shape.astype(np.float32), smpl.shapedirs[gender], axes=[[1], [2]]).numpy() os.makedirs(os.path.join(folder, id, 'mesh_unposed_no_shape_one_hot'), exist_ok=True) for i in range(info['poses'].shape[1]): mesh_unposed = np.load(os.path.join(garment_dir, 'mesh_unposed_one_hot', id + '_' + str(i) + '.npz')) mesh_unposed = mesh_unposed['mesh_unposed'] mesh_unposed[mask, :] = mesh_unposed[mask, :] - np.squeeze(shape_offsets)[smpl.to_keep, :][mask, :] np.savez_compressed(os.path.join(folder, id, 'mesh_unposed_no_shape_one_hot', id + '_' + str(i) + '.npz'), mesh_unposed_no_shape = mesh_unposed) def check_if_skirt(folder): ids = os.listdir(folder) for id in tqdm.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 f != 'info.mat' and not f.endswith('.obj')] # assert len (garments) == 1 # garment = garments[0] # assert garment != 'Skirt' assert "Skirt" not in garments def gather_ids(folder, file_name): ids = os.listdir(folder) to_print = [] for id in tqdm.tqdm(ids): garment_dir = os.path.join(folder, id) if not os.path.isdir(garment_dir): continue if len(id) == 5: to_print.append(id) to_print.sort() with open(file_name, 'w') as f: f.write('\n'.join(to_print)) def verify_poses(folder): ids = os.listdir(folder) for id in tqdm.tqdm(ids): if not os.path.isdir(os.path.join(folder, id)): continue garments = [f.replace('_registered.pc16','') for f in os.listdir(os.path.join(folder, id)) if f.endswith('_registered.pc16') and not f.startswith('Combined')] assert len (garments) == 1 info = sio.loadmat(os.path.join(folder, id, 'info.mat')) poses = info['poses'] verts = readPC2(os.path.join(folder, id, garments[0] + '_registered.pc16'), True)['V'] # print(id) # assert type(poses) == np.ndarray and poses.shape[0] == 72 and poses.shape[1] == verts.shape[0], 'mv ' + os.path.join(folder, id) + ' /home/workspace/data/dataset/invalid/' if not (type(poses) == np.ndarray and poses.shape[0] == 72 and poses.shape[1] == verts.shape[0]): shutil.move(os.path.join(folder, id), '/home/workspace/data/dataset/invalid/') print(id) def fix_split_garment_generated_data(folder): ids = os.listdir(folder) for id in tqdm.tqdm(ids): if not os.path.isdir(os.path.join(folder, id)): continue garments = [f.replace('_registered.pc16','') for f in os.listdir(os.path.join(folder, id)) if f.endswith('_registered.pc16') and not f.startswith('Combined')] assert len (garments) == 1 if len(id) == 6 and id.endswith('2'): frames_folder = os.path.join(folder, id, 'mesh_posed_body_no_orient') # os.path.join(folder, id, 'mesh_unposed_no_shape') frames = os.listdir(frames_folder) frame_id = frames[0].split('_')[0] if len(frame_id) == 6: assert frame_id == id, 'Wrong id found!' if len(frame_id) == 5 and id.startswith(frame_id): for frame in frames: frame_nr = frame.split('.')[0].split('_')[1] shutil.move(os.path.join(frames_folder, frame), os.path.join(frames_folder, id + '_' + frame_nr + '.npz')) def create_garment_type_table(): folder = '/home/workspace/data/dataset' # sub_folders = ['train', 'validation', 'test', 'train_skirt', 'valid_skirt', 'test_skirt'] # sub_folders = ['train', 'validation', 'train_skirt', 'valid_skirt'] sub_folders = ['test_skirt'] # sub_folders = ['test', 'test_skirt'] garment_type_table = dict() for sub_f in sub_folders: ids = os.listdir(os.path.join(folder, sub_f)) for id in tqdm.tqdm(ids): if not os.path.isdir(os.path.join(folder, sub_f, id)): continue garments = [f.replace('_registered.pc16','') for f in os.listdir(os.path.join(folder, sub_f, id)) if f.endswith('_registered.pc16') and not f.startswith('Combined')] assert len (garments) == 1 if id in garment_type_table.keys(): print(id + ': ' + garments[0] + ' ' + garment_type_table[id]) garment_type_table[id] = garments[0] np.save('custom_data/garment_type_table_test.npy', garment_type_table) def generate_faces(): folder = '/home/workspace/data/dataset/test_skirt_ready' output_folder = '/home/workspace/data/results/support/faces_skirt' F = np.load("custom_utils/smpl_faces_skirt.npy") os.makedirs(output_folder, exist_ok=True) ids = os.listdir(os.path.join(folder, 'pose_shape')) ids = [id.split('_')[0] for id in ids] ids = set(ids) for id in tqdm.tqdm(ids): mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy')).astype(bool) _, faces = _maskFaces(np.zeros((14475, 3), np.float32), F, mask) np.save(os.path.join(output_folder, id + '.npy'), faces.astype(np.int32)) def generate_frames_file(folder, ratio=0.1): all_frames = os.listdir(os.path.join(folder, 'pose_shape')) ids = [frame.split('_')[0] for frame in all_frames] all_frames = [frame.split('.')[0] for frame in all_frames] ids = sorted(list(set(ids))) garment_type_table = np.load('custom_data/garment_type_table_train.npy', allow_pickle=True).item() random.shuffle(ids) sub_ids = ids[:int(len(ids) * ratio)] selected = [] for id in tqdm.tqdm(sub_ids): if garment_type_table[id] != 'Dress': continue if len(selected) > 30000: break # frames = glob.glob(os.path.join(folder, 'pose_shape', id + '_*')) frames = [frame for frame in all_frames if frame.startswith(id + '_')] frames = [frame.split('/')[-1].split('.')[0] for frame in frames] selected.extend(frames) with open(os.path.join(folder, 'frames_dress_30k.txt'), 'w') as f: for i, frame in enumerate(selected): f.write(frame + '\n') def generate_frames_file_filtered_cape(): folder = '/home/workspace/local-data/cape_processed/train/female' all_frames = os.listdir(os.path.join(folder, 'pose_shape')) ids = [frame.split('.')[0] for frame in all_frames] all_frames = ['.'.join(frame.split('.')[:-1]) for frame in all_frames] ids = sorted(list(set(ids))) seqs_train_00032_SS = ['00032_shortshort_hips', '00032_shortshort_shoulders_mill'] seqs_test_00032_SS = [ '00032_shortshort_tilt_twist_left'] seqs_train_00032_SL = ['00032_shortlong_hips', '00032_shortlong_shoulders_mill'] seqs_test_00032_SL = [ '00032_shortlong_tilt_twist_left'] seqs_train_00032_SL_2 = ['00032_shortlong_hips', '00032_shortlong_tilt_twist_left'] seqs_test_00032_SL_2 = ['00032_shortlong_shoulders_mill'] seqs_train_03223_SL_2 = ['03223_shortlong_hips', '03223_shortlong_tilt_twist_left'] seqs_test_03223_SL_2 = ['03223_shortlong_shoulders_mill'] seqs_train_03223_SS_2 = ['03223_shortshort_hips', '03223_shortshort_tilt_twist_left'] seqs_test_03223_SS_2 = ['03223_shortshort_shoulders_mill'] ids = [id for id in ids if id in seqs_train_03223_SS_2] random.shuffle(ids) sub_ids = ids selected = [] for id in tqdm.tqdm(sub_ids): # frames = glob.glob(os.path.join(folder, 'pose_shape', id + '_*')) frames = [frame for frame in all_frames if frame.startswith(id + '.')] # frames = [frame.split('/')[-1].split('.')[0] for frame in frames] selected.extend(frames) random.shuffle(selected) with open(os.path.join(folder, 'frames_03223_SS_train_2.txt'), 'w') as f: for i, frame in enumerate(selected): f.write(frame + '\n') if __name__ == '__main__': src_folder = '/home/workspace/local-data/train_skirt' #cloth3d_skirt # src_folder = r'D:\Work\Dataset\validation_2' dest_folder = '/home/workspace/data/dataset/validation' # create_no_orient_body(src_folder) # fix_split_garment_generated_data(src_folder) # create_no_shape_unposed(src_folder) # check_if_skirt(src_folder) # gather_ids(src_folder, 'ids/test_non_skirt_ids.txt') # verify_poses(src_folder) # create_garment_type_table() # generate_faces() generate_frames_file(src_folder, ratio=1) # generate_frames_file_filtered_cape()