one_hot unposing
This commit is contained in:
@@ -0,0 +1,230 @@
|
||||
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()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user