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.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 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] return posed, mask, pose, shape, mesh_static, 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 if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument("-f", "--folder", help = "end of range", default='/home/workspace/data/dataset/train_skirt') parser.add_argument("-s", "--start", help = "start of range", default=0) parser.add_argument("-e", "--end", help = "end of range", default=1) 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') 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 # if len(id) == 5: # continue os.makedirs(os.path.join(folder, id, 'mesh_unposed'), exist_ok=True) posed, mask, pose, shape, mesh_static, mask_idx = read_garment(folder, id) gender = int(shape[-1] > 0) 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 R = smpl_2.run_rodrigues(np.transpose(pose, [1,0])) frames = posed.shape[0] # rot = Rotation.from_euler('z', [90], degrees=True) # rest_pose = np.zeros((1,72), dtype=np.float32) # R_rest = smpl.run_rodrigues(rest_pose) # body = smpl.run(R_rest, shape[tf.newaxis,:])[0].numpy()[:, mask,:] # # verts = smpl.run(R, tf.tile(shape[tf.newaxis,:], [frames,1]))[0].numpy()[:,mask,:] # # verts = tf.stack([rot.apply(verts[i]).astype(np.float32) for i in range(frames)]) # verts = tf.tile(body, [frames,1,1]) # start = time.perf_counter() # result = smpl.run_garment(R, tf.tile(shape[tf.newaxis,:], [frames,1]), verts, filter_idx=mask_idx) # orig_time = time.perf_counter() - start # print(orig_time) # start = time.perf_counter() # result_2 = smpl_2.run_garment(R, shape[:10], gender, verts, filter_idx=mask_idx) # new_time = time.perf_counter() - start # print(new_time) # # test = np.abs(result.numpy() - result_2.numpy()) > 1e-4 # # x = np.sum(test, axis=-1).sum(-1).sum(-1) # euclid_error = tf.reduce_mean(tf.math.sqrt(tf.reduce_sum((result - result_2) ** 2, axis=-1))).numpy() # print('speedup: ' + str(orig_time/new_time) + ' euclid_error: ' + str(euclid_error)) # # continue # batch = np.random.choice(np.array(range(frames))) # _, faces = _maskFaces(np.zeros((14475,3)), F, mask) # fig = go.Figure(data=[ # go.Mesh3d( # # 8 vertices of a cube # x=result[batch,:,0], # y=result[batch,:,1], # z=result[batch,:,2], # # i, j and k give the vertices of triangles # i = faces[:,0], # j = faces[:,1], # k = faces[:,2], # text='speedup: ' + str(orig_time/new_time) + ' euclid_error: ' + str(euclid_error), # ), # go.Mesh3d( # # 8 vertices of a cube # x=result_2[batch,:,0], # y=result_2[batch,:,1], # z=result_2[batch,:,2], # # i, j and k give the vertices of triangles # i = faces[:,0], # j = faces[:,1], # k = faces[:,2], # text='speedup: ' + str(orig_time/new_time) + ' euclid_error: ' + str(euclid_error), # )], # ) # wandb.log({id: fig}) # continue lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(1000.0, decay_steps=40000, decay_rate=0.1, staircase=False) optimizer = tf.keras.optimizers.SGD(learning_rate=lr_schedule) unposed = tf.Variable(tf.tile(init[tf.newaxis,:,:], [frames,1,1])) @tf.function def skinning(verts): result = smpl_2.run_garment(R, shape[:10], gender, verts, filter_idx=mask_idx) return result @tf.function def error(): pred = skinning(unposed) loss = tf.reduce_mean(tf.math.abs(pred - posed)) return loss start = time.perf_counter() euclid_error = 10000 while True: optimizer.minimize(error, [unposed]) i = optimizer.iterations.numpy() if i % 1000 == 0: pred = skinning(unposed) # loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(pred - posed) ** 2, axis=-1)) loss = tf.reduce_mean(tf.math.abs(pred - posed)) euclid_error = tf.reduce_mean(tf.math.sqrt(tf.reduce_sum((posed - 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 > 15000: break if euclid_error < 0.001 and i > 30000: break if i > 40000: break # b = np.random.choice(np.array(range(frames))) # _, faces = _maskFaces(np.zeros((14475,3)), F, mask) # plot_3d(posed[b], pred.numpy()[b], unposed.numpy()[b], faces).show() result = np.zeros((frames, 14475, 3), dtype = np.float32) result[:, mask,:] = unposed.numpy() for frame in range(frames): np.savez_compressed(os.path.join(folder, id, 'mesh_unposed', id + '_' + str(frame) + '.npz'), mesh_unposed=result[frame]) # batch = 0 # go.Figure(data=[ # go.Mesh3d( # # 8 vertices of a cube # x=posed[batch,:,0], # y=posed[batch,:,1], # z=posed[batch,:,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=pred[batch,:,0], # y=pred[batch,:,1], # z=pred[batch,:,2], # # i, j and k give the vertices of triangles # i = faces[:,0], # j = faces[:,1], # k = faces[:,2], # )], # ).show() continue offset = 1 for frame in range(posed.shape[0]): lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(10.0, decay_steps=10000, decay_rate=0.1, staircase=False) optimizer = tf.keras.optimizers.SGD(learning_rate=lr_schedule) unposed = tf.Variable(np.random.rand(offset, posed.shape[1], 3).astype(np.float32) * 2 - 1) # unposed = tf.Variable(init) @tf.function def skinning(verts): # verts = smpl.run(R[frame:frame+offset], tf.tile(shape[tf.newaxis,:], [offset,1]))[0].numpy()[:,mask,:] # start = time.perf_counter() result = smpl.run_garment(R[frame:frame+offset], tf.tile(shape[tf.newaxis,:], [offset,1]), verts, filter_idx=mask_idx) # print(time.perf_counter() - start) # start = time.perf_counter() # result_2 = smpl_2.run_garment(R[frame:frame+offset], shape[:10], gender, verts, filter_idx=mask_idx) # print(time.perf_counter() - start) # test = np.abs(result.numpy() - result_2.numpy()) > 1e-4 # x = np.sum(test, axis=-1).sum(-1).sum(-1) # euclid_error = tf.reduce_mean(tf.math.sqrt(tf.reduce_sum((result - result_2) ** 2, axis=-1))) # batch = 25 # _, faces = _maskFaces(np.zeros((14475,3)), F, mask) # go.Figure(data=[ # go.Mesh3d( # # 8 vertices of a cube # x=result[batch,:,0], # y=result[batch,:,1], # z=result[batch,:,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=result_2[batch,:,0], # y=result_2[batch,:,1], # z=result_2[batch,:,2], # # i, j and k give the vertices of triangles # i = faces[:,0], # j = faces[:,1], # k = faces[:,2], # )]).show() return result @tf.function def error(): pred = skinning(unposed) loss = tf.reduce_mean(tf.math.abs(pred - posed[frame])) return loss start = time.perf_counter() euclid_error = 10000 for i in range(iterations): optimizer.minimize(error, [unposed]) if i % 50 == 0: pred = skinning(unposed) # loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(pred - posed) ** 2, axis=-1)) loss = tf.reduce_mean(tf.math.abs(pred - posed[frame])) euclid_error = tf.reduce_mean(tf.math.sqrt(tf.reduce_sum((posed[frame] - pred) ** 2, axis=-1))) print('id: ' + id + ' frame: ' + str(frame) + ' 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.001: break result = np.zeros((14475, 3), dtype = np.float32) result[mask,:] = unposed.numpy() np.savez_compressed(os.path.join(folder, id, 'mesh_unposed', id + '_' + str(frame) + '.npz'), mesh_unposed=result) init = unposed # pred, unposed = pred.numpy(), unposed.numpy() # _, faces = _maskFaces(np.zeros((14475,3)), F, mask) # k = 0 # plot_3d(posed[k], pred[k], unposed[k], faces).show() # print()