117 lines
4.3 KiB
Python
117 lines
4.3 KiB
Python
import os
|
|
import sys
|
|
import time
|
|
import cv2
|
|
import random
|
|
import numpy as np
|
|
import tensorflow as tf
|
|
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__':
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = "5"
|
|
|
|
uv2mesh = UV2Mesh()
|
|
smpl = SMPL()
|
|
|
|
F = np.load("custom_utils/smpl_faces.npy")
|
|
|
|
factory = DatasetFactory()
|
|
signature_dynamic = { 'id': (),
|
|
'pose_shape': (72 + 11,),
|
|
# 'uv_posed': (3, 512, 256,),
|
|
# 'uv_body_posed': (3, 512, 256,),
|
|
# 'mesh_body_posed': (14475, 3,),
|
|
# 'mesh_posed_body_no_orient': (14475, 3,),
|
|
'mesh_posed': (14475, 3,),
|
|
'mesh_unposed': (14475, 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': (14475,),
|
|
# 'uv_mask': (512, 256,),
|
|
# 'uv_shaped_body': (3, 512, 256,),
|
|
}
|
|
|
|
dataset = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train',
|
|
batch_size=1,
|
|
signature=signature_dynamic)
|
|
static_dataset = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train',
|
|
signature=signature_static)
|
|
|
|
|
|
uv_res = 256
|
|
|
|
|
|
|
|
for i, batch in enumerate(dataset):
|
|
print(i)
|
|
batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
|
|
|
|
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()
|
|
# mesh_posed = np.squeeze(batch['mesh_posed'].numpy())
|
|
gender = int(shape[0, -1] > 0)
|
|
print(gender)
|
|
|
|
shape_offsets = tf.tensordot(shape[:,:10], smpl.shapedirs[gender], axes=[[1], [2]]).numpy()
|
|
|
|
|
|
mesh_unposed[mask, :] = mesh_unposed[mask, :] - np.squeeze(shape_offsets)[smpl.to_keep, :][mask, :]
|
|
|
|
np.savez_compressed(os.path.join(folder, id, 'mesh_posed_body_no_orient', id + '_' + str(i) + '.npz'), mesh_posed_body_no_orient = T[i, :, :])
|
|
|
|
|
|
# _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
|
|
|
|
# x = mesh_unposed[mask, :]
|
|
# y = mesh_unposed[mask, :] - np.squeeze(shape_offsets)[smpl.to_keep, :][mask, :]
|
|
|
|
# # x = smpl.v_template.numpy()[0, smpl.to_keep, :]
|
|
# # R = smpl.run_rodrigues(np.zeros((1, 72), dtype=np.float32))
|
|
# # x = smpl.run(R, shape)[0][0]
|
|
# # y = smpl.v_template.numpy()[gender, smpl.to_keep, :] + np.squeeze(shape_offsets)[smpl.to_keep, :]
|
|
# # 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,1],),
|
|
# yaxis = dict(nticks=4, range=[-1,1],),
|
|
# zaxis = dict(nticks=4, range=[-1,1],),)
|
|
# )
|
|
|
|
# fig.show()
|
|
|
|
print() |