166 lines
6.5 KiB
Python
166 lines
6.5 KiB
Python
import os
|
|
|
|
import os
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = '1'
|
|
import tensorflow as tf
|
|
|
|
gpus = tf.config.experimental.list_physical_devices('GPU')
|
|
if gpus:
|
|
try:
|
|
# Currently, memory growth needs to be the same across GPUs
|
|
for gpu in gpus:
|
|
tf.config.experimental.set_memory_growth(gpu, True)
|
|
logical_gpus = tf.config.experimental.list_logical_devices('GPU')
|
|
print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
|
|
except RuntimeError as e:
|
|
# Memory growth must be set before GPUs have been initialized
|
|
print(e)
|
|
|
|
|
|
import sys
|
|
import time
|
|
import cv2
|
|
import random
|
|
import numpy as np
|
|
import cmapy as cm
|
|
import plotly.graph_objects as go
|
|
import plotly.express as ex
|
|
|
|
from dataset import DatasetFactory, StaticDataset
|
|
from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces
|
|
from custom_utils.smpl import SMPL
|
|
|
|
if __name__ == '__main__':
|
|
|
|
uv2mesh = UV2Mesh(is_smpl_orig=True)
|
|
smpl = SMPL(orig_smpl=True)
|
|
|
|
F = np.load("custom_utils/smpl_faces_orig.npy")
|
|
|
|
factory = DatasetFactory()
|
|
signature_dynamic = { 'id': (),
|
|
'pose_shape': (72 + 11,),
|
|
# 'uv_posed': (3, 512, 256,),
|
|
# 'uv_body_posed': (3, 512, 256,),
|
|
'mesh_body_posed': (3638, 3,),
|
|
'mesh_posed': (3638, 3,),
|
|
'mesh_unposed': (3638, 3,),
|
|
# 'uv_unposed': (3, 512, 256,),
|
|
# 'uv_curvature': (512, 256),
|
|
# 'mesh_curvature': (14475)
|
|
}
|
|
signature_static = { #'uv_static': (3, 512, 256,),
|
|
# 'mesh_static': (14475, 3,),
|
|
'mesh_mask': (3638,),
|
|
# 'uv_mask': (512, 256,),
|
|
# 'uv_shaped_body': (3, 512, 256,),
|
|
}
|
|
|
|
dataset = factory.get_dataset(folder='/home/workspace/local-data/cape_processed/train/female',
|
|
batch_size=1,
|
|
signature=signature_dynamic,
|
|
resolutions=[512],
|
|
is_cape=True)
|
|
static_dataset = StaticDataset(folder='/home/workspace/local-data/cape_processed/train/female',
|
|
signature=signature_static,
|
|
resolutions=[512],
|
|
is_cape=True)
|
|
|
|
|
|
|
|
for i, batch in enumerate(dataset):
|
|
print(i)
|
|
batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
|
|
|
|
J0_old = np.load('custom_data/smpl_root_joint.npy')
|
|
|
|
id = batch['id'][0].numpy().decode("utf-8")
|
|
mesh_unposed = np.squeeze(batch['mesh_unposed'].numpy())
|
|
mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool)
|
|
pose, shape = batch['pose_shape'][:, :72].numpy(), batch['pose_shape'][:, 72:].numpy()
|
|
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
|
|
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])]).numpy()
|
|
mesh_unposed_shaped = batch['mesh_unposed'] + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
|
|
J0_new = smpl.run(smpl.run_rodrigues(np.zeros((1, 72), np.float32)), shape)[1][0,0,:].numpy()
|
|
|
|
min_body = np.load(os.path.join('/home/workspace/local-data/temp/cape_release', 'minimal_body_shape', id.split('_')[0], id.split('_')[0] + '_minimal.npy'))[np.newaxis, :, :].astype(np.float32)
|
|
J0_2 = smpl.get_root_joint(shape, v_shaped=tf.stack([min_body, min_body], axis=-1))[0].numpy()
|
|
# min_body -= J0_2[np.newaxis,np.newaxis,:]
|
|
|
|
mesh_unposed_shaped = mesh_unposed_shaped + J0_old[tf.newaxis,tf.newaxis,:]
|
|
mesh_reposed = smpl.run_garment(pose, shape, mesh_unposed_shaped, v_shaped=tf.stack([min_body, min_body], axis=-1)).numpy()[0]
|
|
mesh_posed = np.squeeze(batch['mesh_posed'].numpy()) + J0_old[np.newaxis,:]
|
|
mesh_body_posed_no_orient = np.squeeze(batch['mesh_body_posed'].numpy())
|
|
mesh_body_posed = smpl.run(smpl.run_rodrigues(pose), shape)[0][0].numpy()
|
|
|
|
# x = smpl.run(smpl.run_rodrigues(pose), shape)[0][0]
|
|
x = mesh_reposed[mask, :]
|
|
y = mesh_posed[mask, :]
|
|
_, faces = _maskFaces(np.zeros((3638, 3), dtype=np.float32), F, mask)
|
|
fig = go.Figure(data=[
|
|
go.Mesh3d(
|
|
x=x[:,0],
|
|
y=x[:,1],
|
|
z=x[:,2],
|
|
i = faces[:,0],
|
|
j = faces[:,1],
|
|
k = faces[:,2]
|
|
),
|
|
go.Mesh3d(
|
|
x=y[:,0],
|
|
y=y[:,1],
|
|
z=y[:,2],
|
|
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()
|
|
|
|
|
|
x = mesh_unposed_shaped.numpy()[0, mask, :]
|
|
# y = smpl.run(smpl.run_rodrigues(np.zeros((1, 72), np.float32)), shape)[0][0]
|
|
y = min_body[0, smpl.to_keep]
|
|
_, faces = _maskFaces(np.zeros((3638, 3), dtype=np.float32), F, mask)
|
|
fig = go.Figure(data=[
|
|
go.Mesh3d(
|
|
x=x[:,0],
|
|
y=x[:,1],
|
|
z=x[:,2],
|
|
i = faces[:,0],
|
|
j = faces[:,1],
|
|
k = faces[:,2]
|
|
),
|
|
go.Mesh3d(
|
|
x=y[:,0],
|
|
y=y[:,1],
|
|
z=y[:,2],
|
|
i = F[:,0],
|
|
j = F[:,1],
|
|
k = F[:,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() |