original unpose sript
This commit is contained in:
@@ -0,0 +1,388 @@
|
||||
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()
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user