diff --git a/unpose_one_hot.py b/unpose_one_hot.py new file mode 100644 index 0000000..d96091f --- /dev/null +++ b/unpose_one_hot.py @@ -0,0 +1,230 @@ +import os +import sys +import time +import argparse + +import numpy as np +import scipy.io as sio +import scipy as sp +import plotly.graph_objects as go +from scipy.spatial.transform import Rotation + +from tqdm import tqdm + +from custom_utils.custom_utils import maskFacesBatch, _maskFaces + + +os.environ["CUDA_VISIBLE_DEVICES"] = "1" +import tensorflow as tf +gpus = tf.config.experimental.list_physical_devices('GPU') +if gpus: + try: + for gpu in gpus: + tf.config.experimental.set_memory_growth(gpu, True) + + except RuntimeError as e: + print(e) + + +from custom_utils.smpl import SMPL +from custom_utils.smpl_2 import SMPL as SMPL_2 +from custom_utils.supersmpl_io import readPC2 + + + +# run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='unpose', name='new_smpl_2', notes='') +# artifact = wandb.Artifact('code', type='code') +# artifact.add_file('Data/unpose.py') +# artifact.add_file('Data/verify_unposed.py') +# artifact.add_file('FeatureLines/superSMPL/smpl.py') +# artifact.add_file('FeatureLines/superSMPL/smpl_2.py') +# run.log_artifact(artifact) + + +def read_garment(folder, id): + garment_dir = os.path.join(folder, id) + + 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, id + + garment = garments[0] + + _posed = readPC2(os.path.join(folder, id, garment + '_registered.pc16'), True)['V'] + info = sio.loadmat(os.path.join(folder, id, 'info.mat')) + _pose, _shape, gender = info['poses'].astype(np.float32), info['shape'], info['gender'] + _shape = np.squeeze(np.concatenate((_shape, gender), axis=1)).astype(np.float32) + mask = np.load(os.path.join(folder, id, garment + '_mask.npy')) + # offsets = sio.loadmat(os.path.join(folder, id, garment + '_offsets_static.mat'))['label'][:,:3] + # template = sio.loadmat(os.path.join(folder, id, 'info_body.mat'))['B'] + # mesh_static = (template + offsets) * mask[:, np.newaxis] + + # posed, mesh_static = posed[:, mask, :], mesh_static[mask, :] + mask_idx = np.array(range(14475))[mask] + _shape = np.tile(_shape[np.newaxis, :], [_pose.shape[1], 1]) + + posed = np.zeros((300, 14475, 3), np.float32) + pose = np.zeros((72, 300), np.float32) + shape = np.zeros((300, 11), np.float32) + + frames = _posed.shape[0] + posed[:frames, :, :] = _posed + pose[:, :frames,] = _pose + shape[:frames, :] = _shape + + return posed, mask, pose, shape, mask_idx + + +def plot_3d(mesh_orig, mesh_pred, unposed, faces): + _mesh_pred = mesh_pred.copy() + _unposed = unposed.copy() + _mesh_pred[:,1] += 1 + _unposed[:,1] += 3 + fig = go.Figure(data=[ + go.Mesh3d( + # 8 vertices of a cube + x=mesh_orig[:,0], + y=mesh_orig[:,1], + z=mesh_orig[:,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=_mesh_pred[:,0], + y=_mesh_pred[:,1], + z=_mesh_pred[:,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=_unposed[:,0], + y=_unposed[:,1], + z=_unposed[:,2], + # i, j and k give the vertices of triangles + i = faces[:,0], + j = faces[:,1], + k = faces[:,2], + # vertexcolor=collision_color + ) + ]) + + fig.update_layout(scene_aspectmode='manual', + scene_aspectratio=dict(x=1, y=2, z=1)) + fig.update_layout( + scene = dict( + xaxis = dict(nticks=4, range=[-1,1],), + yaxis = dict(nticks=4, range=[-1,3],), + zaxis = dict(nticks=4, range=[-1,1],),) + ) + + return fig + + +@tf.function +def smpl_run_wrapper(R, shape): + return smpl.run(R, shape) + +@tf.function +def smpl_run_backward_wrapper(Rt, J, posed): + return smpl.run_backward(Rt, J, posed, smpl.weights_one_hot) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument("-f", "--folder", help = "end of range", default='/home/workspace/data/dataset/test_skirt') + parser.add_argument("-s", "--start", help = "start of range", default=0) + parser.add_argument("-e", "--end", help = "end of range", default=1500) + parser.add_argument("-n", "--name", help = "end of range", default='lhunor') + args = parser.parse_args() + + F = np.load('custom_utils/smpl_faces_skirt.npy') + folder = args.folder + smpl = SMPL() + # smpl_2 = SMPL_2('custom_utils') + + is_skirt = True + + ids = os.listdir(folder)[int(args.start):int(args.end)] + + for id in tqdm(ids): + if not os.path.isdir(os.path.join(folder, id)): + continue + + # if len(id) == 5: + # continue + + os.makedirs(os.path.join(folder, id, 'mesh_unposed_one_hot'), exist_ok=True) + + posed, mask, pose, shape, mask_idx = read_garment(folder, id) + # gender = (shape[:, -1] > 0).astype(np.int32) + # init = Rotation.from_euler('zyz', [90,90,90], degrees=True).apply(mesh_static.copy()).astype(np.float32) + + # existing_frames = os.listdir(os.path.join(folder, id, 'mesh_unposed')) + # if len(existing_frames) == posed.shape[0]: + # continue + + frames = posed.shape[0] + + # R = smpl.run_rodrigues(np.transpose(pose, [1,0])) + + if is_skirt: + pose_dress = np.copy(pose) + p_avg = (pose[:, 3:6] + pose[:, 6:9]) / 2 + pose_dress[:,3:9] = np.concatenate([p_avg, p_avg], axis=-1) + R = smpl.run_rodrigues(np.transpose(pose_dress, [1,0])).numpy() + for b in range(frames): + R[b][[4,5,7,8,10,11],:,:] = np.eye(3) + else: + R = smpl.run_rodrigues(np.transpose(pose, [1,0])) + + Rt = tf.transpose(R, [0, 1, 3, 2]) + T_body, J, _ = smpl_run_wrapper(R, shape) + # zero_body = smpl_run_wrapper(smpl.run_rodrigues(np.zeros((1, 72), np.float32)), shape[:1, :])[0][0].numpy() + + + T0_garment, J = smpl_run_backward_wrapper(Rt, J, posed) + # T_garment, _ = smpl.run_garment_orig(R, shape, T0_garment) # =====> you do not need this and next line for unposing + # metric = tf.reduce_max(tf.abs(mask[tf.newaxis,:,tf.newaxis] * (T_garment - posed))) + + + for frame in range(frames): + np.savez_compressed(os.path.join(folder, id, 'mesh_unposed_one_hot', id + '_' + str(frame) + '.npz'), mesh_unposed=T0_garment[frame]) + + # F = np.load("custom_utils/smpl_faces" + ("_skirt" if is_skirt else "") + ".npy") + # _, faces = _maskFaces(np.zeros((14475, 3)), F, mask.astype(bool)) + + # batch = 0 + # x = T0_garment[batch].numpy()[mask, :] + # y = zero_body + # 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 = F[:,0], + # j = F[:,1], + # k = F[:,2], + # )], + # ).show() + + + # print() + +