import os import sys import pickle import time import json import math import wandb import argparse import numpy as np import tensorflow as tf import scipy.io as sio import scipy as sp import matplotlib.pyplot as plt import plotly.graph_objects as go from scipy.spatial.transform import Rotation from tqdm import tqdm from datetime import datetime sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') from FeatureLines.superSMPL.IO import readPC2, writePC2 from Data.Dataloader import DataLoader from Models.Reconstructor import Reconstructor_v2 as Reconstructor from Models.FeatureLineGenerator import FeatureLineGenerator from Models.losses import L2_loss, bend_loss, collision_loss, normal_loss, L1_loss, feature_lines_loss, edge_loss, normal_loss_v2 # from Models.metrics import L2_error, surf2surf, feature_lines_error from Models.utils import maskFacesBatch, pretty_json, repeat_faces, repeat_laplacian, sparse_to_tensor, getFlattenedSortedUniqueFeatureLineIndices, _maskFaces from FeatureLines.superSMPL.smpl import SMPL from FeatureLines.superSMPL.smpl_2 import SMPL as SMPL_2 os.environ["CUDA_VISIBLE_DEVICES"] = "6" 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) # 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") 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) 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.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]) 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()