new files
This commit is contained in:
@@ -0,0 +1,279 @@
|
||||
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
|
||||
Reference in New Issue
Block a user