import os import sys import pickle import time import json import math import argparse import numpy as np import tensorflow as tf 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 datetime import datetime from custom_utils.supersmpl_io import readPC2, writePC2 from custom_utils.smpl import SMPL from custom_utils.smpl_2 import SMPL as SMPL_2 os.environ["CUDA_VISIBLE_DEVICES"] = "1" 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) 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 def unpose(mesh_static, posed, pose, shape, gender, mask, mask_idx, smpl_2, id): 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]: # return R = smpl_2.run_rodrigues(np.transpose(pose, [1,0])) frames = posed.shape[0] 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() return result if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument("-f", "--folder", help = "end of range") parser.add_argument("-s", "--start", help = "start of range") parser.add_argument("-e", "--end", help = "end of range") parser.add_argument("-n", "--name", help = "end of range") args = parser.parse_args() F = np.load('FeatureLines/superSMPL/smpl_faces.npy') folder = args.folder smpl = SMPL('FeatureLines/superSMPL') smpl_2 = SMPL_2('FeatureLines/superSMPL') 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 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) unposed = unpose(mesh_static, posed, pose, shape, gender, mask, mask_idx, smpl_2) for frame in range(unposed.shape[0]): np.savez_compressed(os.path.join(folder, id, 'mesh_unposed', id + '_' + str(frame) + '.npz'), mesh_unposed=unposed[frame]) 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()