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 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.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): min_body = np.load(os.path.join(folder, id, id + '_minimal.npy')) return min_body 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 def verify(): folder = '/home/workspace/data/dataset/cape_release/minimal_body_shape' F = np.load('/home/workspace/data/dataset/cape_release/misc/smpl_tris.npy') gender_table = np.load('/home/workspace/data/dataset/cape_release/misc/subj_genders.pkl', allow_pickle=True) smpl = SMPL(orig_smpl=True) ids = os.listdir(folder) for id in tqdm(ids): if not os.path.isdir(os.path.join(folder, id)): continue orig_body = read_garment(folder, id) pose = tf.constant(np.zeros((1, 72)), dtype=np.float32) R = smpl.run_rodrigues(pose) shape = np.load(os.path.join(folder, id, 'betas.npy')) gender = [0] if gender_table[id] == 'female' else [1] shape = np.concatenate((shape, gender)).astype(np.float32) new_body = smpl.run(R, shape[np.newaxis,:], remove=False)[0][0] euclid_error = tf.reduce_mean(tf.math.sqrt(tf.reduce_sum((orig_body[smpl.to_keep,:] - new_body.numpy()[smpl.to_keep,:]) ** 2, axis=-1))).numpy() x = orig_body y = new_body faces = F 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() if __name__ == '__main__': verify() parser = argparse.ArgumentParser() parser.add_argument("-f", "--folder", help = "end of range", default='/home/workspace/data/dataset/cape_release/minimal_body_shape') parser.add_argument("-s", "--start", help = "start of range", default=0) parser.add_argument("-e", "--end", help = "end of range", default=200) parser.add_argument("-n", "--name", help = "end of range", default='lhunor') args = parser.parse_args() F = np.load('/home/workspace/data/dataset/cape_release/misc/smpl_tris.npy') F = np.load('custom_utils/smpl_faces_orig.npy') gender_table = np.load('/home/workspace/data/dataset/cape_release/misc/subj_genders.pkl', allow_pickle=True) folder = args.folder smpl = SMPL(orig_smpl=True) iterations = 10000 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 min_body = read_garment(folder, id) min_body = min_body[smpl.to_keep, :] pose = tf.constant(np.zeros((1, 72)), dtype=np.float32) R = smpl.run_rodrigues(pose) zero_body = smpl.run(R, np.zeros((1, 11), np.float32), remove=False)[0] lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(1.0, decay_steps=40000, decay_rate=0.1, staircase=False) optimizer = tf.keras.optimizers.SGD(learning_rate=1.0) betas = tf.Variable(np.zeros((1, 10), np.float32)) # gender = np.zeros((1,1), np.float32) #F gender = 0 if gender_table[id] == 'female' else 1 @tf.function def shaping(shape): result = (tf.tensordot(shape, smpl.shapedirs[gender], axes=[[1], [2]]) + tf.reshape(smpl.v_template[gender], shape=(1,6890,3)))[0] result = tf.gather(result, smpl.to_keep, axis=0) return result @tf.function def error(): pred = shaping(betas) loss = tf.reduce_mean(tf.math.abs(pred - min_body)) return loss start = time.perf_counter() euclid_error = 10000 while True: optimizer.minimize(error, [betas]) i = optimizer.iterations.numpy() if i % 100 == 0: pred = shaping(betas) # loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(pred - posed) ** 2, axis=-1)) loss = tf.reduce_mean(tf.math.abs(pred - min_body)) euclid_error = tf.reduce_mean(tf.math.sqrt(tf.reduce_sum((min_body - pred) ** 2, axis=-1))) print('id: ' + id + ' iteration: ' + str(i) + ' loss: ' + str(loss.numpy()) + ' error: ' + str(euclid_error.numpy()) + ' time: ' + str(time.perf_counter() - start)) start = time.perf_counter() if euclid_error < 0.0001: break if euclid_error < 0.0005 and i > 20000: break if euclid_error < 0.001 and i > 40000: break if i > 60000: break x = min_body y = pred faces = F 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() np.save(os.path.join(folder, id, 'betas_2.npy'), betas[0]) continue