105 lines
4.0 KiB
Python
105 lines
4.0 KiB
Python
|
|
|
|
import os
|
|
import sys
|
|
|
|
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = '6'
|
|
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 glob
|
|
import tqdm
|
|
import copy
|
|
import numpy as np
|
|
from custom_utils.smpl import SMPL
|
|
|
|
# from custom_utils.custom_utils import UV2Mesh, _maskFaces
|
|
import plotly.graph_objects as go
|
|
# import plotly.express as ex
|
|
from scipy.spatial.transform import Rotation as R
|
|
|
|
|
|
if __name__ == '__main__':
|
|
folder = '/home/workspace/local-data/test_skirt_ready'
|
|
# folder = '/home/workspace/data/dataset/test_ready'
|
|
# uv2mesh = UV2Mesh(is_skirt=False)
|
|
smpl = SMPL(use_one_hot=False)
|
|
|
|
F = np.load("custom_utils/smpl_faces.npy")
|
|
|
|
ids = os.listdir(os.path.join(folder, 'pose_shape'))
|
|
ids = [id.split('_')[0] for id in ids]
|
|
ids = set(ids)
|
|
os.makedirs(os.path.join(folder, 'mesh_posed_no_orient'), exist_ok=True)
|
|
|
|
for id in tqdm.tqdm(ids):
|
|
|
|
frames = glob.glob(os.path.join(folder, 'pose_shape', id + '_*'))
|
|
frames = len(frames)
|
|
|
|
mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy')).astype(bool)
|
|
mesh_posed = [np.load(os.path.join(folder, 'mesh_posed', id + '_' + str(i) + '.npz')) for i in range(frames)]
|
|
mesh_posed = [mesh_posed[i]['mesh_posed'] for i in range(frames)]
|
|
# mesh_unposed = [np.load(os.path.join(folder, 'mesh_unposed', id + '_' + str(i) + '.npz')) for i in range(frames)]
|
|
# mesh_unposed = [mesh_unposed[i]['mesh_unposed'] for i in range(frames)]
|
|
pose_shape = [np.load(os.path.join(folder, 'pose_shape', id + '_' + str(i) + '.npy')).astype(np.float32) for i in range(frames)]
|
|
# pose_shape_no_orient = copy.deepcopy(pose_shape)
|
|
# for i in range(len(pose_shape)):
|
|
# pose_shape_no_orient[i][:3] = 0
|
|
|
|
rodr = smpl.run_rodrigues(np.stack([x[:72] for x in pose_shape]).astype(np.float32)).numpy()
|
|
|
|
for frame in range(frames):
|
|
r = R.from_matrix(rodr[frame, 0, :, :])
|
|
# assert r.as_matrix() == rodr[frame, 0, :, :]
|
|
out = r.apply(mesh_posed[frame], inverse=True)
|
|
|
|
# x = smpl.run_garment(pose_shape_no_orient[frame][:72][np.newaxis,:], pose_shape_no_orient[frame][72:][np.newaxis,:], mesh_unposed[frame][np.newaxis,:,:])[0].numpy()[mask, :]
|
|
# y = out[mask, :]
|
|
# _, faces = _maskFaces(np.zeros((14475,3), 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,1],),
|
|
# yaxis = dict(nticks=4, range=[-1,1],),
|
|
# zaxis = dict(nticks=4, range=[-1,1],),)
|
|
# )
|
|
|
|
# fig.show()
|
|
|
|
|
|
np.savez_compressed(os.path.join(folder, 'mesh_posed_no_orient', id + '_' + str(frame) + '.npz'), mesh_posed_no_orient=out)
|