231 lines
7.5 KiB
Python
231 lines
7.5 KiB
Python
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
|
|
|
|
from custom_utils.custom_utils import maskFacesBatch, _maskFaces
|
|
|
|
|
|
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.smpl_2 import SMPL as SMPL_2
|
|
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):
|
|
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, id
|
|
|
|
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]
|
|
_shape = np.tile(_shape[np.newaxis, :], [_pose.shape[1], 1])
|
|
|
|
posed = np.zeros((300, 14475, 3), np.float32)
|
|
pose = np.zeros((72, 300), np.float32)
|
|
shape = np.zeros((300, 11), np.float32)
|
|
|
|
frames = _posed.shape[0]
|
|
posed[:frames, :, :] = _posed
|
|
pose[:, :frames,] = _pose
|
|
shape[:frames, :] = _shape
|
|
|
|
return posed, mask, pose, shape, 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
|
|
|
|
|
|
@tf.function
|
|
def smpl_run_wrapper(R, shape):
|
|
return smpl.run(R, shape)
|
|
|
|
@tf.function
|
|
def smpl_run_backward_wrapper(Rt, J, posed):
|
|
return smpl.run_backward(Rt, J, posed, smpl.weights_one_hot)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("-f", "--folder", help = "end of range", default='/home/workspace/data/dataset/test_skirt')
|
|
parser.add_argument("-s", "--start", help = "start of range", default=0)
|
|
parser.add_argument("-e", "--end", help = "end of range", default=1500)
|
|
parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
|
|
args = parser.parse_args()
|
|
|
|
F = np.load('custom_utils/smpl_faces_skirt.npy')
|
|
folder = args.folder
|
|
smpl = SMPL()
|
|
# smpl_2 = SMPL_2('custom_utils')
|
|
|
|
is_skirt = True
|
|
|
|
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
|
|
|
|
# if len(id) == 5:
|
|
# continue
|
|
|
|
os.makedirs(os.path.join(folder, id, 'mesh_unposed_one_hot'), exist_ok=True)
|
|
|
|
posed, mask, pose, shape, mask_idx = read_garment(folder, id)
|
|
# gender = (shape[:, -1] > 0).astype(np.int32)
|
|
# 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
|
|
|
|
frames = posed.shape[0]
|
|
|
|
# R = smpl.run_rodrigues(np.transpose(pose, [1,0]))
|
|
|
|
if is_skirt:
|
|
pose_dress = np.copy(pose)
|
|
p_avg = (pose[:, 3:6] + pose[:, 6:9]) / 2
|
|
pose_dress[:,3:9] = np.concatenate([p_avg, p_avg], axis=-1)
|
|
R = smpl.run_rodrigues(np.transpose(pose_dress, [1,0])).numpy()
|
|
for b in range(frames):
|
|
R[b][[4,5,7,8,10,11],:,:] = np.eye(3)
|
|
else:
|
|
R = smpl.run_rodrigues(np.transpose(pose, [1,0]))
|
|
|
|
Rt = tf.transpose(R, [0, 1, 3, 2])
|
|
T_body, J, _ = smpl_run_wrapper(R, shape)
|
|
# zero_body = smpl_run_wrapper(smpl.run_rodrigues(np.zeros((1, 72), np.float32)), shape[:1, :])[0][0].numpy()
|
|
|
|
|
|
T0_garment, J = smpl_run_backward_wrapper(Rt, J, posed)
|
|
# T_garment, _ = smpl.run_garment_orig(R, shape, T0_garment) # =====> you do not need this and next line for unposing
|
|
# metric = tf.reduce_max(tf.abs(mask[tf.newaxis,:,tf.newaxis] * (T_garment - posed)))
|
|
|
|
|
|
for frame in range(frames):
|
|
np.savez_compressed(os.path.join(folder, id, 'mesh_unposed_one_hot', id + '_' + str(frame) + '.npz'), mesh_unposed=T0_garment[frame])
|
|
|
|
# F = np.load("custom_utils/smpl_faces" + ("_skirt" if is_skirt else "") + ".npy")
|
|
# _, faces = _maskFaces(np.zeros((14475, 3)), F, mask.astype(bool))
|
|
|
|
# batch = 0
|
|
# x = T0_garment[batch].numpy()[mask, :]
|
|
# y = zero_body
|
|
# 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 = F[:,0],
|
|
# j = F[:,1],
|
|
# k = F[:,2],
|
|
# )],
|
|
# ).show()
|
|
|
|
|
|
# print()
|
|
|
|
|